Title: NEURAL ARITHMETIC UNITS

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. 1 

1. INTRODUCTION:

When studying intelligence, insects, reptiles, and humans have been found to possess neurons with the capacity to hold integers, real numbers, and perform arithmetic operations (Nieder, 2016; Rugani et al., 2009; Gallistel, 2018) . In our quest to mimic intelligence, we have put much faith in neural networks, which in turn has provided unparalleled and often superhuman performance in tasks requiring high cognitive abilities (Silver et al., 2016; Devlin et al., 2018; OpenAI et al., 2018) . However, when using neural networks to solve simple arithmetic problems, such as counting, multiplication, or comparison, they systematically fail to extrapolate onto unseen ranges (Lake & Baroni, 2018; Suzgun et al., 2019; Trask et al., 2018) . The absence of inductive bias makes it difficult for neural networks to extrapolate well on arithmetic tasks as they lack the underlying logic to represent the required operations.

A neural component that can solve arithmetic problems should be able to: take an arbitrary hidden input, learn to select the appropriate elements, and apply the desired arithmetic operation. A recent attempt to achieve this goal is the Neural Arithmetic Logic Unit (NALU) by Trask et al. (2018) .

The NALU models the inductive bias explicitly via two sub-units: the NAC + for addition/subtraction and the NAC • for multiplication/division. The sub-units are softly gated between, using a sigmoid function, to exclusively select one of the sub-units. However, we find that the soft gating-mechanism and the NAC • are fragile and hard to learn.

In this paper, we analyze and improve upon the NAC + and NAC • with respect to addition, subtraction, and multiplication. Our proposed improvements, namely the Neural Addition Unit (NAU) and Neural Multiplication Unit (NMU), are more theoretically founded and improve performance regarding stability, speed of convergence, and interpretability of weights. Most importantly, the NMU supports both negative and small numbers and a large hidden input-size, which is paramount as neural networks are overparameterized and hidden values are often unbounded.

The improvements, which are based on a theoretical analysis of the NALU and its components, are achieved by a simplification of the parameter matrix for a better gradient signal, a sparsity regularizer, and a new multiplication unit that can be optimally initialized. The NMU does not support division.

1-1-1-:

Figure 1 : Visualization of the NMU, where the weights (W i,j ) controls gating between 1 (identity) or x i , each intermediate result is then multiplied explicitly to form z j .

However, we find that the NAC • in practice also only supports multiplication and cannot learn division (theoretical analysis on division discussed in section 2.3).

To analyze the impact of each improvement, we introduce several variants of the NAC • . We find that allowing division makes optimization for multiplication harder, linear and regularized weights improve convergence, and the NMU way of multiplying is critical when increasing the hidden size.

Furthermore, we improve upon existing benchmarks in Trask et al. (2018) by expanding the "simple function task", expanding "MNIST Counting and Arithmetic Tasks" with a multiplicative task, and using an improved success-criterion Madsen & Johansen (2019) . This success-criterion is important because the arithmetic units are solving a logical problem. We propose the MNIST multiplication variant as we want to test the NMU's and NAC • 's ability to learn from real data and extrapolate.

1.1 LEARNING A 10 PARAMETER FUNCTION Consider the static function t = (x 1 + x 2 ) • (x 1 + x 2 + x 3 + x 4 ) for x ∈ R 4 . To illustrate the ability of NAC • , NALU, and our proposed NMU, we conduct 100 experiments for each model to learn this function. Table 1 shows that the NMU has a higher success rate and converges faster.

Table 1 : Comparison of the success-rate, when the model converged, and the sparsity error for all weight matrices, with 95% confidence interval on the t = (x 1 + x 2 ) • (x 1 + x 2 + x 3 + x 4 ) task. Each value is a summary of 100 different seeds. 5.5 • 10 4 5.9 • 10 4 +7.8•10 3 -6.6•10 3 7.5 • 10 -6 +2.0•10 -6

-2.0•10 -6 NALU 26% +9% -8%

7.0 • 10 4 7.8 • 10 4 +6.2•10 3 -8.6•10 3 9.2 • 10 -6 +1.7•10 -6

-1.7•10 -6 × NMU 94% +3% -6%

1.4 • 10 4 1.4 • 10 4 +2.2•10 2 -2.1•10 2 2.6 • 10 -8 +6.4•10 -9 -6.4•10 -9

2. INTRODUCING DIFFERENTIABLE BINARY ARITHMETIC OPERATIONS:

We define our problem as learning a set of static arithmetic operations between selected elements of a vector. E.g. for a vector x learn the function (x 5 + x 1 ) • x 7 . The approach taking in this paper is to develop a unit for addition/subtraction and a unit for multiplication, and then let each unit decide which inputs to include using backpropagation.

We develop these units by taking inspiration from a theoretical analysis of Neural Arithmetic Logic Unit (NALU) by Trask et al. (2018) .

2.1. INTRODUCING NALU:

The Neural Arithmetic Logic Unit (NALU) consists of two sub-units; the NAC + and NAC • . The sub-units represent either the {+, -} or the {×, ÷} operations. The NALU then assumes that either NAC + or NAC • will be selected exclusively, using a sigmoid gating-mechanism.

The NAC + and NAC • are defined accordingly,

EQUATION

where Ŵ, M ∈ R H ×H -1 are weight matrices and z h -1 is the input. The matrices are combined using a tanh-sigmoid transformation to bias the parameters towards a {-1, 0, 1} solution. Having {-1, 0, 1} allows NAC + to compute exact {+, -} operations between elements of a vector. The NAC • uses an exponential-log transformation for the {×, ÷} operations, which works within precision and for positive inputs only.

The NALU combines these units with a gating mechanism z = g NAC + + (1 -g) NAC • given g = σ(Gx). Thus allowing NALU to decide between all of {+, -, ×, ÷} using backpropagation.

2.2 WEIGHT MATRIX CONSTRUCTION AND THE NEURAL ADDITION UNIT Glorot & Bengio (2010) show that E[z h ] = 0 at initialization is a desired property, as it prevents an explosion of both the output and the gradients. To satisfy this property with

W h -1 ,h = tanh( Ŵh -1 ,h )σ( Mh -1 ,h ), an initialization must satisfy E[tanh( Ŵh -1 ,h )] = 0.

In NALU, this initialization is unbiased as it samples evenly between + and -, or × and ÷. Unfortunately, this initialization also causes the expectation of the gradient to become zero, as shown in ( 4).

E ∂L ∂ Mh -1 ,h = E ∂L ∂W h -1 ,h E tanh( Ŵh -1 ,h ) E σ ( Mh -1 ,h ) = 0 (4)

Besides the issue of initialization, our empirical analysis (table 2) shows that this weight construction (1) do not create the desired bias for {-1, 0, 1}. This bias is desired as it restricts the solution space to exact addition, and in section 2.5 also exact multiplication, which is an intrinsic property of an underlying arithmetic function. However, this bias does not necessarily restrict the output space as a plain linear transformation will always be able to scale values accordingly.

To solve these issues, we add a sparsifying regularizer to the loss function (L = L + λ sparse R ,sparse ) and use a simple linear construction, where

EQUATION

EQUATION

NAU :

z h = H -1 h -1 =1 W h ,h -1 z h -1 (7)

2.3. CHALLENGES OF DIVISION:

The NAC • , as formulated in equation 3, has the capability to compute exact multiplication and division, or more precisely multiplication of the inverse of elements from a vector, when a weight in

W h -1 ,h is -1.

However, this flexibility creates critical optimization challenges. By expanding the exp-logtransformation, NAC • can be expressed as

NAC • : z h = H -1 h -1 =1 (|z h -1 | + ) W h ,h -1 . (8) In equation (8), if |z h -1 | is near zero (E[z h -1 ]

= 0 is a desired property when initializing (Glorot & Bengio, 2010) ), W h -1 ,h is negative, and is small, then the output will explode. This issue is present even for a reasonably large value (such as = 0.1), and just a slightly negative W h -1 ,h , as visualized in figure 2 . Also note that the curvature can cause convergence to an unstable area.

This singularity issue in the optimization space also makes multiplication challenging, which further suggests that supporting division is undesirable. These observations are also found empirically in ]. The problem is (

x 1 + x 2 ) • (x 1 + x 2 + x 3 + x 4 ) for x = (1, 1.2, 1.8, 2

). The solution is w 1 = w 2 = 1 in (a), with many unstable alternatives.

2.4. INITIALIZATION OF NAC •:

Initialization is important for fast and consistent convergence. A desired property is that weights are initialized such that E[z h ] = 0 (Glorot & Bengio, 2010) . Using second order Taylor approximation and assuming all z h -1 are uncorrelated; the expectation of NAC • can be estimated as

EQUATION

As shown in equation 9, satisfying E[z h ] = 0 for NAC • is likely impossible. The variance cannot be input-independently initialized and is expected to explode (proofs in Appendix B.3).

2.5. THE NEURAL MULTIPLICATION UNIT:

To solve the the gradient and initialization challenges for NAC • we propose a new unit for multiplication: the Neural Multiplication Unit (NMU)

EQUATION

EQUATION

NMU :

EQUATION

The NMU is regularized similar to the NAU and has a multiplicative identity when W h -1 ,h = 0. The NMU does not support division by design. As opposed to the NAC • , the NMU can represent input of both negative and positive values and is not bounded, which allows the NMU to extrapolate to z h -1 that are negative or smaller than . Its gradients are derived in Appendix A.3.

2.6. MOMENTS AND INITIALIZATION:

The NAU is a linear layer and can be initialized using Glorot & Bengio (2010) . The NAC + unit can also achieve an ideal initialization, although it is less trivial (details in Appendix B.2).

The NMU is initialized with E[W h ,h -1 ] = 1 /2. Assuming all z h -1 are uncorrelated, and E[z h -1 ] = 0, which is the case for most neural units (Glorot & Bengio, 2010) , the expectation can be approximated to

EQUATION

which approaches zero for H -1 → ∞ (see Appendix B.4). The NMU can, assuming V ar[z h -1 ] = 1 and H -1 is large, be optimally initialized with V ar[W h -1 ,h ] = 1 4 (proof in Appendix B.4.3).

2.7. REGULARIZER SCALING:

We use the regularizer scaling as defined in ( 14). We motivate this by observing optimization consists of two parts: a warmup period, where W h -1 ,h should get close to the solution, unhindered by the sparsity regularizer, followed by a period where the solution is made sparse.

EQUATION

2.8. CHALLENGES OF GATING BETWEEN ADDITION AND MULTIPLICATION:

The purpose of the gating-mechanism is to select either NAC + or NAC • exclusively. This assumes that the correct sub-unit is selected by the NALU, since selecting the wrong sub-unit leaves no gradient signal for the correct sub-unit.

Empirically we find this assumption to be problematic. We observe that both sub-units converge at the beginning of training whereafter the gating-mechanism, seemingly random, converge towards either the addition or multiplication unit. Our study shows that gating behaves close to random for both NALU and a gated NMU/NAU variant. However, when the gate correctly selects multiplication our NMU converges much more consistently. We provide empirical analysis in Appendix C.5 for both NALU and a gated version of NAU/NMU.

As the output-size grows, randomly choosing the correct gating value becomes an exponential increasing problem. Because of these challenges we leave solving the issue of sparse gating for future work and focus on improving the sub-units NAC + and NAC • .

3. RELATED WORK:

Pure neural models using convolutions, gating, differentiable memory, and/or attention architectures have attempted to learn arithmetic tasks through backpropagation (Kaiser & Sutskever, 2016; Kalchbrenner et al., 2016; Graves et al., 2014; Freivalds & Liepins, 2017) . Some of these results have close to perfect extrapolation. However, the models are constrained to only work with well-defined arithmetic setups having no input redundancy, a single operation, and one-hot representations of numbers for input and output. Our proposed models does not have these restrictions.

The Neural Arithmetic Expression Calculator (Chen et al., 2018) can learn real number arithmetic by having neural network sub-components and repeatedly combine them through a memory-encoderdecoder architecture learned with hierarchical reinforcement learning. While this model has the ability to dynamically handle a larger variety of expressions compared to our solution they require an explicit definition of the operations, which we do not.

In our experiments, the NAU is used to do a subset-selection, which is then followed by either a summation or a multiplication. An alternative, fully differentiable version, is to use a gumbelsoftmax that can perform exact subset-selection (Xie & Ermon, 2019) . However, this is restricted to a predefined subset size, which is a strong assumption that our units are not limited by.

4. EXPERIMENTAL RESULTS:



4.1. ARITHMETIC DATASETS:

The arithmetic dataset is a replica of the "simple function task" by Trask et al. (2018) . The goal is to sum two random contiguous subsets of a vector and apply an arithmetic operation as defined in (15)

EQUATION

where n (default 100), U [r lower , r upper ] (interpolation default is U [1, 2] and extrapolation default is U [2, 6]), and other dataset parameters are used to assess learning capability (see details in Appendix C.1 and the effect of varying the parameters in Appendix C.4).

4.1.1. MODEL EVALUATION:

We define the success-criterion as a solution that is acceptably close to a perfect solution. To evaluate if a model instance solves the task consistently, we compare the MSE to a nearly-perfect solution on the extrapolation range over many seeds. If W 1 , W 2 defines the weights of the fitted model, W 1 is nearly-perfect, and W * 2 is perfect (example in equation 16), then the criteria for successful convergence is L W1,W2 < L W 1 ,W * 2 , measured on the extrapolation error, for = 10 -5 . We report a 95% confidence interval using a binomial distribution (Wilson, 1927) .

EQUATION

To measure the speed of convergence, we report the first iteration for which

L W1,W2 < L W 1 ,W * 2

is satisfied, with a 95% confidence interval calculated using a gamma distribution with maximum likelihood profiling. Only instances that solved the task are included.

We assume an approximate discrete solution with parameters close to {-1, 0, 1} is important for inferring exact arithmetic operations. To measure the sparsity, we introduce a sparsity error (defined in equation 17). Similar to the convergence metric, we only include model instances that did solve the task and report the 95% confidence interval, which is calculated using a beta distribution with maximum likelihood profiling.

E sparsity = max h -1 ,h min(|W h -1 ,h |, |1 -|W h -1 ,h ||) (17) 4.1.2 ARITHMETIC OPERATION COMPARISON

We compare models on different arithmetic operations • ∈ {+, -, ×}. The multiplication models, NMU and NAC • , have an addition unit first, either NAU or NAC + , followed by a multiplication unit. The addition/substraction models are two layers of the same unit. The NALU model consists of two NALU layers. See explicit definitions and regularization values in Appendix C.2.

Each experiment is trained for 5 • 10 6 iterations with early stopping by using the validation dataset, which is based on the interpolation range (details in Appendix C.2). The results are presented in table 2. For multiplication, the NMU succeeds more often and converges faster than the NAC • and NALU.

For addition and substraction, the NAU and NAC + has similar success-rate (100%), but the NAU is significantly faster at solving both of the the task. Moreover, the NAU reaches a significantly sparser solution than the NAC + . Interestingly, a linear model has a hard time solving subtraction. A more extensive comparison is included in Appendix C.7 and an ablation study is included in Appendix C.3.

4.1.3. EVALUATING THEORETICAL CLAIMS:

To validate our theoretical claim, that the NMU model works better than NAC • for a larger hidden input-size, we increase the hidden size of the network thereby adding redundant units. Redundant units are very common in neural networks, which are often overparameterized.

Additionally, the NMU model is, unlike the NAC • model, capable of supporting inputs that are both negative and positive. 2.8 • 10 6 3.0 • 10 6 +2.9•10 5 -2.4•10 5

5.8

• 10 -4 +4.8•10 -4 -2.6•10 -4 NALU 0% +4% -0% - - - × NMU 98% +1% -5%

1.4 • 10 6 1.5 • 10 6 +5.0•10 4   -6.6•10 4 4.2 • 10 -7 +2.9•10 -8 -2.9•10 -8

NAC + 100% +0% -4% 2.5 • 10 5 4.9 • 10 5 +5.2•10 4 -4.5•10 4 2.3 • 10 -1 +6.5•10 -3 -6.5•10 -3 Linear 100% +0% -4% 6.1 • 10 4 6.3 • 10 4 +2.5•10 3 -3.3•10 3 2.5 • 10 -1 +3.6•10 -4 -3.6•10 -4 NALU 14% +8% -5% 1.5 • 10 6 1.6 • 10 6 +3.8•10 5 -3.3•10 5 1.7 • 10 -1 +2.7•10 -2 -2.5•10 -2 + NAU 100% +0% -4% 1.8 • 10 4 3.9 • 10 5 +4.5•10 4 -3.7•10 4 3.2 • 10 -5 +1.3•10 -5 -1.3•10 -5 NAC + 100% +0% -4% 9.0 • 10 3 3.7 • 10 5 +3.8•10 4 -3.8•10 4 2.3 • 10 -1 +5.4•10 -3 -5.4•10 -3 Linear 7% +7% -4% 3.3 • 10 6 1.4 • 10 6 +7.0•10 5 -6.1•10 5 1.8 • 10 -1 +7.2•10 -2 -5.8•10 -2 NALU 14% +8% -5%

1.9 • 10 6 1.9 • 10 6 +4.4•10 5 -4.5•10 5

2.1 • 10 -1 +2.2•10 -2 -2.2•10 -2 - NAU 100% +0% -4% 5.0 • 10 3 1.6 • 10 5 +1.7•10 4 -1.6•10 4 6.6 • 10 -2 +2.5•10 -2 -1.9•10 -2

Finally, for a fair comparison we introduce two new units: A variant of NAC • , denoted NAC •,σ , that only supports multiplication by constraining the weights with W = σ( Ŵ ). And a variant, named NAC •,NMU , that uses clamped linear weights and sparsity regularization identically to the NMU.

Figure 3 shows that the NMU can handle a much larger hidden-size and negative inputs. Furthermore, results for NAC •,σ and NAC •,NMU validate that removing division and adding bias improves the success-rate, but are not enough when the hidden-size is large, as there is no ideal initialization. Interestingly, no models can learn U[1.1, 1.2], suggesting certain input ranges might be troublesome. q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Hidden size model q q q q q NAC •,NMU NAC •,σ NAC • NALU NMU q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Sparsity error Interpolation range Published as a conference paper at ICLR 2020

U [ -2 , -1 ] U [ -2 , 2 ] U [ 0 , 1 ] U [ 0 . 1 , 0 . 2 ] U [ 1 , 2 ] U [ 1 . 1 , 1 . 2 ] U [ 1 0 , 2 0 ] U [ -2 , -1 ] U [ -2 , 2 ] U [ 0 , 1 ] U [ 0 . 1 , 0 . 2 ] U [ 1 , 2 ] U [ 1 . 1 , 1 . 2 ] U [ 1 0 , 2 0 ] U [ -2 , -1 ] U [ -2 , 2 ] U [ 0 , 1 ] U [ 0 . 1 , 0 . 2 ] U [ 1 , 2 ] U [ 1 . 1 , 1 . 2 ] U [ 1 0 , 2 0 ]

model q q q q q NAC •,NMU NAC •,σ NAC • NALU NMU

4.2. PRODUCT OF SEQUENTIAL MNIST:

To investigate if a deep neural network can be optimized when backpropagating through an arithmetic unit, the arithmetic units are used as a recurrent-unit over a sequence of MNIST digits, where the target is to fit the cumulative product. This task is similar to "MNIST Counting and Arithmetic Tasks" in Trask et al. (2018) 2 , but uses multiplication rather than addition (addition is in Appendix D.2). Each model is trained on sequences of length 2 and tested on sequences of up to 20 MNIST digits.

We define the success-criterion by comparing the MSE of each model with a baseline model that has a correct solution for the arithmetic unit. If the MSE of each model is less than the upper 1% MSE-confidence-interval of the baseline model, then the model is considered successfully converged.

Sparsity and "solved at iteration step" is determined as described in experiment 4.1. The validation set is the last 5000 MNIST digits from the training set, which is used for early stopping.

In this experiment, we found that having an unconstrained "input-network" can cause the multiplication-units to learn an undesired solution, e.g. (0.1 • 81 + 1 -0.1) = 9. Such network do solve the problem but not in the intended way. To prevent this solution, we regularize the CNN output with

R z = 1 H -1 H H h H -1 h -1 (1 -W h -1 ,h ) • (1 -zh -1 ) 2 .

This regularizer is applied to the NMU and NAC •,NMU models. See Appendix D.4 for the results where this regularizer is not used.

Figure 4 shows that the NMU does not hinder learning a more complex neural network. Moreover, the NMU can extrapolate to much longer sequences than what it is trained on. q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Sparsity error Extrapolation length 

model q q q q q q NAC •,NMU NAC •,σ NAC • LSTM NALU NMU

5. CONCLUSION:

By including theoretical considerations, such as initialization, gradients, and sparsity, we have developed the Neural Multiplication Unit (NMU) and the Neural Addition Unit (NAU), which outperforms state-of-the-art models on established extrapolation and sequential tasks. Our models converge more consistently, faster, to an more interpretable solution, and supports all input ranges.

A natural next step would be to extend the NMU to support division and add gating between the NMU and NAU, to be comparable in theoretical features with NALU. However we find, both experimentally and theoretically, that learning division is impractical, because of the singularity when dividing by zero, and that a sigmoid-gate choosing between two functions with vastly different convergences properties, such as a multiplication unit and an addition unit, cannot be consistently learned.

Finally, when considering more than just two inputs to the multiplication unit, our model performs significantly better than previously proposed methods and their variations. The ability for a neural unit to consider more than two inputs is critical in neural networks which are often overparameterized.

A GRADIENT DERIVATIVES

A.1 WEIGHT MATRIX CONSTRUCTION:

For clarity the weight matrix construction is defined using scalar notation

EQUATION

The of the loss with respect to Ŵh ,h -1 and Mh ,h -1 is then derived using backpropagation.

∂L ∂ Ŵh:

EQUATION

As seen from this result, one only needs to consider

∂L ∂W h ,h -1

for NAC + and NAC • , as the gradient with respect to Ŵh ,h -1 and Mh ,h -1 is a multiplication on

∂L ∂W h ,h -1 . A.2 GRADIENT OF NAC •

The NAC • is defined using scalar notation.

EQUATION

)

The gradient of the loss with respect to W h ,h -1 can the be derived using backpropagation.

EQUATION

We now wish to derive the backpropagation term

δ h = ∂L ∂z h , because z h affects {z h +1 } H +1

h +1 =1 this becomes:

EQUATION

To make it easier to derive

∂z h +1 ∂z h

we re-express the z h as z h +1 .

EQUATION

The gradient of

∂z h +1 ∂z h

is then:

∂z h +1 ∂z h = exp H h =1 W h +1 ,h log(|z h | + ) W h +1 ,h ∂ log(|z h | + ) ∂z h = exp H h =1 W h +1 ,h log(|z h | + ) W h +1 ,h abs (z h ) |z h | + = z h +1 W h +1 ,h abs (z h ) |z h | + (24)

abs (z h ) is the gradient of the absolute function. In the paper we denote this as sign(z h ) for brevity. However, depending on the exact definition used there may be a difference for z h = 0, as abs (0) is undefined. In practicality this doesn't matter much though, although theoretically it does mean that the expectation of this is theoretically undefined when E[z h ] = 0.

A.3 GRADIENT OF NMU:

In scalar notation the NMU is defined as:

EQUATION

The gradient of the loss with respect to W h -1 ,h is fairly trivial. Note that every term but the one for h -1 , is just a constant with respect to W h -1 ,h . The product, except the term for h -1 can be expressed as

z h W h -1 ,h z h -1 +1-W h -1 ,h

. Using this fact, the gradient can be expressed as:

EQUATION

Similarly, the gradient ∂L ∂z h which is essential in backpropagation can equally easily be derived as: The desired properties for initialization are according to Glorot et al. (Glorot & Bengio, 2010) :

EQUATION

E[z h ] = 0 E ∂L ∂z h -1 = 0 V ar[z h ] = V ar z h -1 V ar ∂L ∂z h -1 = V ar ∂L ∂z h (28) B.1.2 INITIALIZATION FOR ADDITION

Glorot initialization can not be used for NAC + as W h -1 ,h is not sampled directly. Assuming that

Ŵh ,h -1 ∼ Uniform[-r, r] and Mh ,h -1 ∼ Uniform[-r, r],

then the variance can be derived (see proof in Appendix B.2) to be:

EQUATION

One can then solve for r, given the desired variance (

V ar[W h -1 ,h ] = 2 H -1 +H ) (Glorot & Bengio, 2010).

B.1.3 MOMENTS AND INITIALIZATION FOR MULTIPLICATION:

Using second order multivariate Taylor approximation and some assumptions of uncorrelated stochastic variables, the expectation and variance of the NAC • layer can be estimated to:

EQUATION

This is problematic because E[z h ] ≥ 1, and the variance explodes for

E[z h -1 ] = 0. E[z h -1 ] = 0

is normally a desired property (Glorot & Bengio, 2010) . The variance explodes for E[z h -1 ] = 0, and can thus not be initialized to anything meaningful.

Published as a conference paper at ICLR 2020 For our proposed NMU, the expectation and variance can be derived (see proof in Appendix B.4) using the same assumptions as before, although no Taylor approximation is required:

EQUATION

These expectations are better behaved. It is unlikely that the expectation of a multiplication unit can become zero, since the identity for multiplication is 1. However, for a large H -1 it will be near zero.

The variance is also better behaved, but do not provide a input-independent initialization strategy.

We propose initializing with V ar

[W h -1 ,h ] = 1 4 , as this is the solution to V ar[z h ] = V ar[z h -1 ] assuming V ar[z h -1 ] = 1

and a large H -1 (see proof in Appendix B.4.3). However, more exact solutions are possible if the input variance is known.

B.2 EXPECTATION AND VARIANCE FOR WEIGHT MATRIX CONSTRUCTION IN NAC LAYERS:

The weight matrix construction in NAC, is defined in scalar notation as:

EQUATION

Simplifying the notation of this, and re-expressing it using stochastic variables with uniform distributions this can be written as:

EQUATION

Since tanh( Ŵ ) is an odd-function and E[ Ŵ ] = 0, deriving the expectation E[W ] is trivial.

E[W ] = E[tanh( Ŵ )]E[σ( M )] = 0 • E[σ( M )] = 0 (34)

The variance is more complicated, however as Ŵ and M are independent, it can be simplified to:

Var[W ] = E[tanh( Ŵ ) 2 ]E[σ( M ) 2 ] -E[tanh( Ŵ )] 2 E[σ( M )] 2 = E[tanh( Ŵ ) 2 ]E[σ( M ) 2 ] (35)

These second moments can be analyzed independently. First for E[tanh( Ŵ ) 2 ]:

EQUATION

Then for E[tanh( M ) 2 ]:

E[σ( M ) 2 ] = ∞ -∞ σ(x) 2 f U [-r,r] (x) dx = 1 2r r -r σ(x) 2 dx = 1 2r r -tanh r 2 (37)

Which results in the variance:

EQUATION

B.3 EXPECTATION AND VARIANCE OF NAC •

B.3.1 FORWARD PASS:

Expectation Assuming that each z h -1 are uncorrelated, the expectation can be simplified to:

EQUATION

Here we define g as a non-linear transformation function of two independent stochastic variables:

EQUATION

We then apply second order Taylor approximation of f , around

EQUATION

Inserting the derivatives and computing the inner products yields:

E[f (z h -1 , W h ,h -1 )] ≈ (|E[z h -1 ]| + ) E[W h ,h -1 ] + 1 2 V ar[z h -1 ](|E[z h -1 ]| + ) E[W h ,h -1 ]-2 E[W h ,h -1 ](E[W h ,h -1 ] -1) + 1 2 V ar[W h ,h -1 ](|E[z h -1 ]| + ) E[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 = 1 + 1 2 V ar[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 (43)

This gives the final expectation:

E[z h ] = E g(z h -1 , W h ,h -1 ) H -1 ≈ 1 + 1 2 V ar[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 H -1 (44)

We evaluate the error of the approximation, where

W h ,h -1 ∼ U [-r w , r w ] and z h -1 ∼ U [0, r z ].

These distributions are what is used in the arithmetic dataset. The error is plotted in figure 5 . Variance The variance can be derived using the same assumptions as used in "expectation", that all z h -1 are uncorrelated.

V ar[z:

h ] = E[z 2 h ] -E[z h ] 2 = E   H -1 h -1 =1 (|z h -1 | + ) 2•W h ,h -1   -E   H -1 h -1 =1 (|z h -1 | + ) W h ,h -1   2 = E f (z h -1 , 2 • W h ,h -1 ) H -1 -E f (z h -1 , W h ,h -1 ) 2•H -1 (45)

We already have from the expectation result in (43) that:

E f (z h -1 , W h ,h -1 ) ≈ 1 + 1 2 V ar[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 (46)

By substitution of variable we have that:

E f (z h -1 , 2 • W h ,h -1 ) ≈ 1 + 1 2 V ar[2 • W h ,h -1 ] log(|E[z h -1 ]| + ) 2 ≈ 1 + 2 • V ar[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 (47)

This gives the variance:

V ar[z h ] = E g(z h -1 , 2 • W h ,h -1 ) H -1 -E f (z h -1 , W h ,h -1 ) 2•H -1 ≈ 1 + 2 • V ar[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 H -1 -1 + 1 2 • V ar[W h ,h -1 ] log(|E[z h -1 ]| + ) 2 2•H -1 (48) B.3.2 BACKWARD PASS

Expectation The expectation of the back-propagation term assuming that δ h +1 and ∂z h +1 ∂z h are mutually uncorrelated:

E[δ h ] = E   H +1 h +1 =1 δ h +1 ∂z h +1 ∂z h   ≈ H +1 E[δ h +1 ]E ∂z h +1 ∂z h (49)

Assuming that z h +1 , W h +1 ,h , and z h are uncorrelated:

E ∂z h +1 ∂z h ≈ E[z h +1 ]E[W h +1 ,h ]E abs (z h ) |z h | + = E[z h +1 ] • 0 • E abs (z h ) |z h | + = 0 (50)

Variance Deriving the variance is more complicated:

V ar ∂z h +1 ∂z h = V ar z h +1 W h +1 ,h abs (z h ) |z h | + (51)

Assuming again that z h +1 , W h +1 ,h , and z h are uncorrelated, and likewise for their second moment:

V ar ∂z h +1 ∂z h ≈ E[z 2 h +1 ]E[W 2 h +1 ,h ]E abs (z h ) |z h | + 2 -E[z h +1 ] 2 E[W h +1 ,h ] 2 E abs (z h ) |z h | + 2 = E[z 2 h +1 ]V ar[W h +1 ,h ]E abs (z h ) |z h | + 2 -E[z h +1 ] 2 • 0 • E abs (z h ) |z h | + 2 = E[z 2 h +1 ]V ar[W h +1 ,h ]E abs (z h ) |z h | + 2 (52)

Using Taylor approximation around E[z h ] we have:

EQUATION

Finally, by reusing the result for E[z 2 h ] from earlier the variance can be expressed as: Expectation Assuming that all z h -1 are independent:

EQUATION

EQUATION

Assuming that E[z h -1 ] = 0 which is a desired property and initializing E[W h -1 ,h ] = 1 /2, the expectation is:

EQUATION

Variance Reusing the result for the expectation, assuming again that all z h -1 are uncorrelated, and using the fact that W h -1 ,h is initially independent from z h -1 :

EQUATION

Assuming that E[z h -1 ] = 0, which is a desired property and initializing

EQUATION

B.4.2 BACKWARD PASS:

Expectation For the backward pass the expectation can, assuming that ∂L ∂z h and ∂z h ∂z h -1 are uncorrelated, be derived to:

EQUATION

Initializing

E[W h -1 ,h ] = 1 /2

, and inserting the result for the expectation

EQUATION

Assuming that E ∂L ∂z h = 0, which is a desired property (Glorot & Bengio, 2010) .

EQUATION

Variance For the variance of the backpropagation term, we assume that ∂L ∂z h is uncorrelated with

EQUATION

Assuming again that E ∂L ∂z h = 0, and reusing the result

EQUATION

Focusing now on V ar

∂z h ∂z h -1

, we have:

V ar ∂z h ∂z h -1 = E z h W h -1 ,h z h -1 + 1 -W h -1 ,h 2 E[W 2 h -1 ,h ] -E z h W h -1 ,h z h -1 + 1 -W h -1 ,h 2 E[W h -1 ,h ] 2 (64)

Inserting the result for the expectation E

z h W h -1 ,h z h -1 +1-W h -1 ,h

and Initializing again

EQUATION

Using the identity that

E[W 2 h -1 ,h ] = V ar[W h -1 ,h ] + E[W h -1 ,h

] 2 , and again using

EQUATION

To derive E

z h W h -1 ,h z h -1 +1-W h -1 ,h 2 the result for V ar[z h ] can be used, but for Ĥ -1 = H -1 -1, because there is one less term. Inserting E z h W h -1 ,h z h -1 +1-W h -1 ,h 2 = V ar[W h -1 ,h ] + 1 4 H -1 -1 V ar[z h -1 ] + 1

H -1 -1 , we have:

EQUATION

Inserting the result for V ar

∂z h ∂z h -1 into the result for V ar ∂L ∂z h -1 : V ar ∂L ∂z h -1 ≈ V ar ∂L ∂z h H 1 2 2•H -1 + V ar[W h -1 ,h ] + 1 4 H -1 V ar[z h -1 ] + 1 H -1 -1 - 1 2 2•H -1 = V ar ∂L ∂z h H • V ar[W h -1 ,h ] + 1 4 H -1 V ar[z h -1 ] + 1 H -1 -1 (68) B.4.3 INITIALIZATION The W h -1 ,h should be initialized with E[W h -1 ,h ] = 1 2

, in order to not bias towards inclusion or exclusion of z h -1 . Using the derived variance approximations (68), the variance should be according to the forward pass:

EQUATION

And according to the backward pass it should be:

EQUATION

Both criteria are dependent on the input variance. If the input variance is know then optimal initialization is possible. However, as this is often not the case one can perhaps assume that V ar[z h -1 ] = 1. This is not an unreasonable assumption in many cases, as there may either be a normalization layer somewhere or the input is normalized. If unit variance is assumed, the variance for the forward pass becomes:

EQUATION

And from the backward pass:

EQUATION

The variance requirement for both the forward and backward pass can be satisfied with

V ar[W h -1 ,h ] = 1 4 for a large H -1 .

C ARITHMETIC TASK:

The aim of the "Arithmetic task" is to directly test arithmetic models ability to extrapolate beyond the training range. Additionally, our generalized version provides a high degree of flexibility in how the input is shaped, sampled, and the problem complexity.

Our "arithmetic task" is identical to the "simple function task" in the NALU paper (Trask et al., 2018) . However, as they do not describe their setup in details, we use the setup from Madsen & Johansen (2019) , which provide Algorithm 3, an evaluation-criterion to if and when the model has converged, the sparsity error, as well as methods for computing confidence intervals for success-rate and the sparsity error. 

EQUATION

Algorithm 1 defines the exact procedure to generate the data, where an interpolation range will be used for training and validation and an extrapolation range will be used for testing. Default values are defined in table 3 . x ← UNIFORM(R lower , R upper , i) Sample i elements uniformly 3: 

k ← UNIFORM(0, 1 -2s -o) Sample offset 4: a ← SUM(x[ik : i(k + s)]) Create sum a from subset 5: b ← SUM(x[i(k + s -o) : i(k + 2s -0)])

NAC • NAC + NAC • - - - NAC •,σ NAC + NAC •,σ - - - NAC •,NMU NAC + NAC •,NMU 10 10 6 2 • 10 6 NAC + NAC + NAC + - - - NALU NALU NALU - - - Linear Linear Linear - - - ReLU ReLU - - - ReLU6 ReLU6 ReLU6 - - - C.3 ABLATION STUDY

To validate our model, we perform an ablation on the multiplication problem. Some noteworthy observations:

1. None of the W constraints, such as R sparse and clamping W to be in [0, 1], are necessary when the hidden size is just 2.

2. Removing the R sparse causes the NMU to immediately fail for larger hidden sizes. Doing this gives only a slightly slower convergence. Although, this can not be guaranteed in general, as the regularizer is omitted during the initial optimization.

4. Removing both constraints, gives a somewhat satisfying solution, but with a lower successrate, slower convergence, and higher sparsity error.

In conclusion both constraints are valuable, as they provide faster convergence and a sparser solution, but they are not critical to the success-rate of the NMU.

• • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • • •

Success rate Solved at iteration step Sparsity error 

C.4 EFFECT OF DATASET PARAMETER:

To stress test the models on the multiplication task, we vary the dataset parameters one at a time while keeping the others at their default value (default values in table 3 ). Each runs for 50 experiments with different seeds. The results, are visualized in figure 8 .

In figure 3 , the interpolation-range is changed, therefore the extrapolation-range needs to be changed such it doesn't overlap. For each interpolation-range the following extrapolationrange is used: [10, 20] uses U [20, 40] .

U[-2, -1] uses U[-6, -2], U[-2, 2] uses U[-6, -2] ∪ U[2, 6], U[0, 1] uses U[1, 5], U[0.1, 0.2] uses U[0.2, 2], U[1.1, 1.2] uses U[1.2, 6], U[1, 2] uses U[2, 6], U

q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q qq q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Sparsity error Size of the input vector model q q q q q q NAC •,NMU NAC •,σ NAC • Gated NAU/NMU NALU NMU q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Sparsity error 0.00 0.25 0.50 0.75 1.00 0.00 0.25 0.50 0.75 1.00 0.00 0.25 0.50 0.75 1.00 0.000 The ratio of which subsets overlap with each other model q q q q q q NAC •,NMU NAC •,σ NAC • Gated NAU/NMU NALU NMU q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Relative size of subsets compared to input size 

model q q q q q q NAC •,NMU NAC •,σ NAC • Gated NAU/NMU NALU NMU

C.5 GATING CONVERGENCE EXPERIMENT:

In the interest of adding some understand of what goes wrong in the NALU gate, and the shared weight choice that NALU employs to remedy this, we introduce the following experiment.

We train two models to fit the arithmetic task. Both uses the NAC + in the first layer and NALU in the second layer. The only difference is that one model shares the weight between NAC + and NAC • in the NALU, and the other treat them as two separate units with separate weights. In both cases NALU should gate between NAC + and NAC • and choose the appropriate operation. Note that this NALU model is different from the one presented elsewhere in this paper, including the original NALU paper (Trask et al., 2018) . The typical NALU model is just two NALU layers with shared weights.

Furthermore, we also introduce a new gated unit that simply gates between our proposed NMU and NAU, using the same sigmoid gating-mechanism as in the NALU. This combination is done with seperate weights, as NMU and NAU use different weight constrains and can therefore not be shared.

The models are trained and evaluated over 100 different seeds on the multiplication and addition task.

A histogram of the gate-value for all seeds is presented in figure 9 and table 5 contains a summary. Some noteworthy observations: 1. When the NALU weights are separated far more trials converge to select NAC + for both the addition and multiplication task. Sharing the weights between NAC + and NAC • makes the gating less likely to converge for addition. 2. The performance of the addition task is dependent on NALU selecting the right operation. In the multiplication task, when the right gate is selected, NAC • do not converge consistently, unlike our NMU that converges more consistently. 3. Which operation the gate converges to appears to be mostly random and independent of the task. These issues are caused by the sigmoid gating-mechanism and thus exists independent of the used sub-units.

These observations validates that the NALU gating-mechanism does not converge as intended. This becomes a critical issues when more gates are present, as is normally the case. E.g. when stacking multiple NALU layers together. 1.5 • 10 6 1.5 • 10 6 +3.9•10 4 The simple function task is a purely synthetic task, that does not require a deep network. As such it does not test if an arithmetic layer inhibits the networks ability to be optimized using gradient decent.

The sequential MNIST task takes the numerical value of a sequence of MNIST digits and applies a binary operation recursively. Such that t i = Op(t i-1 , z t ), where z t is the MNIST digit's numerical value. This is identical to the "MNIST Counting and Arithmetic Tasks" in Trask et al. (2018, section 4.2) . We present the addition variant to validate the NAU's ability to backpropagate, and we add an additional multiplication variant to validate the NMU's ability to backpropagate.

The performance of this task depends on the quality of the image-to-scalar network and the arithmetic layer's ability to model the scalar. We use mean-square-error (MSE) to evaluate joint image-to-scalar and arithmetic layer model performance. To determine an MSE threshold from the correct prediction we use an empirical baseline. This is done by letting the arithmetic layer be solved, such that only the image-to-scalar is learned. By learning this over multiple seeds an upper bound for an MSE threshold can be set. In our experiment we use the 1% one-sided upper confidence-interval, assuming a student-t distribution.

Similar to the simple function task we use a success-criteria as reporting the MSE is not interpretable and models that do not converge will obscure the mean. Furthermore, because the operation is applied recursively, natural error from the dataset will accumulate over time, thus exponentially increasing the MSE. Using a baseline model and reporting the successfulness solves this interpretation challenge.

D.2 ADDITION OF SEQUENTIAL MNIST:

Figure 13 shows results for sequential addition of MNIST digits. This experiment is identical to the MNIST Digit Addition Test from Trask et al. (2018, section 4.2) . The models are trained on a sequence of 10 digits and evaluated on sequences between 1 and 1000 MNIST digits.

Note that the NAU model includes the R z regularizer, similarly to the "Multiplication of sequential MNIST" experiment in section 4.2. However, because the weights are in [-1, 1], and not [0, 1], and the idendity of addition is 0, and not 1, R z is

EQUATION

)

To provide a fair comparison, a variant of NAC + that also uses this regularizer is included, this variant is called NAC +,Rz . Section D.3 provides an ablation study of the R z regularizer. q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q q Success rate Solved at iteration step Sparsity error Extrapolation length model q q q q q NAC +,Rz NAC + LSTM NALU NAU 

:

Implementation is available on GitHub: https://github.com/AndreasMadsen/stable-nalu.

The same CNN is used, https://github.com/pytorch/examples/tree/master/mnist.

Figures and Tables:

Figure 2: RMS loss curvature for a NAC + unit followed by a NAC • . The weight matrices are constrained to W 1 = w1 w1 0 0 w1 w1 w1 w1 , W 2 = [ w2 w2]. The problem is (x 1 + x 2 ) • (x 1 + x 2 + x 3 + x 4 ) for x = (1, 1.2, 1.8, 2). The solution is w 1 = w 2 = 1 in (a), with many unstable alternatives.

Figure 3: Multiplication task results when varying the hidden input-size and when varying the input-range. Extrapolation ranges are defined in Appendix C.4.

Figure 4: MNIST sequential multiplication task. Each model is trained on sequences of two digits, results are for extrapolating to longer sequences. Error-bars represent the 95% confidence interval.

Figure Error between theoretical approximation and the numerical approximation estimated by random sampling of 100000 observations at each combination of r z and r w .



Figure 6: Shows how the dataset is parameterized.

Removing the clamping of W does not cause much difference. This because R sparse also constrains W outside of [0, 1]. The regularizer used here is R sparse = min(|W |, |1 -W |), which is identical to the one used in other experiments in [0, 1], but is also valid outside [0, 1].

Figure 7: Ablation study where R sparse is removed and the clamping of W is removed. There are 50 experiments with different seeds, for each configuration.

Figure 8: Shows the effect of the dataset parameters.

Figure 9: Shows the gating-value in the NALU layer and a variant that uses NAU/NMU instead of NAC + /NAC • . Separate/shared refers to the weights in NAC + /NAC • used in NALU.

• 10 5 3.1 • 10 5 +4.3•10 4 -5.4•10 4 1.8 • 10 -1 +2.3•10 -2 -2.1•10 -2 D SEQUENTIAL MNIST D.1 TASK AND EVALUATION CRITERIA

Figure 13: Shows the ability of each model to learn the arithmetic operation of addition and backpropagate through the arithmetic layer in order to learn an image-to-scalar value for MNIST digits. The model is tested by extrapolating to larger sequence lengths than what it has been trained on. The NAU and NAC +,Rz models use the R z regularizer from section 4.2.

Comparison of: success-rate, first iteration reaching success, and sparsity error, all with 95% confidence interval on the "arithmetic datasets" task. Each value is a summary of 100 different seeds.

Table content: <table><tr><td>Op Model</td><td>Success</td><td colspan="2">Solved at iteration step</td><td>Sparsity error</td></tr><tr><td/><td>Rate</td><td>Median</td><td>Mean</td><td>Mean</td></tr><tr><td colspan="2">NAC • 31% +10% -8%</td><td/><td/></tr></table>

Default dataset parameters for "Arithmetic task"

Table content: <table><tr><td>Parameter name Default value</td><td>Parameter name Default value</td></tr><tr><td>Input size 100</td><td>Interpolation range U [1, 2]</td></tr><tr><td>Subset ratio 0.25</td><td>Extrapolation range U [2, 6]</td></tr><tr><td>Overlap ratio 0.5</td><td/></tr><tr><td colspan="2">Algorithm 1 Dataset generation algorithm for "Arithmetic task"</td></tr><tr><td colspan="2">1: function DATASET(OP(•, •) : Operation, i : InputSize, s : SubsetRatio, o : OverlapRatio,</td></tr><tr><td>R : Range)</td><td/></tr></table>

• 10 6 iterations. Training takes about 8 hours on a single CPU core(8-Core Intel Xeon E5-2665 2.4GHz). We run 19150 experiments on a HPC cluster.The training dataset is continuously sampled from the interpolation range where a different seed is used for each experiment, all experiments use a mini-batch size of 128 observations, a fixed validation dataset with 1 • 10 4 observations sampled from the interpolation range, and a fixed test dataset with 1 • 10 4 observations sampled from the extrapolation range.

Table content: <table><tr><td>Create sum b from subset Perform operation on a and b Models are defined in table 4 and are all optimized with Adam optimization (Kingma &amp; Ba, 2014) 6: t ← OP(a, b) 7: return x, t C.2 MODEL DEFINTIONS AND SETUP Model Layer 1 Layer 2 λsparse λ start λ end NMU NAU NMU 10 10 6 2 • 10 6 using default parameters, and trained over 5 Table 4: Model definitions NAU NAU NAU 0.01 5 • 10 3 5 • 10 4</td></tr></table>

Comparison of the success-rate, when the model converged, and the sparsity error, with 95% confidence interval on the "arithmetic datasets" task. Each value is a summary of 100 different seeds.

Table content: <table><tr><td>Op</td><td>Model</td><td>Success</td><td colspan="2">Solved at iteration step</td><td>Sparsity error</td></tr><tr><td/><td/><td>Rate</td><td>Median</td><td>Mean</td><td>Mean</td></tr><tr><td/><td colspan="2">Gated NAU/NMU 62% +9% -10%</td><td/><td/></tr></table>

Peer review (iclr2020):
The authors propose the Neural Multiplication Unit (NMU), which can learn to solve a family of arithmetic operations using -, + and * atomic operations over real numbers from examples. They show that a combination of careful initialization, regularization and structural choices allows their model to learn more reliably and efficiently than the previously published Neural Arithmetic Logic Unit. The NALU consists of two additive sub-units in the real and log-space respectively, which allows it to handle both additions/subtractions and multiplications/divisions, and combines them with a gating mechanism. The NMU on the other hand simply learns a product of affine transformations of the input. This choice prevents the model from learning divisions, which the authors argue made learning unstable for the NALU case, but allows for an a priori better initialization and dispenses with the gating which is empirically hard to learn. The departures from the NALU architecture are well justified and lead to significant improvements for the considered applications, especially as far as extrapolation to inputs outside of the training domain. The paper is mostly well written (one notable exception: the form of the loss function is not given explicitly anywhere in the paper) and well executed, but the scope of the work is somewhat limited, and the authors fail to properly motivate the application or put it in a wider context. First, divisions being difficult to handle does not constitute a sufficient justification for choosing to exclude them: the authors should at the very least propose a plausible way forward for future work. More generally, the proposed unit needs to be exposed to at least 10K examples to learn a single expression with fewer than 10 inputs (and the success rate already drops to under 65% for 10 inputs). What would be the use case for such a unit? Even the NMU is only proposed as a step on the way to a more modular, general-purpose, or efficient architecture, its value is difficult to gauge without some idea of what that would look like.