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 (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:

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

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

@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)}
}

Downloads last month: -; Downloads are not tracked for this model. How to track

Ajay0728
/

Compositional-reasoning-arithmetic