Source-linked AI summary
Neural Arithmetic Logic Units
Andrew Trask, Felix Hill, Scott Reed, Jack Rae, Chris Dyer, Phil Blunsom
TL;DR
Neural networks often fail to generalize numerical computations beyond the values seen during training. The paper introduces NALU, which uses linear number representations, primitive arithmetic operators, and learned gates, and reports broad extrapolation across tasks and domains.
Problem
Neural networks frequently fail to systematically generalize numerical representations and functions outside the numerical range observed during training.
Method
The paper introduces NAC and NALU modules that represent quantities with linear activations and manipulate them using differentiable arithmetic operators controlled by learned parameters or gates.
Results
Across synthetic, image, text, code, and reinforcement-learning tasks, the models generalize to values several orders of magnitude beyond training, including a 54% image-counting error-margin advantage.
Takeaways & Limitations
NAC and NALU provide a design strategy for biasing neural networks toward systematic numerical computation across diverse settings.
Takeaways & Limitations
Extrapolation remains limited in some settings, including time-tracking stimuli above 20 for the NAC agent and division, which is especially challenging to extrapolate.
Abstract
from arXiv · showhide
Neural networks can learn to represent and manipulate numerical information, but they seldom generalize well outside of the range of numerical values encountered during training. To encourage more systematic numerical extrapolation, we propose an architecture that represents numerical quantities as linear activations which are manipulated using primitive arithmetic operators, controlled by learned gates. We call this module a neural arithmetic logic unit (NALU), by analogy to the arithmetic logic unit in traditional processors. Experiments show that NALU-enhanced neural networks can learn to track time, perform arithmetic over images of numbers, translate numerical language into real-valued scalars, execute computer code, and count objects in images. In contrast to conventional architectures, we obtain substantially better generalization both inside and outside of the range of numerical values encountered during training, often extrapolating orders of magnitude beyond trained numerical ranges.
1 Introduction
Standard neural networks can represent numerical quantities but often fail to generalize systematically beyond training ranges. The paper introduces a differentiable arithmetic module and evaluates it across diverse domains, reporting extrapolation far beyond observed values.
- Motivation: Neural networks often fail when test quantities fall outside the training range, even for simple counting or linear-extrapolation functions.This failure suggests memorization rather than systematic abstraction.
- Contribution: The proposed module represents numbers in individual linear neurons and applies differentiable arithmetic operators controlled by learned parameters.It is designed to work with standard architectures such as LSTMs and convolutional networks.
- Evaluation: Experiments span synthetic, image, text, and code tasks using supervised and reinforcement learning with feed-forward and recurrent structures.The experiments test whether learned functions capture the underlying numerical nature of data.
- Results: Several orders of magnitude: the proposed models generalize to numbers much larger than those observed during training.In one image-counting case, the model exceeds a state-of-the-art network by an error margin of 54%.
- Motivation: Nonlinear MLPs fail to learn the scalar identity function outside the training range, despite architectures that could theoretically represent it.The paper uses this simple task to expose failures of systematicity.
2 The Neural Accumulator & Neural Arithmetic Logic Unit
The paper develops NAC and NALU modules that bias neural networks toward systematic arithmetic. NAC supports additive operations, while gated combinations of NACs extend this bias to multiplicative and related functions.
- NAC: NAC represents outputs as additions or subtractions of input rows using transformation weights constrained toward −1, 0, and 1.This preserves numerical scale across chained operations.
- NAC: W = tanh(Ŵ) ⊙ σ(M̂) provides a continuous differentiable parameterization biased toward the discrete NAC weights.The parameterization supports gradient-based learning while keeping matrix elements within [−1, 1].
- NALU: NALU combines additive and multiplicative subcells through a learned gate, extending NAC-style computation to new numerical functions.The design retains a bias against rescaling during input-to-output mapping.
- NALU: The NALU uses two tied-weight NACs: one computes addition/subtraction, while the other computes multiplication/division and related operations.The gate interpolates between the two subcells.
- NALU: Multiplication, addition, subtraction, division, and power functions can be learned in a form that extrapolates beyond training values.The logarithmic multiplicative pathway includes ε to prevent log 0.
3 Related Work
The related work connects numerical reasoning to image counting, code execution, systematic representation, architectural linearity, and cognitive theories of magnitude. The paper distinguishes its emphasis on numerical extrapolation beyond training ranges.
- Image counting: Image-counting systems either segment instances before post-processing or predict counts end-to-end with regression.The paper is more closely related to the end-to-end regression strategy.
- Code and question answering: Prior work also models numerical representations and arithmetic functions for code execution and bounded-range question answering.Existing counting tasks generally do not explicitly require counts beyond the training range.
- Architectural context: ResNets, Highway Networks, and DenseNets use linear connections to improve learning bias and reduce exploding or vanishing gradients, with added depth-related computational overhead.The paper places its linear-activation approach within this broader architectural trend.
- Systematicity: The work aligns with efforts to encode behavior-governing equations so models extrapolate to unseen parts of a problem space.This includes systematic representations in recurrent memory for longer-than-training sequences.
- Cognitive context: The single-neuron number representation recalls cognitive theories of continuous magnitude manipulation and proposed number-selective neurons.The paper notes that biological magnitude representations are often approximate.
4 Experiments
Across synthetic, image-based, language, program-evaluation, reinforcement-learning, and parity tasks, NACs and NALUs were evaluated for interpolation and extrapolation. NALUs generally extended numerical generalization beyond standard architectures, while some settings exposed task-specific limitations.
- Experimental scope: Experiments covered arithmetic functions, MNIST counting and addition, number-string translation, program evaluation, time tracking, and MNIST parity.The tasks included both direct numerical inputs and numerical reasoning from images, language, programs, and environment states.
- Simple and image-based tasks: Standard architectures often interpolated successfully but failed to extrapolate, whereas NACs and NALUs generalized beyond training ranges across synthetic and MNIST tasks.For simple functions, NACs handled addition and subtraction while NALUs also handled multiplicative operations, except recurrent division.
- Language to number translation: LSTM + NALU achieved the best generalization by a wide margin when translating number strings into scalar values.The dataset used numbers from 0 to 1000, with training concentrated on 0–19 plus sampled values from the remainder.
- Language to number translation: NALU intermediate states tracked estimates of unseen number expressions and supported compositional use of the “and” token without direct supervision.The model predicted tokens absent in isolation from training while combining expressions into addition-like structures.
- Program evaluation: Only NALU extrapolated stably to larger values in addition and program evaluation, including a domain increase of two orders of magnitude.UGRNN, LSTM, and DNC solved tasks at fixed input domains, but did not match NALU’s larger-number extrapolation.
- Learning to track time: In gridworld time tracking, the NAC agent performed well for T ≤19, while standard A3C performance deteriorated for T >13 after training on T ∼U{5, 12}.For stimuli greater than 20, the NAC agent also failed to arrive at the destination; the authors hypothesize residual LSTM numeracy contributed to this limit.
- MNIST parity and ablations: The NAC reduced prior MNIST parity error by 54%, with removing bias and applying nonlinearities to weights increasing end-to-end accuracy.The NAC replaced the original model’s final affine layer, while most parameters remained outside the NAC itself.
5 Conclusions
The NAC and NALU address poor numerical extrapolation by combining single-neuron number representations with learned gates over arithmetic functions, while remaining a general design strategy rather than a universal solution.
- NAC and NALU rectify failures in numerical representations and functions that generalize beyond training ranges across varied domains.
- The proposed single-neuron representation permits arbitrary differentiable numerical functions to be added and controlled through learned gates.
- NALU exemplifies learned gating between addition/subtraction and multiplication/division within this broader design strategy.
- The authors caution that NAC or NALU is unlikely to be the perfect solution for every task.
A Learning the Identity Function
Identity-function experiments test whether standard nonlinear networks learn systematic numerical representations, revealing severe extrapolation failures outside the training range.
- Identity reconstruction evaluates models trained on −5 to 5 and tested on −1000 to 1000, with Table 5 reporting average errors across nonlinearities.An error score of 500 corresponds to predicting 0 for every target.
- Sharply nonlinear activations produce greater extrapolation error than mildly nonlinear activations on the identity task.
B Synthetic Arithmetic Tasks
Synthetic arithmetic tasks compare standard nonlinear networks with NAC and NALU models on interpolation and extrapolation, including recurrent sequence-length generalization.
- Static arithmetic tasks: Static tasks sum two random subsections of x ∈R100 and apply arithmetic functions to the resulting scalars a and b to produce y.Each model must learn this process end-to-end.
- Static arithmetic tasks: Extrapolation tests require at least one of a, b, or y to exceed every value observed during training, unlike interpolation tests.
- Recurrent arithmetic tasks: Recurrent tasks replace vector sums with sums over timeseries steps xt ∈R10; training and interpolation use length 10, while extrapolation uses length 1000.Two NALUs are stacked end-to-end for comparison with one-hidden-layer MLP baselines.
- Static arithmetic tasks: Most common nonlinearities fail to extrapolate numerical functions beyond training ranges even under strong supervision, especially sigmoid and tanh.
- Recurrent arithmetic tasks: NAC and NALU significantly outperform recurrent baselines when extrapolating functions from sequences of length 10 to 1000.All models interpolate successfully, but no model solves division extrapolation; NAC and NALU underestimate the denominator.
- Conclusion: Systematic numerical representations in NAC and NALU provide the learning bias associated with accurate arithmetic extrapolation.
C MNIST Counting
The MNIST Counting and Addition evaluation reports accuracy across series lengths 1, 10, 100, and 1000 for the considered models.
- MNIST Counting and Addition performance is reported for series lengths 1, 10, 100, and 1000.
D Language To Number Translation Tasks
The summed-state LSTM was used as the comparison baseline against NAC and NALU, with a slight generalization advantage over using only the final state. Models were trained with fixed optimization and selection procedures.
- The summed-state LSTM served as the LSTM comparison for NAC and NALU in Table 3.
- All models were trained for 300K gradient-descent steps on the whole training set using Adam.
- The best model was selected by validation loss across two layer sizes, two learning rates, and 10 initializations.
- Summing LSTM states improved generalization slightly compared with outputting only the final state.
E Program Evaluation
In program evaluation, addition is easy for all models, but systematic generalization fails for every model except the recurrent NALU. The UGRNN performs better than the LSTM and DNC, likely because of its simple linear state update.
- All models solve the simpler addition task, but only the recurrent NALU generalizes beyond the training range.
- The UGRNN performs much better than the LSTM and DNC on the addition task.
- The authors attribute the UGRNN’s advantage likely to its simple linear state update.
- Figure 6 averages results over 10 independent runs and displays 2σ confidence bands.