[Paper Review] Neural Arithmetic Units

Andreas Nygaard Madsen, Alexander Rosenberg Johansen|arXiv (Cornell University)|Jan 14, 2020

Neural Networks and ApplicationsComputer Science17 references20 citations

TL;DR

This paper introduces two novel neural network components—Neural Addition Unit (NAU) and Neural Multiplication Unit (NMU)—designed to enable exact arithmetic operations like addition, subtraction, and multiplication with improved stability, convergence speed, and extrapolation. The NMU, the first such unit to support large hidden sizes and both positive and negative values, achieves 94% success rate on a complex arithmetic task, significantly outperforming prior methods like NALU and NAC.

ABSTRACT

Neural networks can approximate complex functions, but they struggle to perform exact arithmetic operations over real numbers. The lack of inductive bias for arithmetic operations leaves neural networks without the underlying logic necessary to extrapolate on tasks such as addition, subtraction, and multiplication. We present two new neural network components: the Neural Addition Unit (NAU), which can learn exact addition and subtraction; and the Neural Multiplication Unit (NMU) that can multiply subsets of a vector. The NMU is, to our knowledge, the first arithmetic neural network component that can learn to multiply elements from a vector, when the hidden size is large. The two new components draw inspiration from a theoretical analysis of recently proposed arithmetic components. We find that careful initialization, restricting parameter space, and regularizing for sparsity is important when optimizing the NAU and NMU. Our proposed units NAU and NMU, compared with previous neural units, converge more consistently, have fewer parameters, learn faster, can converge for larger hidden sizes, obtain sparse and meaningful weights, and can extrapolate to negative and small values.

Motivation & Objective

• Address the poor extrapolation and instability of neural networks in arithmetic tasks due to lack of inductive bias for arithmetic operations.
• Improve upon the Neural Arithmetic Logic Unit (NALU) by redesigning its components to enhance stability, convergence, and interpretability.
• Develop a new multiplication unit (NMU) capable of handling large hidden sizes and both positive and negative values, which previous units could not.
• Enable reliable backpropagation through arithmetic layers in real-world tasks like MNIST digit sequence processing.
• Introduce a new success criterion based on empirical baselines to fairly evaluate arithmetic generalization beyond mean squared error.

Proposed method

• Propose the Neural Addition Unit (NAU) as a more stable and theoretically grounded alternative to the NALU's NAC₊, using a simplified parameter matrix and sparsity regularization.
• Introduce the Neural Multiplication Unit (NMU), which explicitly multiplies selected vector elements using a differentiable gating mechanism and optimal initialization.
• Apply a sparsity regularizer (Rz) that penalizes non-zero weights in the NAU to encourage meaningful, sparse weight patterns.
• Use a success criterion based on the 1% upper confidence interval of a baseline model to evaluate generalization, avoiding misleading MSE metrics.
• Optimize the NMU with careful initialization and weight constraints to improve gradient flow and prevent vanishing/exploding gradients.
• Evaluate models on both synthetic arithmetic tasks and real-world MNIST-based sequence tasks to test extrapolation and robustness.

Experimental results

Research questions

• RQ1 Can a neural network component be designed to reliably learn and extrapolate exact arithmetic operations like addition and multiplication?
• RQ2 How does the performance of the NAU and NMU compare to the NALU and NAC₊ in terms of convergence speed, success rate, and generalization to out-of-distribution values?
• RQ3 What role does sparsity regularization and proper initialization play in stabilizing training of arithmetic neural units?
• RQ4 Can the NMU effectively handle large hidden sizes and negative numbers, which previous units struggle with?
• RQ5 Does the proposed success criterion provide a more reliable evaluation of arithmetic generalization than standard MSE loss?

Key findings

• The NMU achieved a 94% success rate on the complex function task $t = (x_1 + x_2) imes (x_1 + x_2 + x_3 + x_4)$, compared to only 13% for NAC• and 26% for NALU.
• The NAU converged in a median of 14,000 iterations, significantly faster than NALU (78,000 iterations) and NAC• (59,000 iterations).
• The NMU achieved a sparsity error of $2.6 \times 10^{-8}$, indicating highly sparse and interpretable weights, compared to $7.5 \times 10^{-6}$ for NAC•.
• Removing the Rz regularizer reduced the NAU's success rate, especially on long sequences, showing its critical role in preventing mode collapse.
• The NMU successfully extrapolated to sequences of up to 1000 MNIST digits in the sequential multiplication task, demonstrating robust generalization.
• The NMU outperformed the NAC• in both success rate and convergence speed on the MNIST multiplication variant, confirming its superiority in real-world settings.