neural-λ

Embedding-Space Alignment Enables Unbounded Compositional Generalization in Tiny Recursive Interpreters

A sub-5M-parameter recursive interpreter that:

• Reaches 0.9818 ± 0.003 on canonical LRA ListOps with a fully learned parser + REINFORCE fine-tune, +34.4 points over the prior best (S7, 0.6377).
• Achieves 100% zero-shot generalization to compositions of depth 500 from depth-1 primitive training, at 417K parameters.
• Traces both to one measurable mechanism: an MSE alignment loss that creates an algebraic homomorphism between symbolic operations and continuous representations (β-reduce outputs land on the value-encoder's image, cosine similarity 0.997).

Paper draft: paper/paper.pdf.

TL;DR

Compositional generalization is not a scale problem. It is a structural one: the recursive identity that makes composition tractable, the operator's outputs lying in the same space as its inputs, is not enforced by standard architectures. We enforce it with an MSE alignment loss and get unbounded depth generalization at 417K parameters and canonical-benchmark SOTA at 4.5M parameters.

Results table

LRA ListOps (3 seeds: 42, 7, 123)

Operating point	Params	Test acc.
Vanilla Transformer (Tay et al. 2021)	—	0.3637
S4, S5, MEGA	—	0.596–0.631
S7 (prior best)	—	0.6377
neural-λ, learned parser, argmax	4.5M	0.5017 ± 0.006
neural-λ, learned parser, K=4 majority vote	4.5M	0.8628 ± 0.018
neural-λ, learned parser, REINFORCE 10k steps	4.5M	0.9818 ± 0.003
neural-λ, deterministic parse + reasoner (diagnostic)	2.7M	0.8628 ± 0.018
neural-λ, reasoner-only 30k (diagnostic)	2.7M	0.9273 ± 0.035

Arithmetic depth extrapolation (417K-param model, depth-1 training)

Composition depth	Accuracy
2, 3, 4, 10, 50, 100, 200, 500	1.000

Reproduction

All scripts assume a single GPU. H100 is what we used; anything with ≥24GB should work.

Quickstart: verify the homomorphism claim (~2 min on H100)

# Trains one model with alignment loss, one without. Generates figure 9 (manifold).
python3 scripts/make_figure_manifold.py
# Expected: cos sim without alignment ≈ 0, with alignment ≈ 0.996.

Arithmetic depth-500 (~3 min on H100)

python3 scripts/train_depth_scaling.py
# Trains on 1,200 depth-1 primitives; evaluates at depth 2..500.
# Expected: 100% at every depth.

LRA ListOps reproduction

# 1. Get the canonical LRA ListOps data (Kaggle mirror)
bash data/lra/download.sh  # or see data/lra/README.md

# 2. Train parser + reasoner (3 seeds, ~35 min per seed on H100)
python3 scripts/lra_train.py
# Saves checkpoints to ./checkpoints/seed_{42,7,123}.pt
# Expected argmax TEST: 0.50 ± 0.01

# 3. REINFORCE fine-tune (~35 min × 3 seeds)
for s in 42 7 123; do
  python3 scripts/lra_finetune_reinforce.py \
    checkpoints/seed_${s}.pt \
    checkpoints/seed_${s}_reinforce10k.pt 10000
done
# Expected TEST: 0.98 ± 0.003 (this is the headline 0.9818)

# 4. Eval with all inference strategies (argmax / random / majority vote / oracle)
for s in 42 7 123; do
  python3 scripts/eval_margin_rerank.py checkpoints/seed_${s}.pt data/lra/test.tsv 4
done

Architecture

Three learned modules + one dispatcher (hardcoded OR learned):

• fe(op, c): function encoder, 2-layer transformer
• ae(x): argument encoder, embedding lookup
• β(f, v): β-reduction operator, 2-layer MLP (or variable-arity transformer for LRA), shared across all positions
• dec(v): value decoder

The recursive evaluator:

def evaluate(expr_tree, x):
    if expr_tree is leaf:
        return ae(x)
    op, c, children = expr_tree
    child_vals = [evaluate(c, x) for c in children]
    return β(fe(op, c), child_vals)  # variable-arity β for LRA

The same β is invoked at every node. There is exactly one β per model.

For LRA, we additionally train a tokens→tree parser (1.8M-param transformer predicting parent-pointers) with cross-entropy against the deterministic parse derivable from the bracket grammar.

Training signal

Two losses:

1. Cross-entropy on dec(evaluate(...)) against the target value.
2. MSE alignment: ||β(fe(op,c), ae(x)) − ae(target)||², applied at depth-1 primitives.

The alignment loss creates an algebraic homomorphism: it forces β(f, ae(x)) to land on ae(target), measured cosine similarity 0.997. With this, β(f₂, β(f₁, ae(x))) ≈ β(f₂, ae(intermediate)) is always in distribution, and depth generalization is automatic.

For the LRA pipeline's REINFORCE fine-tune: sample parses from the parser distribution, build trees, run the reasoner, reward = −CE(reasoner_output, digit_target), update parser via REINFORCE (with a batch-mean baseline) and reasoner via CE on the sampled tree. No parse supervision during fine-tune.

Repository layout

paper/         # LaTeX source + compiled paper.pdf (NeurIPS 2026 submission)
scripts/
  smoke_test.py          # Initial validation
  train_align_loss_discovery.py            # Discovery of the alignment loss trick
  train_depth_scaling.py          # Headline arithmetic depth-500 run
  train_2x2_ablation.py            # 2×2 ablation (architecture × loss)
  train_depth_500.py          # Depth-500 + homomorphism check
  lra_train.py              # LRA end-to-end training (parser + reasoner)
  lra_train_reasoner_only.py            # LRA reasoner-only ablation
  lra_finetune_reinforce.py       # REINFORCE fine-tune (headline number)
  eval_margin_rerank.py           # K-sweep + rerank comparison
  eval_deterministic.py           # Deterministic-parse eval
  eval_per_length.py              # Per-sequence-length breakdown
  eval_failure_decomp.py          # Per-operator failure analysis
  make_figure_manifold.py          # Figure 9 (homomorphism manifold)
docs/figures/   # Compiled PDF/PNG figures
results/        # Raw training logs
data/lra/       # LRA ListOps canonical data (downloadable)

Honest limitations

• Algebraic closure is required: β-reduce outputs must live in the same space as inputs. This rules out variable-length sequence generation benchmarks such as SCAN. Extending closure to autoregressive decoding is our primary open direction.
• Fairness on LRA: our parser uses parent-pointer supervision derived from the bracket grammar (deterministic, computable from the input, but not used as auxiliary loss by sequence-model baselines). The REINFORCE fine-tune uses only the digit target as reward signal, but starts from a parse-supervised initialization. We discuss this explicitly in §3.X of the paper.
• β is specialized to its training operation family (frozen β cannot host new operators without retraining).
• SUM_MOD is the dominant failure mode on LRA: selection-based ops (MIN/MAX/MED) are solved at ~97%, SUM_MOD drops to ~80% due to compounding modular-arithmetic error. 76% of errors are off-by-one mod 10.

What this paper is about

We isolate a mechanism — an MSE alignment loss that creates an algebraic homomorphism — and show it is sufficient for (a) unbounded depth generalization on primitive-only algebraic training and (b) canonical-benchmark SOTA on LRA ListOps with a fully learned parser. The mechanism is verified mechanistically (direct geometric measurement), causally (α-sweep phase transition), and across five algebraic domains plus one canonical benchmark.

Within the niche of compositional benchmarks, our results are consistent with structural inductive bias being complementary to scale, not a substitute for it.

Citation

@inproceedings{fareed2026neuralambda,
  title   = {Embedding-Space Alignment Enables Unbounded Compositional Generalization in Tiny Recursive Interpreters},
  author  = {Fareed, Kareem},
  booktitle = {Submitted to NeurIPS 2026},
  year    = {2026},
}