Exact Learning of Arithmetic with Differentiable Agents

🕒 Published (v1): 2025-11-27 20:51 UTC · Source: Arxiv · Venue: NEURIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper introduces Differentiable Finite-State Transducers (DFSTs), a Turing-complete, end-to-end differentiable model family that achieves near-exact length generalization on arithmetic tasks by training on policy-trajectory observations (PTOs) from expert grid agents. Models trained on datasets of 20–10,000 examples generalize without error to inputs thousands of times longer than training examples. This constitutes a concrete step toward gradient-based exact learning of algorithmic skills.

Problem

Prior neural architectures capable of Turing-completeness suffer from either exploding computation time (unbounded-precision RNNs, growing-context Transformers) or uninformative gradients (non-differentiable adaptive computation, discrete environment interactions), making true gradient-based exact learning of algorithms intractable. Transformers and modified RNNs trained on input-output data achieve at best 2–3× length generalization with imperfect accuracy.

Method

DFSTs are linear recurrent agents with weight tensors A, B, C ∈ ℝ^(|Σ|×d×d) operating over a symbolic 2D grid. At each step, the DFST reads a cell symbol x_t, updates its hidden state via h_{t+1} = A(x_t)h_t, and emits a symbol token (argmax of B(x_t)h_t) and a motion token (argmax of C(x_t)h_t) that both edit and navigate the grid. Training uses MSE loss under a next-action prediction (NAP) objective on PTOs — full computational traces (symbol + motion tokens) generated by hand-coded expert grid agents performing binary/decimal addition and multiplication. The linear structure enables log-parallel training via parallel scans. Model hidden-state dimensions match the exact state count of the target ground-truth agent (12–108 states).

Key Contributions

• DFST architecture: a Turing-complete, constant-precision, constant-time-per-step, end-to-end differentiable agent that decouples the control unit from an external symbolic grid.
• Policy-trajectory observations (PTOs): structured intermediate supervision (full computational traces) enabling reduction of learning recursive functions to learning finite-state transducers.
• Unprecedented length generalization: RLG and PLG of ≥3850 digits (binary add), ≥2450 digits (decimal add), ≥600 digits (binary mult), ≥180 digits (decimal mult) — trained on operands of ≤3–5 digits.
• Data efficiency: add2 trained on only 20 examples; prior Neural GPU required ~200k samples and complex training tricks for weaker generalization.
• Zero observed errors: no incorrect computation found on any tested input up to GPU memory limits, including symmetric/same-digit hard cases that break Neural GPU.

Results

• add2: RLG ≥ 3850 digits, PLG ≥ 3850 digits; 20 training samples, 1020 parameters; Neural GPU baseline: RLG ≤ 20, PLG ≥ 2000 with ~200k samples and 10,368 parameters.
• add10: RLG ≥ 2450 digits, PLG ≥ 2450 digits; 225 training samples, 8900 parameters; Neural GPU: not reported (N/A).
• mult2: RLG ≥ 600 digits, PLG ≥ 600 digits; 750 training samples, 5280 parameters.
• mult10: RLG ≥ 180 digits, PLG ≥ 180 digits; 10,000 training samples, 162,108 parameters.
• Longer training time monotonically improves RLG across all tasks; all results memory-bounded by A100-SXM4-80GB at test time (PTO lengths reach ~30M operations for 3850-digit binary addition).

Limitations

• Length generalization upper bounds are GPU-memory-constrained, not algorithmic — true generalization ceiling is unknown.
• Verifying exact learning (equivalence of two Turing machines) is undecidable; correctness is established only via random empirical testing.
• Requires PTOs (full computational traces from expert agents), not just input-output pairs — limits applicability where intermediate steps are unavailable.
• Decimal multiplication (mult10) achieves substantially lower RLG (≥180) versus binary addition (≥3850), suggesting harder algorithmic structure strains the approach.
• No experiments beyond the four arithmetic tasks; generalization to other algorithmic domains is undemonstrated.

Relevance to Agentic AI / LLM Agents

DFSTs represent a principled alternative to LLM-based agents for exact algorithmic execution: rather than prompting a language model to simulate computation, a differentiable agent is trained end-to-end to internalize and execute an algorithm with provable generalization. The PTO training paradigm — learning from expert agent traces rather than final answers — directly parallels scratchpad/chain-of-thought supervision in LLM agents, but with formal guarantees. For researchers building agentic systems that must perform reliable tool use, multi-step reasoning, or symbolic computation, this work suggests that lightweight structured agents trained on intermediate supervision can far outperform large neural models on algorithmic sub-tasks, potentially as callable modules within a broader LLM agent pipeline.