Exact Learning of Arithmetic with Differentiable Agents

AI summary

Overview

Research area: Machine learning — algorithmic learning, length generalization, and differentiable automata (with connections to computational learning theory and neurosymbolic AI).

Technical level: Advanced. The paper assumes familiarity with Turing machines, finite-state transducers, formal learning theory (learning in the limit), and gradient-based training.

Scope: The paper argues and empirically demonstrates that Differentiable Finite-State Transducers (DFSTs), a Turing-complete model family trained end-to-end by gradient descent on expert policy-trajectory observations, achieve exact-looking length generalization on binary and decimal addition and multiplication.

What This Paper Is About

Most neural sequence models learn arithmetic from input-output examples and then fail when asked to compute on numbers longer than those seen in training. This paper asks whether a differentiable agent can instead learn the underlying algorithm, and it answers by training on the intermediate steps an expert takes while solving the problem rather than on final answers alone. The goal is "exact learning" in Angluin's sense — correct application of the ground-truth rule to every possible input — approximated in practice by measuring length generalization far beyond the training lengths.

Key Contributions

1. A differentiable, Turing-complete agent family. The paper introduces Differentiable Finite-State Transducers (DFSTs), control units parameterized by three weight tensors and an initial hidden state that both edit a symbolic grid and choose where to look next, with constant precision, constant-time output generation, and end-to-end log-parallel training.

2. A universality proof. Theorem 2.1 states that the map from p-precision DFSTs to grid agents is surjective, that the DFST family is Turing-complete when interacting with an external symbolic grid, and that any d-state grid agent can be emulated by DFST agents of dimension d.

3. A framework for learning from intermediate supervision. Building on the formalism of Papazov and Flammarion [31], the paper trains DFSTs on policy-trajectory observations (PTOs) — the observed symbol and motion tokens of an expert grid agent — rather than on input-output pairs, which Gold [13] showed is insufficient for exactly learning recursive functions.

4. State-of-the-art length generalization on four arithmetic tasks. DFSTs trained on tiny datasets (20, 225, 750, and 10,000 samples) reach robust and probabilistic length generalization of 3850, 2450, 600, and 180 digits for add2, add10, mult2, and mult10, with no error found as far as testing was possible.

Main Findings

• Length generalization thousands of times beyond training length: Models for add2 and add10 trained on summands with at most 3 digits achieved both RLG and PLG of 3850 and 2450 digits respectively; models for mult2 and mult10 trained on multiplicands with up to 5 digits achieved RLG and PLG of 600 and 180 digits respectively. Testing beyond these lengths produced OutOfMemory errors on the authors' A100-SXM4-80GB GPU.

• No errors observed within tested range: For each arithmetic task, the authors report they could not find a single error in the computation of their best trained models as far as they tested.

• Few training samples, few parameters: The DFST models used 20 (add2), 225 (add10), 750 (mult2), and 10000 (mult10) samples and 1020, 8900, 5280, and 162108 parameters respectively. By comparison, Table 1 lists the Neural GPU at roughly 200k samples and 10368 parameters for add2 and mult2.

• Comparison with the Neural GPU: Table 1 reports the Neural GPU with input-output data reaching Robust LG of ≤20 (add2) and ≤8 (mult2) and Probabilistic LG of ≥2000 for both, with no reported values for add10 or mult10. The authors note that Price et al. [33] found the Neural GPU's generalization depends heavily on random seed, with only a few seeds reaching 2000-digit PLG, and that all trained models fail on highly structured symmetric examples.

• Perfect accuracy on symmetric hard cases: The authors report their models achieve perfect accuracy on long symmetric examples as far as they could test, whereas the Neural GPU fails on such examples. They identify same-digit numbers as the hard test instances for generalization.

• Simpler training recipe: The DFST training uses only a standard Adam optimizer with a cosine annealing scheduler, batch size 32, and float32 precision, in contrast to the Neural GPU's use of curriculum learning, gradient-noise injection, gradient scheduling, relaxation-pull, dropout, and gate cut-off.

• Training data matters more than model class: The authors attribute the improved robustness to observable intermediate steps (PTOs), which "greatly diminishes the number of training samples, and enables the learning of decimal arithmetic."

• Training-time trade-off: Appendix A reports that increased training times correlate with lower training loss and improved robust length generalization across model sizes and tasks.

• Verification is fundamentally limited: Because checking whether two grid agents follow the same deterministic policy reduces to the undecidable language EQ_TM, no general procedure can prove π_M = π_D from the model descriptions; the authors rely on random tests for confidence.

Methodology in Plain English

The authors model the problem as an agent moving on an unbounded two-dimensional grid of symbols. An expert "grid agent" solves an arithmetic problem by reading a cell, writing a symbol, and moving up, down, left, right, or staying — much like a teacher working at a blackboard. A student watching this expert does not see the teacher's internal state, only the sequence of symbols read and moves made. That sequence is a policy-trajectory observation (PTO).

The trainable model, a DFST, is deliberately simple and linear. It keeps a hidden vector, updates it by multiplying with a matrix chosen by the currently observed symbol, and from that state predicts two things: the symbol to write and the direction to move. Both predictions are made with an argmax over linear outputs, and the whole thing is trained with a mean-squared-error loss under a next-action-prediction objective, using one-hot expert actions as targets. Because there are no nonlinearities, the state updates can be parallelized across time with a parallel scan, giving fast log-parallel training.

Initialization is designed to make learning stable: state-transition matrices are initialized as identities, output tensors as zeros, and the initial hidden state as a random vector normalized to unit Euclidean norm. The authors match each DFST's hidden dimension to the number of states in the corresponding expert grid agent (12, 20, 32, and 108). Training data is generated by a sampling function that mixes exhaustively enumerated short problems, same-digit problems, and random longer problems, so that a large portion of the curriculum uses extremely short inputs.

Generalization is measured two ways. Probabilistic length generalization (PLG) is the largest digit count at which the model is perfect on 5 random pairs of exactly that length plus 5 random pairs of at most that length. Robust length generalization (RLG) adds the harder requirement of perfect exact-match accuracy on all R² pairs of numbers made of a single repeated digit.

Why This Matters

The paper connects formal learning theory to practical deep learning, showing that the theoretically motivated choice of training data — intermediate computational traces rather than input-output pairs — can produce dramatically better generalization than architectural scaling or elaborate training tricks. It also introduces a model family with a rare combination of properties: Turing completeness, constant precision, constant-time generation, and fully differentiable training, which the authors argue prior universal architectures sacrificed for optimization.

Real-world applications implied by this work:

• Verifiable arithmetic and numerical libraries: A model that provably executes a known algorithm on arbitrarily long inputs would be suitable for computation where an approximate answer is unacceptable.
• Program synthesis and algorithm induction: Training from execution traces is the same signal available from debuggers, interpreters, and instrumented programs, suggesting a route to learning algorithms from observed runs.
• Hardware and accelerator design: Constant-precision, constant-time, log-parallel computation maps naturally onto parallel hardware, unlike architectures whose internal precision or memory must grow with input length.
• Education and tutoring systems: The grid-agent model of a teacher moving over a blackboard is a direct formalization of worked-example pedagogy, which is exactly the data these systems can record.

Industry relevance is strongest for groups building reasoning or algorithmic components into larger systems, and for anyone who needs sequence models whose correctness on long inputs can be increased by changing the supervision signal rather than by scaling the model.

Future Directions

• Theoretical analysis of the loss landscape: The authors explicitly hope their linear, minimal DFST will enable theoretical study of the optimization landscape of algorithmic learning.

• Removing the reliance on random testing: Since EQ_TM is undecidable, the paper can only gain confidence in exactness through random tests. Finding restricted, decidable settings or certified verification procedures is an open problem.

• Scaling beyond GPU memory: The reported RLG and PLG values are bounded by OutOfMemory errors rather than by model failure, so testing on longer inputs — or with more memory-efficient PTO handling — is a direct next step.

• Relations to other sequence models: The authors note any sequential model capable of state tracking (RNNs, LSTMs, GRUs) could in principle learn a grid agent's policy, but not Transformers or diagonal and non-gated SSMs. Whether these architectures can be adapted to the DFST setting remains open.

• Broader algorithmic coverage: The paper covers four arithmetic tasks; whether the same framework extends to other algorithmic skills such as searching and sorting is untested here.

Target Audience

Researchers working on algorithmic reasoning, length generalization, and neurosymbolic methods will find the direct comparison with the Neural GPU and the PTO training framework most useful. Theoreticians interested in learning in the limit, computability, and differentiable automata will benefit from the universality theorem and the deliberate linearity of the model. Practitioners building systems that must compute correctly on inputs far longer than their training distribution — especially those already using scratchpad or chain-of-thought supervision — are the most likely to apply the ideas, though the material assumes comfort with formal automata and gradient-based training.

Authors’ abstract

We explore the possibility of exact algorithmic learning with gradient-based methods and introduce a differentiable framework capable of strong length generalization on arithmetic tasks. Our approach centers on Differentiable Finite-State Transducers (DFSTs), a Turing-complete model family that avoids the pitfalls of prior architectures by enabling constant-precision, constant-time generation, and end-to-end log-parallel differentiable training. Leveraging policy-trajectory observations from expert agents, we train DFSTs to perform binary and decimal addition and multiplication. Remarkably, models trained on tiny datasets generalize without error to inputs thousands of times longer than the training examples. These results show that training differentiable agents on structured intermediate supervision could pave the way towards exact gradient-based learning of algorithmic skills. Code available at \href{https://github.com/dngfra/differentiable-exact-algorithmic-learner.git}{https://github.com/dngfra/differentiable-exact-algorithmic-learner.git}.