Teaching Neural Module Networks to Do Arithmetic

Taking one module compare-date as a case study: it performs comparisons between two references queried by the question. A key reasoning step inside, is the find-date module that obtains appropriate a date token distribution $D$ related to each reference: $\verb|find-date|(P)\rightarrow D$. It is worth noting that there is no interaction with the question, which could contain essential information (e.g., entities) that is useful to correctly answer the question. Therefore, we revise the find-date module as follows: $\verb|find-date|(P,Q)\rightarrow D$:

where P and Q represent the contextualized embeddings of the paragraph and question, and $\textbf{P}_{d_{j}}$ of the $j^{th}$ date tokens in the paragraph, $\textbf{W}_{date}$ is a trainable parameter, $P,Q$ are the expected attention distribution of the paragraph and the question respectively.

In Equation 1, we concatenate the paragraph embeddings P and question embeddings Q that output from a pre-trained BERT Devlin et al. (2019) model to construct the context representation. A hyper-parameter $\alpha$ is used to adjust their contributions, whose value is empirically determined (Appendix A.1). The context representation is provided to compute the improved similarity matrix $\textbf{S}_{date}$. We concatenate the paragraph and question attention inputs in the same way to calculate the final expected distribution over the date tokens $D$ (Eq. 3). Now the interpreter is equipped with question information to make the prediction.

In the NMNs’ modelling paradigm, for addition/subtraction operations, the programmer takes as input two number distributions and produces an output number distribution over all possible result values: $\verb|add/sub|(N_{1},N_{2})\rightarrow RL$. $N_{1}$ and $N_{2}$ represent the probability distributions of the $1^{st}$ and $2^{nd}$ operands over all numbers that are extracted from the paragraph and collected into a sorted operand list $OL$. The positive and negative values of these numbers are exhaustively combined in pairs, from which the possible results of addition/subtraction operations are compiled into a sorted result list $RL$. For each input number distribution $N_{i},i=1,2$, a matrix $\textbf{C}_{i}\in\mathbb{R}^{m\times n}$ is constructed, where $m$ is the total number of possible results, and $n$ is the maximum number of unique combinations. Each value $\textbf{C}_{i}[j,k]$ is found by looking up the probability value in $N_{i}[k]$ where $OL[k]$ is the $i^{th}$ operand in any pair that produces result $RL[j]$. The probability that the $j^{th}$ number in $N_{i}$ is the correct operand of the $k^{th}$ pair. We compute the marginalized joint probability by summing over the product of $C_{i}$ as the expected distribution over result list $RL$. For the addition module, it is:

For instance, assume the sorted operand list $OL$ from a paragraph is [1, 5, 7, 11] and $N_{1}=[0.1,0.4,0.2,0.3]$. Different combinations are formed, e.g., (+$n_{1}$, +$n_{2}$) for addition and (+$n_{1}$, -$n_{2}$) for subtraction, and all possible results of the combinations are compiled into two result lists, one for addition and one for subtraction. For subtraction in this case, $RL=[0,2,4,6,10]$. The value of $\textbf{C}_{1}[2,1]$ is 0.4, which is found from $N_{1}[1]$ because the result 4 can be calculated from (+5, -1); and $\textbf{C}_{1}[2,3]=0.3$ which equals to $N_{1}[3]$ as 4 is the result of (+11, -7) as well. $\textbf{C}_{2}$ is computed in the same way to further obtain final distribution over $RL$.

We compose add/sub modules in programs to perform 3-number arithmetic. The key to our approach is to construct and distinguish appropriate $C_{i}$ and $RL$ in different reasoning steps. In the second arithmetic step, we should combine the operand list from the paragraph and the result list from the previous step to obtain a new result list $RL^{\prime}$, $\verb|add/sub|(RL,N)\rightarrow RL^{\prime}$. Due to the changes in operands and results, the modules should refer to a different $\textbf{C}_{i}^{\prime}\in\mathbb{R}^{m^{\prime}\times n^{\prime}}$ in the computation. We extend 2-number add/sub modules to recognize the participation of the third number by conditional statement, in order to differentiate the operand and result lists the interpreter should refer to in different steps. Taking the last example in Figure 1, the addition module would first compute the distribution over result list for ‘Albanian and Bulgarian citizens’. The subtraction module can identify itself in the second step calculation and take the correct input to construct the new matrix $\textbf{C}_{i}^{\prime}$. The expected distribution over new result list $RL^{\prime}$ now represent the difference of ‘Greek citizens’ and the previous result.

Instead of introducing specific modules for multi-number arithmetic such as ‘3-num-add’, the structure of NMNs allows us to recursively execute basic operations several times in a compositional program. This design is in accord with the reasoning process of the CQA task, and natural for NMNs to perform complex computations.