[Paper Review] Neural Arithmetic Logic Units

Andrew Trask, Felix Hill+4 more|arXiv (Cornell University)|Aug 1, 2018

Neural Networks and ApplicationsComputer Science80 citations

TL;DR

The paper introduces Neural Accumulator (NAC) and Neural Arithmetic Logic Unit (NALU), differentiable modules biasing neural nets toward systematic numerical computation to improve extrapolation, demonstrated across synthetic, image, language, code, and reinforcement learning tasks.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for systematic numerical extrapolation in neural networks beyond training ranges.
• Propose architectures (NAC and NALU) that bias representations toward arithmetic operations.
• Demonstrate improved extrapolation generalization across multiple domains (synthetic, vision, language, program execution, RL).
• Show that replacing standard linear layers with NAC/NALU can yield substantial gains in numeral reasoning tasks.

Proposed method

• Define Neural Accumulator (NAC) as a constrained linear layer with W = tanh(What) * sigma(Mhat) to encourage inputs to map via additions/subtractions without rescaling.
• Extend NAC to Neural Arithmetic Logic Unit (NALU) by combining an add/subtract NAC and a multiply/divide NAC through a learned gate g, enabling arithmetic functions (add, subtract, multiply, divide, power).
• Train end-to-end with backpropagation on tasks requiring numerical manipulation, using both supervised and reinforcement learning signals.
• Apply NAC/NALU inside various architectures (feed-forward, LSTM-based, and RL agents) to test extrapolation and generalization beyond training ranges.

Experimental results

Research questions

• RQ1 Can NAC and NALU enable neural networks to extrapolate numerically beyond the training range?
• RQ2 Do NAC/NALU provide a numeracy bias that improves arithmetic reasoning across diverse modalities (synthetic, vision, language, code, RL)?
• RQ3 Does replacing the last linear layer with NAC/NALU improve performance on numeral-related tasks such as counting, translation of number words, and program evaluation?
• RQ4 How does multiplicative capability (via NALU) compare to additive capability (via NAC) across tasks?
• RQ5 Can NAC/NALU support end-to-end learning with backpropagation for tasks requiring precise numeric reasoning?

Key findings

• NAC enables robust additive/subtractive computations and improves extrapolation over standard linear/MLP baselines.
• NALU adds multiplicative/division capabilities via a controlled gate, enabling extrapolation for multiplicative functions as well.
• Across tasks (synthetic arithmetic, MNIST counting, language-to-number translation, program evaluation, and grid-world time-tracking), NAC/NALU demonstrate superior extrapolation (often orders of magnitude beyond training ranges).
• On MNIST parity, NAC/NALU reduce error by notable margins compared to prior linear-layer baselines, with NAC/NALU providing strong interpolation performance as well.
• Replacing the last linear layer with NAC/NALU can yield substantial performance gains, illustrating the benefit of a numerically biased architectural component.
• In RL, NAC-enabled control can improve extrapolation in time-tracking tasks, suggesting utility of numerically biased modules beyond supervised settings.