Abrupt Learning in Transformers: A Case Study on Matrix Completion

We thank the reviewer for their comments, and are excited to know that they found our approach interesting!

• The paper lacks a deep theoretical ...

We would like to point out that apart from the analysis in our paper, to the best of our knowledge there has been no work on analyzing such an algorithmic shift either mechanistically or theoretically. Ours is the first step in this direction where we start by formulating a mathematical setup, since this gives us a clearer picture of mechanisms learnt by the model before and after the algorithmic shift. We provide several detailed experiments on this front, as seen in Sections 4–6 of the paper.

To address reviewer concerns, we have also added 2 additional probing experiments that analyze the hidden states in the model in more detail than our existing probing experiment. To understand why the model undergoes algorithmic shift, we have now additionally added interventional experiments wherein various combinations of model components are frozen to their final parameter states (i.e., from the final step of training). These results further indicate rapid structural changes in attention layers drive sudden loss drops. We have also added several other experiments: e.g., OOD generalization experiments, training on different matrix distributions, scaling model width / depth. All these results further demonstrate the generalizability and robustness of our results. Please see our global response for more details on these experiments.

• The experiments are conducted ...

Please see our global response on “Larger matrices’, where we discuss possible ways to scale up our approach to larger input matrices.

• Does the model operate …

Our results show that this is most likely the case, since the mechanisms for non-MASK and MASK prediction remain distinct after the model has converged (copying observed entries vs. computation). We demonstrate this empirically through intervention experiments on the model weights in Sec 5.1 for observed (non-MASK) entries, and in Sec 5.2 for MASK entries. One of the main takeaways is that prediction at observed entries is not affected significantly by attention layers (since uniform ablating them does not lead to a drastic change in performance), whereas prediction at masked entries is significantly affected. Please see Sec 5 for more details about these experiments.

• What would be the model's behavior …

Please refer to Fig. K in the attached PDF; we find that freezing only the attention layer to weights obtained at the last training step and training only MLPs and embeddings (labeled 'mlp+embed’ in the plot) still leads to a plateau and then sudden drop. However, on freezing both Attention and Embedding layers, and training only MLP layers (labeled 'mlp’), we do not see plateau or sudden drop in loss, and the training loss decreases in a continuous manner to optimal value.

Our hypothesis is that the sudden-ness of the drop is contributed to by the Attention / Embedding layers, due to these components encoding ‘structure’ of the matrix completion problem. To corroborate this hypothesis, we also do the following experiments. Below, training all components is labeled `all’, which is the loss plot in the submitted manuscript – it is a reference for visualizing “sudden-ness” of drop in loss, and the length of loss plateau.

Train only Attention layers (labeled ‘att’) – smaller loss plateau than “all” with sharp drop

Train only Embeddings (both positional and token, labeled ‘embed’ in the plot) - smaller loss plateau than “all” with sharp drop

Train only positional embeddings (labeled ‘pos’) – this has the longest plateau of all settings

Train only token embeddings (labeled ‘token’) – there is no observable plateau or sudden drop in loss. Please note that this also matches the observation in our submission (Section 6, “Do embeddings change abruptly?”) where we have plotted the projection of token embeddings on the principal components of the final model’s token embeddings. We observed that the structure of this 2D projection does not change abruptly at step 15000, and that is further supported by this modified training process, which shows that token embeddings are learnt by the model without any significant observable plateau or sudden drop. Whether this early learning of token embeddings is also a driving force behind the sudden drop, is still an open question in our setup.

From these observations, a mechanistic hypothesis emerges: positional embedding layers and attention layers are the main contributors towards the sudden drop dynamics of the training loss. Formalizing and rigorously verifying these observations is an important direction for future research.

• How well does the model generalize to different types …

We find that the model performs well on test data with matrix rank smaller than what it was trained on – hence indicating that the model can find the lowest rank matrix that fits the observed entries. Also please see Fig. A in the attached PDF for results on different (OOD) rank levels, and Fig. 3 in the manuscript for OOD sparsity levels (larger p_mask denotes larger sparsity). We have compared our model in these settings to the nuclear norm minimization solution, to highlight that the BERT model is comparable to, or better than this solution. Please also see the section “OOD performance” in our global response for more information.

• Can the approach be scaled …

Please see section “Larger matrices” in our global response for addressing this concern.

We hope that our responses have justifiably addressed the reviewer’s concerns and they will consider increasing their score to support the acceptance of our work.

We thank the reviewer for their feedback, and are glad that they found our problem formulation and analysis interesting!

• I think the related work could be improved. I think the paper would benefit from a deeper discussion on the experiments and findings of relevant works cited (e.g., 4, 6, 20)

We are certainly happy to expand on the related work! We will also add the following addendum that discusses references [4, 6, 20] in more detail to the manuscript.

[4] show that in an in–context learning framework, a transformer based model can learn to select the most optimal statistical method for solving the task in the prompt, without explicitly being provided any information about the optimal method (called `in–context algorithm selection’ in their work). We emphasize that our setup is not in–context learning, and is quite distinct from [4] as far as the task being solved is concerned. However, whether the framework of layer-wise in-context gradient descent as in this work can also be shown in our setup is a plausible and open direction for future work (we provide some preliminary investigation in the Probing section in Global response).

In [6], the author shows empirically that a encoder-decoder transformer can be trained to solve various linear algebraic tasks, such as eigendecomposition, SVD, matrix inverse etc. They support their findings by showing that the model generalizes to matrix distributions outside the training distribution to some extent, and that OOD performance can be improved by training on non-Wigner matrices. While many experiments in [6] also show a sudden jump in accuracy (Fig 1,2), they do not analyze why such a sudden jump occurs during optimization. In our work, we analyze the sudden drop and the model before and after it to derive insights into the sudden drop in loss in our training setup.

[20] show that even a small transformer model can be trained to perform arithmetic operations like add, subtract, multiply accurately through appropriate data selection and formatting, and using Chain-of-Thought prompting. They further show that learning addition is connected to rank-2 matrix completion, and that the sudden jump in accuracy with increasing number of observed entries of the matrix is recovered when their model is trained on datasets of different sizes. This is because the size of the dataset for addition can be seen as the number of observed entries of the rank-2 matrix representing the addition table. We point out that while the task in this case is related to matrix completion, ours is a completely different setup, where the sudden drop happens with the number of training steps with each step consisting of 256 low-rank matrices, each with a fixed fraction (p_mask) of observed entries.

• The experimental setting is limited. Synthetic data generation is only carried out using only the uniform distribution. What about other distributions? What about OOD performance?

Thanks for this suggestion! We have now conducted new experiments to address your comment.

Changed training distribution. We now use a standard normal distribution to sample inputs X i.e. X_ij ~ N(0, 1). We find that a sudden drop in the loss occurs in this setting as well, and the MSE before and after the drop matches the values obtained for U[-1, 1] (Fig. D in attached PDF).

Changed test distribution. Replacing the distribution of input matrix entries from U[-1, 1] (on which the model is trained) to OOD entries from a Gaussian (mean=0, stdev=0.5) and Laplace distribution (location=0, scale=0.5) does not lead to considerable changes in mean-squared-error (MSE). Specifically, the Gaussian entries yield an MSE of 4e-3 and the Laplace entries yield an MSE of 2e-3 averaged over 1024 samples of 7x7 rank-2 matrix inputs. These results indicate the model has learned a general algorithm to solve the task that is robust to changes in the inputs’ distribution at test time.

Other interventions. While we already analyzed OOD generalization performance when varying p_mask (Fig 3 in paper), we have now also conducted experiments by changing the number of rows / columns and rank of the input matrix at test time – see Fig. (A–C) in attached PDF. We find performance is either comparable or better than nuclear norm minimization on the same input, except for the case when we modify the number of columns. This is likely because positional embeddings in the model depend on the column index of the element, and hence (expectedly) changing the number of columns adversely affects model performance.

• Did you consider other distributions?

Please see our response to the previous question.

We hope that our responses have justifiably addressed the reviewer’s concerns and they will consider increasing their score to support the acceptance of our work.

We thank the reviewer for their extensive suggestions, and are happy to know that they found our work important and interesting!

• How does convergence of BERT …

Thanks for this suggestion! We have now performed experiments with a 4-layer, 8-head model on rank–1 matrices of size 5x5, 7x7, 9x9 (at each training step, the size is sampled randomly, and then the training data consists of 256 matrices of that size). Interestingly, we find that it matches the intuitive learning order i.e. the sudden drop for 5x5 matrices occurs slightly before that for 7x7 matrices, which in turn is before that for 9x9 matrices (Fig. L in attached PDF). Further, the model performs similarly for 5x5 and 7x7, but is slightly worse for 9x9.

For the question about the effect of rank on convergence, we train our 4 layer 8 head model on 7x7 rank-1 inputs (lower than rank-2 in the manuscript), and find that the sudden drop occurs earlier in this case (Fig. M in attachment). Hence this is an indication that the problem structure affects the rate of convergence, however we leave verifying this for larger matrices and rank to future work.

Finally, in our experiments in the submission, the 12-layer, 12-head model was able to converge on matrices of size upto 15x15, rank–4 (please see Fig. 14 in manuscript for attention maps, and Fig. E in PDF attachment for convergence plot). We find that going larger than a certain rank for a specific matrix size does not lead to convergence in the usual training setup we report in the paper. This further reinforces the indication above that problem structure affects convergence. Due to limited time for rebuttal and limited computational resources, we are unable to extensively check the largest rank for larger matrix sizes for which our method converges.

(Note that we have modified the input distribution to U[-1.2, 1.2] in rank-1 to match the loss magnitude rank-2 case before sudden drop (~0.06-0.07) for meaningful comparison.)

• An interesting ablation would be to check …

Thanks for this suggestion! We have now carried out ablations on model size / architecture to empirically test how it affects convergence. We change model width and depth while keeping other hyperparams the same as the model analyzed in the submitted draft (4–layer, 8 heads, width=768). We also analyze a 1–head, 12–layer model to check for convergence rate when attention layers are altered (please note that 1 head models generally do not converge, likely due to insufficient capacity, but do so for the 12-layer case only). Broadly, our observations are as follows.

Width (Embedding size): we changed the hidden state (embedding) size in the model to 64, 128, 256, 1024, keeping number of layers fixed at 4, and number of heads fixed at 8. For each value of width w, the MLP hidden layer width is 4.w as used in the original w=768 case. We observe (Fig. F in attached PDF) that for embedding size less than 256, the model could not converge to the optimal MSE attained by the larger models.

Depth (Number of layers): we also train a model with depth 2, 6, 8, 12, keeping the number of heads fixed at 8, and width fixed at 768 (Fig. G in attachment; L = number of layers). For depth 2, the model did not converge to optimal MSE and is omitted. We note that for larger models (6, 8, 12), the sudden drop happens earlier than that for depth 4, likely due to larger model capacity in those cases.

1–head model: Using a 12–layer, width=768, 1–head model, we find the training converges to MSE ~ 4e-3, with a sudden drop in loss at around step 8000 (Fig J). Interestingly, we find that the attention heads obtained after training (Fig. N; layer 1-12 from left to right) are quite similar to those in 4-layer 8-head models. Possibly, the model exploits residual connections to simulate parallel attention heads within a single layer!

• In the experiments section ...

Please see section “Structure mask” section in our global response for a response to this concern.

• In the "Probing" experiments ... This check is for the full masked row corresponding to the element being probed, and not only the row or column number. Specifically, for 768-dimensional hidden state at element (i,j), we map it through a linear model to the 7-dimensional vector X_mask[i, :] where X_mask = X at observed entries, and 0 at masked entries. We hypothesize that this strong correlation indicates that the model stores the masked input in some form in the hidden states, since this information can be probed through a simple linear map. However, we have not claimed that this is the “only” component for the model’s computation.

For probing results on singular vectors (i.e. low rank decomposition), please see section “Probing” in global response.

• Any discussion on how …

Great question! We have now carried out preliminary experiments training a GPT model on the matrix completion task. Specifically, the input is now a concatenation of (X_mask, [SEP], X), and the objective is next–token prediction using cross–entropy loss. Due to the changed task structure, we now measure accuracy (instead of MSE) of output tokens at masked, observed, and all entries. Similar to the MLM setting of our main paper, we find that there is an initial plateau in loss (and accuracy), followed by a point of sudden drop in loss! However, the model merely learns to copy the observed entries (Fig. Q in attachment) at this point, and is still struggling to fill in the missing entries. We observe that even this sudden drop corresponds to the model learning specific structure in attention heads for copying (Fig. R in attachment, left: step 200, right: step 600). We are confident that with some hyperparameter and model size tuning, the model will be able to completely solve the task. Due to the limited timeframe for rebuttal, we were unable to finish these experiments, but we promise to include them in the final version of the paper.

We thank the reviewer for their comments, and are glad to know they found our results interesting!

• The main weakness is how significant the contribution is. They have some interesting results … how the transformer does math.

We re–emphasize that our focus is studying the sudden drop in loss from interpretability perspective, and not solving matrix completion itself. We did cite papers on using transformers for mathematical tasks, since our task is also mathematical in nature and wanted to give a review of existing work in this area to the reader. We will definitely consider improving this section to better match the focus of our paper.

We would like to humbly disagree about the significance of our contribution – we emphasize that there are currently no explanations for the phenomenon of sudden loss drops in language modeling. To address this problem, our simplified but rich abstraction offers two concrete observations: (1) the optimization of the model on this task shows an abrupt jump in the model performance and (2) the sudden jump can be analyzed via interpretable evidence that it has learnt the problem structure. These observations allow us to implicate sudden acquisition of task structure by attention layers as the cause behind sudden loss drops. To the best of our knowledge, both our observations and the proposed hypothesis are novel results.

Furthermore, we note that the sudden drop observed in our setting occurs without any changes to the optimization procedure (e.g. step-size, warmup etc) during training – prior work studying such loss drops [1] uses learning rate warmup for 10K steps, while we are able to obtain sudden drop in loss without warmup, and hence have arguably demonstrated a stronger empirical result.

[1] https://arxiv.org/abs/2309.07311

• Does this behavior generally to other mathematical tasks?

It does! To address your question, we have now extended our experiments to another mathematical task: histogram computation. Specifically, we follow to setup proposed by [2] and train a 2-layer, 2-head GPT model on sequences of the form (x_1, x_2, …, x_32, [SEP] c_1, c_2, …, c_32), where (x_i) are integers in range [1, 16] and (c_i) is the total number of occurrences of x_i in the full sequence. We use such sequences in a next-token prediction (online) training setup for our GPT model, using cross–entropy loss. At test time we prompt the model with incomplete sequences of the form (x_1, x_2, …, x_32, [SEP]) and expect it to complete the sequence with the correct sequence (c_1, c_2, …, c_32).

We find that our observations in matrix completion extend to this setup (Fig. O, P in attached PDF). Specifically, there is a sudden drop in training loss, and corresponding increase in accuracy of the model in predicting c_i (Fig. O), and attention heads transition from a nearly uniform pattern + attending to [SEP] (step 7000, Fig P - left) to a pattern that attends to various x_i for predicting counts (step 30000, Fig P - right) – as expected for the task.

This shows that such sudden drop in loss, and corresponding emergence of structure in attention heads is not limited to matrix completion and BERT architecture, but generalizes to a different mathematical task of histogram and on a different model (GPT) with autoregressive training, rather than masked language modeling. We are optimistic that such behavior would also extend to other mathematical tasks. We will add these results to the final version of the paper.

We also remark that some experiments in (Fig 1,2) [3] show a sudden increase in accuracy during training, however they do not analyze why that sudden increase occurs, and that is the main focus of our work.

[2] https://arxiv.org/abs/2402.03902 [3] https://openreview.net/forum?id=L2a_bcarHcF

• Can we get any insight about the actually algorithm or implicit optimization procedure?

As we show in Section 5, the algorithm learnt by the model shows many interpretable characteristics, indicating that the model learns (1) the problem structure through token and positional embeddings (Fig 7), and (2) how to combine elements at various positions through attention heads (Fig 4,5). Moreover, in our probing experiments (Fig. 6), we find that the model 'stores’ the masked rows of each element in the 3rd and 4th layer, giving insight into the algorithm used by the model. Please also see the section “Probing” in our global response for additional experiments on probing, including a possible hypothesis for implicit optimization in our setup.

Overall, we note that despite our best efforts, we could not precisely nail down the algorithm learned by the model. However, we emphasize this characterization was not the main objective of our paper; instead, we aimed to capture and understand sudden loss drops in the model’s learning dynamics. We did investigate if naively expected hypotheses for how the model might be solving the task check out, e.g., whether the model implicitly learns to perform nuclear norm minimization. Our results were however in the negative, as reported in Fig 3 of the main paper: we found that the model in fact outperforms nuclear norm minimization in terms of MSE.

• What is the generalization to larger matrices?

Please see section “Larger matrices” in global response.

• Is there a way to more rigorously verify the hypotheses from the Attention Heads with Structured Mask section?

Please see section “Structured masks” in our global response – we find that uniform ablation on each category of attention heads leads to increase in L_mask to different extents, largest being for the observed–only heads i.e. (2,2–4), (2,6), (3,2), (3,3), (3,5), (3,1), (3,6).

We hope that our responses have justifiably addressed the reviewer’s concerns and they will consider increasing their score to support the acceptance of our work.

We thank the reviewer for their insightful comments, and are happy to know they found our work interesting!

• Though using matrix completion as an "abstraction" of MLM is interesting and novel … interpretation and insights.

We emphasize that our goal while casting matrix completion as an abstraction of MLM was to define a simple, controllable system that captures the phenomenon that we aim to investigate: sudden drops in the loss. First reported by Chen et al. [1] in BERT training, the complexity of natural language disallowed authors to develop precise mechanistic hypotheses for what factors lead to these sudden drops. However, the simplicity of our setting affords us the ability to investigate model internals, showing that the model not only undergoes a change in the algorithm used to solve the given task, but also completely morphs at a mechanistic level (e.g., attention heads specialize to inferring structures relevant to the task). We hypothesize such rapid structure acquisition, specifically driven by attention, drives sudden loss drops. While defining this hypothesis was the goal of this work, given natural language and related practical domains such as code and mathematical problems are filled with rich, structured patterns (e.g., syntactical ones), we expect our hypothesis to hold true therein too. We plan to investigate this more concretely in future work.

We will make sure to add the discussion above to the paper and better clarify why our setup is a faithful abstraction and why we expect our results to transfer across domains.

[1] https://arxiv.org/abs/2309.07311

• Despite adapting a new task, the main patterns the paper found in through experiments are similar to prior works.

We emphasize that there are currently no explanations for the phenomenon of sudden loss drops in language modeling. To address this problem, our simplified but rich abstraction offers two concrete observations: (1) the optimization of the model on this task shows an abrupt jump in the model performance and (2) the sudden jump can be analyzed via interpretable evidence that it has learnt the problem structure. These observations allow us to implicate sudden acquisition of task structure by attention layers as the cause behind sudden loss drops. To the best of our knowledge, both our observations and the proposed hypothesis are novel results.

Furthermore, we note that the sudden drop observed in our setting occurs without any changes to the optimization procedure (e.g. step-size, warmup etc) during training – prior work studying such loss drops [1] uses learning rate warmup for 10K steps, while we are able to obtain sudden drop in loss without warmup, and hence have arguably demonstrated a stronger empirical result.

[1] https://arxiv.org/abs/2309.07311

• What is the purpose of training the model to predict all values instead of just the masked ones?

We find that it is difficult to train the model only to predict the masked entries, and predicting all entries substantially aids the training process. We hypothesize that learning to copy the observed entries is helpful for learning to complete the masked entries. We will add this observation to the paper.

• What do the two panels in figure represent? The right panel caption appears incorrect and the loss surprisingly goes down as p_mask increases.

We note that the caption is indeed correct. Specifically, the right panel in Fig. 3 compares the nuclear norm of the solutions obtained by BERT and nuclear-norm minimization – in this case, as p_mask increases, MSE displayed in the left panel goes up, as expected. We apologize for the lack of clarity though, and promise to make the caption clearer in the final version of the paper.

• Is there any consideration for tokenizing the matrix entries, besides mimicking language sequences? Since transformers are … if the same patterns hold for continuous inputs.

Indeed, the motivation behind tokenization is to keep the setup similar to language sequences, where BERT has been shown to perform well empirically. Importantly, as we show in Fig. 7, the model can learn the token embeddings representing the real values with expected structure (principal components separated by sign of the input, and continuously varying with magnitude of input). Hence it should not significantly affect the results if we use specialized embeddings instead of our approach; extending our work to other embedding methods is an interesting avenue for future work.