Tracing and Manipulating intermediate values in Neural Math Problem Solvers (2022)

Abstract: How language models process complex input that requires multiple steps of inference is not well understood. Previous research has shown that information about intermediate values of these inputs can be extracted from the activations of the models, but it is unclear where that information is encoded and whether that information is indeed used during inference. We introduce a method for analyzing how a Transformer model processes these inputs by focusing on simple arithmetic problems and their intermediate values. To trace where information about intermediate values is encoded, we measure the correlation between intermediate values and the activations of the model using principal component analysis (PCA). Then, we perform a causal intervention by manipulating model weights. This intervention shows that the weights identified via tracing are not merely correlated with intermediate values, but causally related to model predictions. Our findings show that the model has a locality to certain intermediate values, and this is useful for enhancing the interpretability of the models.

## introduction recent language models (lms) can solve complex input such as math word problems @xcite . to obtain the correct output from such complex (latent structured) inputs, it is necessary for multiple steps of inference via intermediate values. however, how lms process their inputs and capture latent structure is still not well understood. in previous studies, @xcite and @xcite showed that the neural models can capture some implicit hierarchical structure, but it is unclear where that information is encoded. @xcite observed that in lms trained with dyck language and showed some activations are highly correlated with the depth of their syntactic tree. however, even if such features can be extracted, there is no guarantee that it is used by the model @xcite . given these considerations, to better understand lm predictions for latent structured inputs, it is necessary to: (a) to find where information about intermediate values of the latent structured inputs is encoded. (b) to evaluate the impact of the features when the model makes predictions. in this work, we introduce a method for analyzing the relationship between internal representations in transformer @xcite -based models and intermediate values of latent structured inputs by using simple math problems. we choose them as a formal language because their intermediate values of the latent (tree) structure are clear and continuous, and it is easy to investigate their relationship to the internal representation of the model. the intermediate value of (154 -38) -(290 -67) can be clearly defined as 154, 290, 154 -38 =116, and so on. we take up a transformer model trained to solve math equations. an overview of our experiments is shown in fig. 1 . first, we search which directions of internal representations are highly correlated with intermediate values in equations by pca (tracing) to find where the information about intermediate values is encoded. we find some directions correlate very well with the intermediate values. second, we observe how the model prediction changed when we manipulate the weights along its direction (manipulation) to conduct a causal intervention. the result of this experiment suggests that some directions of them are indeed used by the model. these two results show that a transformer model has a locality to certain intermediate values, and it could help enhance the interpretability of the models. our contributions are as follows: (a)we show that intermediate values of equations are encoded in particular directions in internal representation. (b)we show that some features representing intermediate values are used during inference. 1 𝑎 -(𝑏 -𝑐) =? 105 -(293 -131) 203 -(610 -225) 567 -(661 -725) trained model ① pca ③ inference 568 -(215 -328) 15 10 5 0 5 10 15 20 weight of principal component p 3 2 0 200 400 600 800 1000 intermediate value: b |r|=0.973 2 related work intermediate values. previous work has examined the representation of intermediate values in neural models. linzen et al. (2016), bowman et al. @xcite and @xcite found that lms capture implicit hierarchical structures to some extent, e.g., when performing logical inference over formal languages. closest to this work are @xcite , who trained lms on the dyck language and observed hidden units that are highly correlated with nesting depth. in contrast to their work, we analyze representations of more complex inputs, i.e., equations, and also manipulate these representations to understand the impact of correlated activations on model predictions. @xcite have shown that they can reach the state-of-the-art performance of numerical reasoning by using large pre-trained lm. several studies have shown that a transformer model can solve more complex problems such as linear algebra and elementary mathematics to some extent @xcite . based on their findings, we use simple mathematical equations as problems that can be solved by a transformer model in this study. ## experiments we conduct two types of experiments. first we trace the representation of intermediate values in the model. as a result we find directions in activation space that are highly correlation with intermediate values. then we manipulate activations along these directions and observe if model predictions change as expected. as neural math problem solving model, we train a 6-layer transformer using the settings by sajjad et al. (2021) on synthetic data. we generate 200k equations involving up to five steps of addition or subtraction of integers between 1 and 1000, e.g., (154 -38) -(290 -67). following geva et al. (2020) inputs are split into digits, e.g., "123" is tokenized into 1, ##2, ##3. model predictions are obtained via linear regression on the final layer's [cls] token representation. after training on 190k equations we evaluate the model on 10k equations and obtain a regression score of r 2 = 0.9988, i.e., the model solves the equations almost perfectly.