README.md

---
license: apache-2.0
datasets:
  - Ajay0728/Compositional-reasoning-arithmetic-datasets
library_name: pytorch
pipeline_tag: text-generation
tags:
  - arithmetic
  - compositional-generalization
  - length-generalization
  - lstm
  - transformer
  - seq2seq
  - interpretability
---

# Compositional Reasoning on Arithmetic: LSTM and Transformer Checkpoints

This repo holds the 12 trained checkpoints from the paper *Partial Competence and Limited Transfer: Comparing LSTM and Transformer Generalization on Arithmetic Expressions* (arXiv ID pending).

The models are small character-level encoder-decoders, trained from scratch. Each one reads a fully parenthesized integer arithmetic expression and outputs its value:

```
Input:  ((14 - 1) * 10)
Output: 130
```

They are research artifacts for studying generalization. They are not useful calculators: the best model gets about 35% exact match on held-out in-distribution expressions.

- **Code:** https://github.com/bcode0127-debug/Compositional-Reasoning-Arithmetic
- **Data:** [Ajay0728/Compositional-reasoning-arithmetic-datasets](https://huggingface.co/datasets/Ajay0728/Compositional-reasoning-arithmetic-datasets) (`datasets_v2`)

## Checkpoints

There are 2 architectures × 2 studies × 3 seeds = 12 runs. Each file is the checkpoint with the best validation accuracy from its run.

| Path | Best epoch | Val | Test |
|---|---|---|---|
| `lstm/study1/seed0/best_model.pt` | 43 | 35.1% | 33.6% |
| `lstm/study1/seed1/best_model.pt` | 46 | 35.9% | 35.5% |
| `lstm/study1/seed2/best_model.pt` | 51 | 36.1% | 34.6% |
| `lstm/study2/seed0/best_model.pt` | 57 | 32.5% | 31.8% |
| `lstm/study2/seed1/best_model.pt` | 195 | 30.1% | 31.6% |
| `lstm/study2/seed2/best_model.pt` | 66 | 32.0% | 28.0% |
| `transformer/study1/seed0/best_model.pt` | 114 | 7.5% | 7.9% |
| `transformer/study1/seed1/best_model.pt` | 160 | 6.8% | 7.7% |
| `transformer/study1/seed2/best_model.pt` | 160 | 7.5% | 7.9% |
| `transformer/study2/seed0/best_model.pt` | 34 | 3.7% | 1.9% |
| `transformer/study2/seed1/best_model.pt` | 162 | 4.2% | 2.8% |
| `transformer/study2/seed2/best_model.pt` | 51 | 3.8% | 2.1% |

Each `.pt` file is a dict with these keys:

- `model_state_dict`: the weights
- `optimizer_state_dict`: the Adam state, kept so training can be resumed. It makes the files about 3× larger than the weights alone.
- `epoch`, `val_loss`, `val_accuracy`, `seed`
- `dataset_version` (always `"datasets_v2"`) and `generator_commit` (the commit that generated the training data)

## Models

| | LSTM | Transformer |
|---|---|---|
| Structure | Bidirectional LSTM encoder, unidirectional LSTM decoder | 3 encoder + 3 decoder layers, 8 heads, sinusoidal positions |
| Size | embedding 128, hidden 256, ~2.1M params | d_model 256, ~5.5M params |
| Learning rate | 1e-3 | 3e-4 |

Both architectures share the same training setup:

- Character-level vocabulary of 20 tokens: `<PAD> <SOS> <EOS>`, then `0123456789+-*/() ` (the last one is a space)
- Teacher forcing, Adam, batch size 64
- 400 warmup steps, then cosine decay down to 1e-5
- Up to 200 epochs, early stopping on validation accuracy with patience 50
- `max_input_len=64` and `max_output_len=12`. Inputs that are too long raise an error instead of being truncated.

## Usage

You need the model and tokenizer code from the GitHub repo:

```bash
git clone https://github.com/bcode0127-debug/Compositional-Reasoning-Arithmetic.git
cd Compositional-Reasoning-Arithmetic
pip install -r requirements.txt huggingface_hub
```

```python
import torch
from huggingface_hub import hf_hub_download
from data.tokenizer import create_tokenizer
from models.lstm import create_lstm_model
from models.transformer import create_transformer_model

tok = create_tokenizer()
path = hf_hub_download("Ajay0728/Compositional-reasoning-arithmetic",
                       "lstm/study1/seed0/best_model.pt")
ckpt = torch.load(path, map_location="cpu")

model = create_lstm_model(vocab_size=tok.vocab_size, embedding_dim=128, hidden_size=256)
# Transformer:
# model = create_transformer_model(vocab_size=tok.vocab_size, d_model=256, nhead=8,
#                                  num_encoder_layers=3, num_decoder_layers=3,
#                                  pad_idx=tok.pad_idx)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

def solve(expr, max_in=64, max_out=12):
    ids = tok.encode_checked(expr, max_in)
    src = torch.tensor([ids + [tok.pad_idx] * (max_in - len(ids))])
    out = [tok.sos_idx]
    with torch.no_grad():
        for _ in range(max_out - 1):
            dec = torch.tensor([out + [tok.pad_idx] * (max_out - len(out))])
            nxt = model(src, dec)[0, len(out) - 1].argmax().item()
            if nxt in (tok.eos_idx, tok.pad_idx):
                break
            out.append(nxt)
    return tok.decode(out[1:])

print(solve("((14 - 1) * 10)"))
```

`solve` does greedy autoregressive decoding. It is the same procedure behind the numbers below, implemented in `scripts/eval/cold_verify.py`. Inputs must use the training format: fully parenthesized, with spaces around each operator.

## Training data

The training data is `datasets_v2`, generated with `data/generate_controlled.py`. Every expression is a fully parenthesized integer expression. Division always comes out exact (no remainders), and every result satisfies |result| ≤ 10,000.

| Study | Train / Val / Test (in-distribution) | OOD |
|---|---|---|
| 1: Length | 2–3 operations, depth 2–3 (8,000 / 1,000 / 1,000) | 4, 5, 6, 7 operations, depth 3 (1,000 each) |
| 2: Depth | 3 operations, depth 2 (8,000 / 1,000 / 1,000) | 3 operations, depth 3 chains (1,000) |

Validation was used to pick each checkpoint. The test split was generated after training finished and is the primary in-distribution result.

Division makes up only about 3–5% of the operators in every split, because the exact-division requirement rejects most division candidates.

## Results

All numbers are exact-match accuracy under greedy autoregressive decoding. Val and OOD show the mean across seeds 0/1/2 with [min–max]. Test shows mean ± SD.

**Study 1: length generalization**

| Model | Val | Test | OOD ops4 | OOD ops5 | OOD ops6 | OOD ops7 |
|---|---|---|---|---|---|---|
| LSTM | 35.7% [35.1–36.1] | 34.57% ± 0.95 | 14.8% [13.7–16.4] | 8.2% [7.5–9.5] | 2.7% [2.5–2.8] | 2.1% [1.8–2.5] |
| Transformer | 7.3% [6.8–7.5] | 7.83% ± 0.12 | 0.7% [0.4–1.2] | 0.6% [0.5–0.7] | 0.3% [0.1–0.5] | 0.4% [0.2–0.6] |

**Study 2: depth generalization**

| Model | Val | Test | OOD |
|---|---|---|---|
| LSTM | 31.5% [30.1–32.5] | 30.47% ± 2.14 | 7.0% [6.3–7.8] |
| Transformer | 3.9% [3.7–4.2] | 2.27% ± 0.47 | 2.7% [2.0–3.1] |

The LSTM reaches partial competence in-distribution, and its accuracy falls steadily as expressions get longer or deeper. The Transformer stays weak in every split at this data scale.

## Limitations

- **Low accuracy:** these are small models trained on 8,000 examples, so they are wrong most of the time even in-distribution.
- **Narrow input format:** they only work on the exact input format and range they were trained on. Any other format is out of scope.
- **Training environment not recorded:** the torch/numpy versions used for training were never saved, so training can't be reproduced bit for bit. Evaluation only reloads the saved weights, so it doesn't depend on that environment. The test numbers were produced with torch 2.13.0 on CPU.
- **Older results are retired:** an earlier version of this project reported results from a pipeline that silently cut inputs to 20 characters. Those numbers are not valid. Every number on this card comes from the corrected v2 pipeline.

## Citation

```bibtex
@misc{jamanjyothi2026partial,
  title  = {Partial Competence and Limited Transfer: Comparing LSTM and Transformer
            Generalization on Arithmetic Expressions},
  author = {Jamanjyothi, Ajay},
  year   = {2026},
  note   = {arXiv preprint (ID pending)}
}
```