Abrupt Learning in Transformers: A Case Study on Matrix Completion
Authors: Pulkit Gopalani
Ekdeep Singh Lubana
Wei Hu
Venue: NeurIPS 2024 poster (2024)
Keywords: Science of language models
matrix completion
BERT
phase transition
interpretability
Primary area: interpretability_and_explainability
TL;DR: BERT solves low rank matrix completion in an interpretable manner with a sudden drop in the loss.

Abstract: Recent analysis on the training dynamics of Transformers has unveiled an interesting characteristic: the training loss plateaus for a significant number of training steps, and then suddenly (and sharply) drops to near--optimal values. To understand this phenomenon in depth, we formulate the low-rank matrix completion problem as a masked language modeling (MLM) task, and show that it is possible to train a BERT model to solve this task to low error. Furthermore, the loss curve shows a plateau early in training followed by a sudden drop to near-optimal values, despite no changes in the training procedure or hyper-parameters. To gain interpretability insights into this sudden drop, we examine the model's predictions, attention heads, and hidden states before and after this transition. Concretely, we observe that (a) the model transitions from simply copying the masked input to accurately predicting the masked entries; (b) the attention heads transition to interpretable patterns relevant to the task; and (c) the embeddings and hidden states encode information relevant to the problem. We also analyze the training dynamics of individual model components to understand the sudden drop in loss.

Abrupt Learning in Transformers: A Case Study on Matrix Completion
Pulkit Gopalani University of Michigan gopalani @umi ch. edu
Ekdeep Singh Lubana
Wei Hu University of Michigan vvh@umich edu
Harvard University
ekdeeplubana@fas. hww.aaaaddaaaa.aad.a.ao..mm.. harvard. edu
Abstract
Recent analysis on the training dynamics of Transformers has unveiled an inter- esting characteristic: the training loss plateaus for significant number of training steps, and then suddenly (and sharply) drops to near-optimal values. To understand this phenomenon in depth, we formulate the low-rank matrix completion problem as a masked language modeling (MLM) task, and show that it is possible to train a BERT model solve this task to low error. Furthermore, the loss curve shows a plateau early in training followed by a sudden drop to near-optimal values, despite changes in the training procedure or hyper-parameters To gain interpretability insights into this sudden drop, we examine the model's predictions, attention heads, and hidden states before and after this transition. Concretely, we observe that (a) the model transitions from simply copying the masked input to accurately predict- ing the masked entries; (b) the attention heads transition to interpretable patterns relevant to the task; and (c) the embeddings and hidden states encode information relevant to the problem. We also analyze the training dynamics of individual model components to understand the sudden drop in loss.
Introduction
Large Language Models (LLMs) have revolutionized the field of natural language processing (NLP). However, there are still gaps in our understanding of these models, leading challenges controlling their behavior. As pertinent example, the training of these models appears to demonstrate sudden improvements metrics correlated with various capabilities [8. prompting questions about whether learning of given capability can be predicted by tracking predefined progress measures and why such sudden changes occur. undesirable capabilities can suddenly 'emerge' (despite any explicit supervision for them) [66. such sudden changes can be challenge for AI regulation [21]]
To better understand such sudden changes during model training, this work investigates training BERT [12) on the classical mathematical task of low-rank matrix completion (LRMC) (6). Making an analogy with masked language modeling (MLM), where sudden learning syntactical structures was recently demonstrated [8), we argue matrix completion captures the core aspect of this learning problem (Fig. I: given some relevant context (observed tokens). fill the missing elements (masked tokens). Specifically, we assume access to a matrix with some fraction of its entries missing, and would like complete the missing entries of this matrix assuming the ground truth matrix is low-rank. We find that despite being a simplified abstraction of MLM, this setting already demonstrates sharp decrease loss as the model undergoes training (Fig. [ (B)), preceded by a loss plateau for
significant number of training steps (akin to Chen et al. g ). The simplicity of our setting further
affords us interpretability, as we find that the point of sudden drop coincides with a precise change in show
how the model solves the task-we call this change algorithmic transition. Specifically, that the pre-transition model simply copies the input (predicting 0 at masked positions), while the post-transition model accurately predicts missing values at masked positions. To perform the latter,
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Paris the capital of
11, 2..23.123..4.-0
Copying phose
Completion phase
0.08
BERT
1
-2
-2
0.00
Parts the capital of MABK
[1, MASK, 3. 3,1,1,1, KASK, 22 MAMM
JCK
206
33%
#Train
-2
Algorithmic Shift
-2
Figure 1: (A) Matrix completion using BERT. Similar completing missing words in English sentence in MLM, we complete missing entries in masked low-rank matrix. (B) Sudden drop in loss. During training, the model undergoes an algorithmic shift marked by a sharp decrease in mean-squared-ong (MSE) loss. Here, the model shifts from simply copying the input (copying phase) to computing missing entries accurately (completion phase).
distinctive changes occur in the model's attention heads during the period of sudden drop, wherein the model learns to identify relevant positional information to combine various elements in the input matrix and compute missing entries for matrix completion. We perform range of interventions on the input, model (before and after the transition), and training process to further understand this phenomenon, leading to the following observations.
Pre-transition: Copying the Input Matrix Before the transition, the model is simply copying the input matrix both at observed entries as well as missing entries, predicted value for missing entries being nearly O. The attention maps at this stage do not correspond to particularly interpretable structure, and contribute little to the model output.
Tost-Transition: Computing Missing Entries After the transition, the model accurately completes the missing entries, while still copying observed entries. The attention maps at this stage clearly demonstrate that the model 'attends" relevant tokens in the input, and the attention layers are crucial for accurate matrix completion. Interestingly, the post- transition model can outperform the classical nuelear norm minimization algorithm for matrix completion, suggesting that does not simply recover this algorithm.
Model Components and Sudden Drop We analyze the training dynamics of individual components, keeping other components fixed to their final values. We find that different components converge to their optimal values at quite different points during this training.
Preliminaries
2.1 Problem Setup
MLM and LRMC In masked language modeling (MLM), a fraction of tokens in the input sequence are masked out and the model is required to predict the correct token for those masked entries. In this setup, the model has access to both the tokens before and after the current token for computing the prediction. Low-rank matrix completion has a similar structure: given matrix (assumed low-rank) with fraction its elements available, the goal is to predict missing entries. For a matrix X R denote its observed entries by the set Ω C [m] [ni], and the set of missing entries f2 [m] [n] (2 Formally, the problem is
min rank(U) s.t. Uy = Xix v(i.j) Ω
Importantly, both problems (MLM and LRMC) have the same goal -predict the missing entries in the input, i.e., either the language tokens (MLM) matrix elements (LRMC).
Matrix Completion using BERT BERT 12 is an encoder-only Transformer architecture used widely for MLM. For input sequence of tokens t2,.2..... tt., the output is sequence of

---

D ---iinnsionaal 'hidden states' [ei....,ez] E R that is used for prediction. We train BERT model T to predict missing entries in low-rank masked matrix X. For model output X = TFx(X) E Rux", , the training objective L == L(0) is the mean-squared-error (MSE) loss over all entries,
L(
1 (X n² 1,j=1
Σ
X,y)²
In our experiments, data for matrix completion is generated as
X UVUV; U, V.V E Rax" Uei Vij Unif[-1,1] vi,j [x]] x [r]
so that X has rank at most To mask entries at random, sample binary matrices M [0,1)nxP such that Mi = 0 with probability Pmask,, and otherwise; that is, Ω {(i,j) Maj = 1]
Nuclear norm minimization Nuclear norm minimization [6] is a widely used convex optimization approach to LRMC; for completeness, not a convex function of the matrix, one modifies the low rank completion problem by defining the nuclear norm |U]|, i.e., sum of singular values of a matrix U. The overall optimization problem is as follows.
compare our trained models to this approach. Since rank is
min |0|| e
s.t. Uú = Xi v(i.j) E Ω.
2.2 Experiments
Training We use a 4-layer, 8-head BERT model 40 for 7 x 7 (rank-2) matrices, with 'absolute' positional embeddings, no token-type type embeddings, and no dropout. We fix Pmask 0.3 for training. and 256 matrices are sampled as training data at each step (in an 'online' training setup). We use Adam optimizer with constant step size 1e-4 for 50000 steps, without weight decay or warmup. In addition to L, we track MSE over observed and masked entries,
Lobs
1
(Xu, - Xig)3, and Limask (L,))
Σ
Σ (Xuj - X.z)²
Ω
[22
(K)))
Please see Appendix for details on tokenizing matrices and other experimental details. Code is com/
available at this ttt.ggtp../ggggbooo.oo////gggggg -mattttommmmmmmmmmmmmm
Compute Resources For 7 x 7 matrices (training and testing), we used a single (V100 / A100 / I.40S) GPU. A single (A40/ / A100 L40SS) GPU was used for matrices of order 10, 12, 15.
Sudden Drop in Loss
In our training setup, the model converges to fi- nal MSE of approximately 4e-3 that is, it can solve matrix completion well (as in Fig. 3] this MSE is lower than nuclear norm minimization). Fig. 2 demonstrates the loss dynamics over the
Traning BERT Matrix Completion
0.25
0.20
Mean Squared Bner (M52)
course of training the model on this task.
0.5
Interestingly, we observe sudden drop in train- ing loss at approximately step 15000. This sud- den drop in loss is reminiscent phase transi- tions in physical systems, that are characterized by sudden observable changes in the system on continuous variation of some parameter (here equivalent to the number of training steps). Mo- tivated by this similarity, we analyse the 'pre- shift' model step 4000, and 'post-shift' model
0.20
0:00
Step (22000)
Figure 2: Sharp reduction in training loss.
at the end of training, i.c., step 50000 understand model properties and sudden drop in loss.

---

3.1 Before the Algorithmic Shift Copying Phase
Since the value of Lobs remains quite low in the first phase of model training (Fig. 2 algorithm does the model use for predicting matrix entries in this phase?
we ask: what
We find that the model learns to copy the imput verbatim in the first phase (with output 0 for missing entries), verified through token interventions (Sec. 3.1.1 and by investigating the contribution of attention heads (Sec. 3.1.2) towards the output.
3.1.1 Verifying Copying via Token Intervention
To rigorously verify that the pre-shift model indeed copies the input, we replace the masked elements in the 7 x 7, rank-2 input by the token corresponding some m R. For such input, would like to see whether the model implements copying and outputs m at the masked positions. In this setup for model output X. MSE at observed positions is Lobs, and for masked positions the MSE is defined as
mask
(X1, - (x)2 (A.))
Σ
|S29
Lobs and mask for this experiment averaged over 512 samples are compiled in Table (Appendix Ab. The small loss values confirm that model output matches the ground truth at observed positions, while at masked positions it outputs a value nearly equal to m. When the mask token is MASK (i.e., no replacement), we set m 0, indicating that the model outputs 0 at the masked locations.
To generalize this observation to OOD matrices, we sample uniform random 7 x matrices for input; i.e., all entries in the matrix are i.i.d. uniformly in [-1, 1J. Importantly, these matrices do not necessarily have a low-rank structure. With these matrices as input to the same pre- pre-s model as before, we find that model still copies the input (Table [.). This confirms that the model is indeed not 'computing' any entries in the sense of low-rank matrix completion, and simply copies all entries, masked observed.
3.1.2 Attention Heads - Mostly Inconsequential
Attention heads at this stage (Fig. 22a do not appear to attend to tokens an interpretable manner. Since the model is copying the input, and does not need to combine different tokens, Attention heads should not affect the model output at this stage. To confirm that this is indeed the case, we do the following tests.
Uniform Ablation Uniform ablation entails replacing the softmax probabilities an n x n attention head by 1/n² for all elements i.e. 'force' the model to equally attend to all tokens (Sec. 4.6, (22). On such an intervention in our case, there negligible change in MSE at both observed and masked positions. Averaged over 256 samples, Lobe = 3.4e- 4 and Lmakk 0.2236 when using all attention heads; whereas, on ablating all heads, these values are 3.2c and 0. 0.2236 respectively. The negligible change in MSE supports the hypothesis that attention does not contribute to the model output at this stage.
Model Switching In the extreme case, what if we replace the model weights for some component to check for changes to the output? In model switching, we 'transplant' the attention key, query and value weights in the pre-shift model to those from the post-shift model. Averaged over 256 samples, Lass is 5e-3, that similar to the optimal total MSE (L) obtained at the end training, while Lmaak 0.2246, similar to the values obtained without such replacement. This shows that replacing the phi-shift attention weights by the optimal ones does not significantly affect Lobes Lmack highlighting that attention layers have little effect the model output at this stage.
3.2 After the Algorithmic Shift Matrix Completion Phase
In this section, we focus on the model properties in the post-shift phase (specifically, the end of training at 50000 steps). Since L are oear-optimal in this setting, we ask What algorithm the model using for completing missing entries? For example, is implementing the classical nucleur-nome minimization algorithm? For the second question, we show below that the BERT model not implicitly implementing nuclear norm minimization for completing missing entries in the input.

---

Nuclear Norm Minimization We use CVXPY 13 to solve low-rank matrix com-
MSE
Nuclear Norm
10
pletion using nuclear-norm minimization at various levels of Pinak: comparing it to the out- put of a BERT model trained on Pmask = 0.3. find that BERT performs better than nuclear norm minimization with respect to MSE; the same time, the nuclear norm BERT solution is larger (Fig.
1.0
1077
4 40'0
38 10P
10
To verify if the model implicitly optimizes dif- ferent objective for nuclear norm minimization, we also compare to the regularized version of the above problem (A > 0),
rutlear norm min
100
BERT
02
044
05
02
644
05
Pmask
Pmask
Figure 3: BERT Nuclear Norm Minimization. Comparing our model (trained with Pmask 0.3) and nuclear norm minimization the matrix com- pletion task at various levels of Pmask- The differ- ence in MSE and nuclear norm of solutions ob- tained using these two approaches indicates that BERT not implicitly doing nuelear norm mini- mization to complete missing entries.
min
Σ (1J)
(Us x1))² I|
We find that this is not the case, as for various values of X, BERT still outperforms regularized MSE minimization w.r.t. MSE (Appendix [B}. This confirms that the model not implementing nuclear norm minimization as its algorithm for computing missing entries.
We now move an interpretability based analysis of the model behavior, to attempt to extract useful signal about the implemented algorithm, analysing model behavior for observed and missing entries separately in the following sections.
3.2.1 Observed Entries
Uniform Ablation As in Sec. 3.1.2 to quantify the effect of attention heads at this stage, we uniformly ablate all attention heads in the post-shift model. Averaged over 256 samples, this leads to Labs = 9.2e-5 without ablation, and 3. 3.7e-3 with ablation (close to the value of L at the end of training). However, Lmask increases from 0.0128 to 0.2183, approximately the value of Lmmask in the loss plateau before sudden drop. This difference in effect of ablating attention heads confirms that they are much more important for predicting missing entries than for observed entries.
Model Switching We repeat the model switching experiments from Sec. 3.1.2 in the reverse direction Le. 'transplant' attention key, query, value weights from phe-shift model to the post-shift model. Note that this direction of weight switching is stronger, in the sense that the learnt information in attention layers removed. We find that on this modification, Lobs 9.5e-4 averaged over 256 samples; that is, the observed loss is still not too large. This test confirms that the prediction at observed entries not substantially affected by the attention layers.
Position Sensitivity Finally, since the attention mechanism crucially depends on token positions, we intervene on this component of the model by randomly permuting its positional embeddings. Formally, the embedding originally for position in the input now represents position π(i) for some random permutation æ [n²] [n²]. Averaged over 256 samples, Lobe 2.4c-4. whereas Lmaok = 0.5687, indicating that the observed positions are negligibly affected compared to masked positions due this intervention.
These results support our 'sub aub-aalgoithmm' hypothesis; (a) since positional information intuitively required for the copying sub-algorithm, Lobs remains low; and (b) Lmaas increases significantly, demonstrating that removing positional information is detrimental to accurately computing missing
entries.
3.2.2 Missing Entries
To confirm that attention heads causally affect the model output for missing entries, in addition to uniform ablations, we perform causal interventions (activation patching) [42 on the hidden states just after the attention heads. This involves replacing the hidden state after attention head for input
5

---

Head 1
Head 2
Head
Head 4
Head
Head
Head 7
Head
40
48
Layer
Layer2
66
Layer
68
Figure 4: Attention heads in post-shift model attend to specific positions. For example, (Layer Head 1) attends to elements in the same row as the query element, and (Layer 2, Head 2) attends to elements in the same column as the query element. (These attention matrices are an average over multiple independent matrix and mask samples.)
with the hidden state obtained at the same attention head, but a different input A'. Ideally, if that head is causally relevant to the output, then such an intervention should steer the model towards the output for A', instead of A. We find our case that for X and -X, such an intervention on all attention heads clearly steers the model output missing entries towards -X (more details in Appendix F
Structure in Attention Heads Denote attention head H in layer L by the tuple (I, H). We can group the attention heads depending on the specific regions of the input matrix they attend to,
1. [Row Head] same row as the query element 'block diagonal' patterns, e.g. (2, 1);
2. [Column Head] same column query element 'off diagonall patterns e.g. (2, 2); and
3. [Identity Head] query element itself - 'diagonal' patterns in the last layer, e.g. (4, 3),
There are also some other attention heads that do not neatly fit into either of these 3 categories-for example, all heads in layer except (1.3), (1.4); (3.3); (4,2), (4.5-7). In this context, we note that uniformly ablating heads (3,3), (4.2), (4,5-7) gives Lobs = 9.3.6e-.. Lmast = 0.01575 compared Lobe 9.44c Lmaak 0.01428 without ablation, i.c. these uninterpretable heads do not significantly affect the output.
Attention Heads with 'Structured Masking" Since the maps in Fig. are averages over multiple random masks and input matrices, it is difficult to derive more fine-grained insights into the model
computation. To address this, we generate inputs with specific mask structure, see for example Fig. 5 This implies that while averaging the attention probabilities over different input matrices, the mask
iee. Ω remains the same. This step helps us highlight how an attention head attends to input elements
based on the element being masked or observed. From the results in Fig. 5 we find clear evidence that different attention heads focus on specific parts of the input. For instance,
[Masked-Row Head] (2, 1), (3,4) and (4,8) are mainly active only at the masked rows, and therein attends to the only observed position in those rows.

---

Head
Head Head
Layerz
45
Layer
450
42
Layer
Figure Attention heads with specific mask structure in inputs. We can derive fine-graine -reigrained insights about the functions of individual heads in this setup by using specific mask structure for all input matrices. (Mask appended below each plot, blue denotes missing entries). For example, multiple attention heads like (Layer 2, Head 2) have negligible attention weight at missing positions in the input matrix, implying that these heads attend only to observed entries in the column of the query element. Further, (Layer 2 Head 1) and similar heads have larger attention weights for the rows with missing entries, and those rows they attend the sole observed element.
2. [Observed-Copp Head] (4.3) and (4.4) correspond roughly to an identity map, slightly deviating in the masked rows. In these cases, again the maximal attention score corresponds to the only observed position these rows.
3. [Mask-Ignore Heads] Further, there are multiple 'parallel off-diagonal heads that com- pletely ignore the masked rows for their computation. These heads include (2,2-4), (2,6); (3,2), (3,3), (3,5), Additionally, there also attention heads like (3, (3.6) that attend to only the observed element of each masked row. Collectively these heads act as 'mask-ignon heads, attending only observed entries, and using this information to compute missing entries.
4. ConggsstCContguuuu--Coooooo
Heads] There also exist attention heads that respond systematically changes in the mask. For example, consider attention heads (2, 5), (2, 7), (2, 8) in Fig. 23 For each row, these heads attend to the element in the 6th and 2nd column respectively for part (a) and (b). On closer look, the connecting link between these two mask patterns is that, the longest contiguous unmasked column is exactly the column that these heads attend to. We hypothesize that this information is somehow used by the model in its inner computation for masked entries.
5. [Input-Processing Heads] Finally, Heads (1,1-2). (1, 5-8) do not fall in any of the categories above These heads are mostly static across different mask input variations (for example, comparing Fig 4 and [5), and the patterns suggest that these heads almost exclusively focus on the middle row of the input matrix and some other elements. A possible function of these heads is to process positional and token embeddings (input to the first layer) so that this information can be used appropriately in the subsequent layers.
To quantitatively assess the effect of these attention heads the model output, also perform uniform ablations on each sub-group separately (Appendix L). and find that the groups significantly affect the output, varying degrees depending on the specific group.

---

Probing We probe for properties of the input matrix in the hidden states of the model, to concretely determine how the model computes the output. We use our 12-layer model in this case, to enhance contrast between probing results in different layers.
MSE
Average Cosine Similarity
0:13
0.75
Tran
0.12
Text
0.70
0:11
0.65
0.10
0.60
0.35
0.55
Specifically, for every element in the imput, we fit a linear probe [ on its hidden state after given layer, mapping the hidden state to the n--dimensional masked row that this element belongs to (missing en- tries are replaced by 0). That is, el- ement at position (i,j) maps to the 7 dimensiionnl vector X, The results for this experiment in Fig. demon- strate that the hidden states layer 3 and 4 in the model correlate quite strongly with the probe target, compared other layers. This result suggests that the model tracks input information in its intermediate layers and possibly uses it for computing missing entries.
0.08
0.50
9.37
0.45
Layer
Figure 6: Hidden states encode input information. We probe for the row of each element in the masked matrix input, replacing missing entries by O. We find that layers and 4 have a much lower MSE and a much larger cosine similarity compared to other layers in the model. Hence, these layers somehow 'store information' about the masked input matrix.
We also probe for the true matrix element at missing entries, and find that the hidden states at these positions get gradually more correlated with depth (through linear probing). Further, we also attempt to extract information about singular vectors of the ground truth matrix from the hidden states through linear probing, though are unable to conclusively do so. We discuss these results in Appendix
3.3
Role of Embeddings
Token Embeddings The l2 norm of token embeddings corresponding values from 1.5 1.5 is symmetric w.r.t. 0 as seen in Fig. 7a Further, the PCA of token embeddings in Fig. 7b shows that the embeddings have separable structure based on the sign of the real-valued input (y-axis), and continuous variation w.r.t. magnitude of input (x-axis). Importantly, unlike other metrics, token embeddings do seem to abruptly change only at step 15000; rather, the final structure appears before the sudden drop in loss. Similar to (27), we compute the top-2 principal components of the token embeddings at the final step (50000), and project the token embeddings at intermediate training
0.4
0.95
0.2
nomm
[-100 -o.]
0.85
a.)
0.0
(.5.
0.80
-0.24
D. 18)
0012
-1.5
-1.0
-0.5
0.5
1.0
1.5
-0.4
-0.50
-0.25
-25 33
05
1.0
25
(b)
(C)
Figure 7: Embeddings demonstrate relevant structure. In the post-shift model, positional and token embeddings exhibit properties demonstrating that the model has learnt relevant information about the matrix completion problem. (a) E2 norm of token embeddings is symmetric around 0. This aligns with the intuition that the norm of token embeddings should depend only on the magnitude of the input, and not on its sign. (b) Top-2 principal components of token embeddings correspond to the magnitude and sign of the real valued input. In our case, the 'y-axis' denotes sign of input, and the "x-axis" denotes the magnitude of the input value. (c) Positional embeddings of elements in the same column cluster together the SNE projection, showing that the model uses positional information relevant to the matrix completion problem.
8

---

steps on these components. The results (Fig. D Appendix C show that the embeddings align very closely to the final arrangement before the actual drop. This as expected, since the model needs to learn what the tokens actually represent on the real line, before it can use those values for completing missing entries. This also explains to some extent why the model implements copying before the sudden drop, since accurately learning token embedddinguunemeeddingg is sufficient for that task.
Positional Embeddings In the t-SNE projection of positional embeddings, positions in the same column tend cluster together seen in Fig. [7c This is non-trivial because we have used any marker tokens to mark the end of a row or column. Further, note that in contrast to token embeddings, positional embeddings do have a continual evolution in structure - Fig. IO (Appendix (C shows that the clustering appears only after the sudden drop (step 20000 and after). This along with the evolution of attention heads (Sec. 3.1 3.2 aligns with how the pre-shift model copies observed entries with little effect from ablating attention heads or positional embeddings (Sec. 3.2.
4 Sudden Drop in Loss Role of Model Components
Is it possible to analyse training dynamics of individual model components to derive insights about the full model training? This is motivated by the findings in the previous section on embeddings, and in Section [.]] the pre-shift model does not use Attention layers for its computation in that stage, and relies on other components to copy input entries. Hence, in our case, the sudden drop corresponds in large part to learning the right Attention patterns (see Appendix [..1) To analyse training dynamics of different model components, we choose (a set of) components Attention layers, MLP layers, Positional Embeddings and Token Embeddings, randomly initialize them and freeze the weights of other components to their values at the final step of training (Fig. 8).
We find that (a) MLPs and Token Embeddings converge without any observable plateau or sudden drop in loss; (b) for other components, the dynamies resemble those for the full model training (i.e. plateau and then sudden drop), and (c) Positional embeddings show the longest plateau in loss.
Training Dynamics of Individual Components
All components
Positional Embed.
0.10
Token Embed.
Mean Squared Errrr (MSE)
Embeddings
Attention layers
0.08
MUP layers
MLP+Emboddingg
0.06
0.04
0.02
0.00
10
30
50
Step (x1000)
Figure 8: Individual model components have distinct training dynamies. Training individual model components, initializing others to their final value ('All components' indicates normal training). There no loss plateau for token embeddings and MLP layers, in contrast to positional embeddings. where the sudden drop occurs just before step 40000. In all other cases the sudden drop occurs before the sudden drop in usual training.
Additional Results To further understand the effect of data and model properties on the sudden drop in loss, we train
a 2 layer, h-head GPT model on the matrix completion task (Appendix G
models of different depth (number of layers) and width (hidden state dimension) (App. H

---

our model on (mixture of) matrices of different sizes keeping the rank fixed at 1 (App-)) D
a 12-laeer, 12-head model on 10 x 10 matrices of multiple ranks (separately) (App. D
our model on input matrix entries (U,,, Viy) i.i.d. ~ N 1) instead of Unif[-1, (App.
J.I
we also analyze test-time OOD performance of our model in App. D.2
5
Related Work
28 38] analyse 'grokking'. the sudden emergence of generalization during model training. In the context of training dynamics of MLM, [8) analyses 'breakthroughs (sudden drop in loss and associated improvement in generalization capabilities of the model), specifically for BERT. They show that the breakthrough marks the transition of the model however is focused on language tasks, distinct from our setting which is mathematical (and hence more controllable) in nature. We also note that their work is not in the online training setting; our setup is online in the sense of sampling new data at every step of training.
generalizing one. Their work
Mathematical problem solving capabilities of Transformers have been topic of interest lately 22 4.. In fact, [24] show that learning addition from samples is equivalent to low-rank matrix completion. Further, ? show that it is possible to train transformer based model to solve various linear algebraic tasks e.g. eigendecomposition, matrix inversion, etc.; however, to the best of our knowledge, interpretability studies for such tasks have not been conducted before. For interpretability in simpler math tasks, [88 mechanistically analyse GPT-2 small on predicting whether a number is "greater-than" given number, by formulating the problem as natural language task. 35 36 [0]
analyse BERT from an interpretability perspective. More recently, there has been a line of research works analysing decoder based models to ennenner the mechanisms employed by these [4]
models, termed as 'mechanistic interpretability' that our setting distinct from the recent work on solving mathematical tasks like linear regression
31
32
39
I
33 25
26 23 34 20). We note
through 'in 'ontccttext learning in transformers [ 9 15
17
2
30 [3]. Whether our model learns to
implicitly 'implemnn' an optimization procedure as shown in some of these is an open question. We discuss related work in more detail in Appendix E
6 Conclusion
We trained a BERT model on matrix completion, and analyzed it before and after the sudden drop in training loss (algorithmic shift) to interpret the algorithm being learnt by the model, and gain insight on why such a sharp drop in loss occurs.
It is evident in our analysis that both before and after the shift, the model does not really compute anything at observed positions, and simply copies these entries. For missing entries, we have shown that the model learns useful abstractions rapidly through the algorithmic shift. Mathematically formulating what algorithm the model employs to implement matrix completion for missing entries is a direction for future work.
Since our work is primarily interpreting model training and mechanism, all experiments are with small scale matrices (largest being 15 x 15), and the current method would likely need modifications scale to larger matrices. Finally, we only intended study Transformers on matrix completion as a toy task from an interpretability viewpoint, and do not advocate replacing existing efficient solvers for matrix completion with our approach.
We study Transformer based models on their ability to solve a mathematical task (matrix completion) and the associated training dynamics. The work focuses on Transformer interpretability, aiding in improving our understanding of these models and their training and thus we do not foresee any negative societal impact of our work.
Societal Impact
Acknowledgements We thank Yu Bai, Andrew Lee, Naomi Saphra and anonymous reviewers for their helpful comments. WH acknowledges support from the Google Research Scholar Program. ESL's time at University of Michigan was supported by NSF under award CNS-221 1509 and at CBS, Harvard by the CBS-NTT Physies of Intelligence program.
10

---

References
[1] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[2] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference Learning Representations, 2023.
[3] Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes, 2017.
[4] Yu Bai, Fan Chen, Huan Wang. Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[5] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave. K. Cho, and A. Oh. editors, Advances in Neural Information Processing Systems, volume 35, pages 21750-21764 Curran Associates, Inc., 2022.
[6] Emmanuel J. Candès and Benjamin Recht. Exact matrix completion via convex optimization. Found. Comput. Math., 9(6):11-7772 2009.
[7] Francois Charton. Linear algebra with transformers. Transactions on Machine Learning Research, 2022.
[8] Angelica Chen, Ravid Shwartz-Ziv, Kyunghyun Cho, Matthew L Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in MLMs. In The Twelfth International Conference Learning Representations, 2024.
[9] Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent learn non-linear functions in context, 2024.
[10] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? an analysis of BERT's attention. In Tal Linzen, Grzegorz Chrupata, Yonatan Belinkov, and Dieuwke Hupkes, editors, Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, Florence, Italy, August 2019. Association for Computational Linguistics.
[11] Arthur Conmy, Augustine N. Mavor-Parker. Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alonso. Towards automated circuit discovery for mechanistic interpretability. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[12] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171-4186 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[13] Steven Diamond and Stephen Boyd CVXPY: Python-embeddeed modeling language for convex optimization. Journal of Machine Learning Research, 17(83):1-5,, 2016.
11

---

[14] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, 7ac Hatfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https//ttansformer--iciiisssss http:/ppp//b////a/ooooaoaaoo //frameworr/iddee..nen
[15] Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models, 2023.
[16] Deep Ganguli, Danny Hernandez, Liane Lovitt. Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova Dassarma, Dawn Drain, Nelson Elhage, et al. Predictability and surprise in large generative models. In Proceedings of the 2022 ACM Conference on Fairness, Accountability, and Transparency, pages 1747-1764, 2022.
[17] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[18] Michael Hanna, Ollie Liu, and Alexandre Variengien. How does GPT-2 compute greater- than?: Interpreting mathematical abilities pre-trained language model. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[19] David T. Hoffmann, Simon Schrodi, Jelena Bratulić, Nadine Behrmann, Volker Fischer, and Thomas Brox. Eureka-moments in transformers: Multi-step tasks reveal softmax induced optimization problems. In International Conference on Machine Learning, July 2024.
[20] Yifan Hou, Jiaoda Li, Yu Fei, Alessandro Stolfo, Wangchunshu Zhou, Guangtao Zeng, Antoine Bosselut, and Mrinmaya Sachan. Towards a mechanistic interpretation of multi-step reasoning capabilities of language models. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. Singapore, December 2023. Association for Computational Linguistics.
[21] Margot E Kaminski. Regulating the risks of ai. Forthcoming, Boston University Law Review, 103:22-21, 2023.
[22] Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of BERT. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4365-4374, Hong Kong, China, November 2019. Association for Computational Linguisties.
[23] János Kramár, Tom Lieberum, Rohin Shah, and Neel Nanda. Atp An efficient and scalable method for localizing I1m behaviour to components, 2024.
[24] Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers, 2023.
[25] Michael A. Lepori, Thomas Serre, and Ellie Pavlick. Uncovering intermediate variables in transformers using circuit probing, 2023.
[26] Tom Lieberum, Matthew Rahtz, János Kramár, Neel Nanda, Geoffrey Irving, Rohin Shah, and Vladimir Mikulik. Does circuit analysis interpretability scale? evidence from multiple choice capabilities in chinchilla, 2023.
[27] Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 34651-34663. Curran Associates, Inc., 2022.
(28) Ziming Liu, Eric Michaud, and Max Tegmark, Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference Learning Representations, 2023,
12

---

[29] Ekdeep Singh Lubana. Kyogo Kawaguchi, Robert P. Dick, and Hidenori Tanaka. A percolation model of emergence: Analyzing transformers trained on formal language, 2024.
[30] Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference Learning Representations, 2024.
[31] Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Circuit component reuse across tasks in transformer language models, 2024.
[32] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, In The Eleventh International Conference on Learming Representations, 2023.
[33] Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models. In Yonatan Belinkov, Sophie Hao, Jaap Jumelet, Najoung Kim, Arya McCarthy, and Hosein Mohebbi, editors, Proceedings of the 6th Black- boxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, Singapore, December 2023. Association for Computational Linguistics.
[34] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Haffield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. heads/index.html.
[35] Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in BERTology: What we know about how BERT works. Transactions of the Association for Computational Linguistics, 2020.
[36] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi- head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Anna Korhonen, David Traum, and Lluís Mårquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguisties, Florence, Italy, July 2019. Association for Computational Linguistics.
[37] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 35151-35174, PMLR, 23 29 Jul 2023.
[38] Boshi Wang, Xiang Yue, Yu Su, and Huan Sun. Grokked transformers are implicit reasoners: mechanistic journey to the edge of generalization, 2024.
[39] Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023.
[40] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface 's transformers: State-of-the natural language processing, 2020.
[41] Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in reLU networks for XOR cluster data. In The Twelfth International Conference on Learning Representations, 2024.
13

---

[42] Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metrics and methods. In The Twelfth International Conference on Learning Representations. 2024.
[43] Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Initial- ization is critical to whether transformers fit composite functions by inference or memorizing, 2024.
14

---

A
Copying in Pre-Shift Model
Mask "MASK"
Mask "0.44"
Mask 00224
Step
Input Samples
Linask
Lobe
L'mask
Lobs
L/mask
Lobs
Rank 2 matrices
1.4e-3
9.3e-4
7.6e-4
1e-3
6.7e-4
9.6e-4
1000
Random matrices
1.5e-3
8.3e-4
7.8c-4
le-3
6.8c-4
9.6e-4
Rank matrices
3e-4
3.3e-4
4e-4
2.8e-4
3.7e-4
2.7e-4
4000
Random matrices
2.8e-4
3.5e-4
3.7e-4
3e-4
3.6e-4
2.8e-4
Rank 2 matrices
1.6e-5
3.4e-5
1.8e-5
4.9e-5
5.1e-6
6.7e-5
14000
Random matrices
1.le-5
3.7e-5
3.0e-5
8.5e-5
4.1e-6
1.1e-4
Table I: Models at different steps before sudden drop implement copying, predicting the value for mask token at missing entries.
B Nuclear Norm Minimization
We use the regularized version of the nuclear norm minimization problem as detailed in Sec. 3.2 and obtain the following L, Lobs: Lmaas for various values of X. We average our results over 256 samples generated in the same way as the training data for BERT (including rounding off to 2 decimal places) for the sake of comparison.
Lobs
Lmmas
L
0.0005
1.015e-5
0.040728
0.012173
0.001
3.686e
0.040456
0.01211
0.0015
7.959e
0.040505
0.012155
0.002
0.00013769
0.040734
0.012264
0.005
0.00078591
0.043402
0.013516
C Embeddings Progress during Training
Stop
32
Figure 9: Token embedding structure appears before sudden drop. Projection of token embed- dings along principal components of embeddings at step 50000. (Labels same Fig. 77. Embeddings align with the principal components early on in training before the sudden drop loss.
Stop 4000
Stop 10000
Step 14000
Step 18000
Step 20008
1.
-30
25
12
1.0
Figure 10: Positional embedding structure appears after sudden drop. Evolution of t-SNE projection of positional embeddings with training. (Labels follow the same color coding as Fig. ( In this case, the positional embeddings show clustering some time after the sudden drop in loss has occurred step 15000.
15

---

D Experimental Details
Online Training In online training, the data is sampled afresh from the distribution at every step.
Since data is not partitioned into fixed train and test sets, we only analyze the training loss in all cases.
Tokenizing Matrices For tokenizing real values, we discretize the range [-10, 10] in steps of size <= 0.01, and assign token IDs starting from 1; the mask token (MASK) is assigned ID 0. Input to the model is the tokenized masked sequence Xaak = TOK (Vec(XX @ M)), where @ denotes the element-wiss product, Vec denotes vectorizing the n n matrix to a n² -i-dimensional vector, and TOK denotes tokenization. Due to this preprocessing, in all cases X is rounded to 2 decimals for computing L.
E Related Work
In the online training setup, [5] study the parity learning problem using a variety of model archi- tectures, and show that 'hidden progress measures' can be used to track abrupt changes in model performance during training. [29] study abrupt learning in an autoregressive (GPT), language data setup, connecting learning the grammar to percolation on graphs. [99] discuss abrupt learning dynamies in the context of transformers and claim that the softmax function in Attention leads to longer training loss plateaus however, reducing the length of plateau does not explain why the drop in loss sudden and sharp when it occurs. [43] show that initialization of the model affects whether it learns to infer the compositional structure of the task, or simply memorizes the solution.
4 show that in an in-context learning framework, transformer based model can learn 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 layer-wise in-context gradient descent can also be used in our setup is plausible and open direction for future work.
In (7). the author shows empirically that encoder-decoder transformer 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 DonnWr matrices. While many experiments in 7 also show sudden jump in accuracy (Fig 1.2), they do analyze why such 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 setup.
[44] 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- 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 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 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 (Pmask) of observed entries.
F Causal Intervention on Attention heads
In the uniform ablation setup, it is possible that setting the softmax probabilities given value might change the distribution of resultant hidden states, and consequently degrade model performance. more principled technique to analyze the effect of specific component is to replace the hidden state just after that component by hidden states on a different input, and analyze how this affects the final output 42). our case, we intervene on attention heads by replacing the hidden state after an attention head for input matrix X by the hidden state for input (-X). Importantly, this change does not affect properties like rank of the input, and hence the hidden states obtained are from the same distribution those for input X.
16

---

Step I Extract the hidden states for all attention layers from the model on some imput matrix X; call these h4- Concretely, these hidden states are obtained just after the matrix product of the softmax attention probabilties and the value matrix and hence before the output matrix product.
Step 2 Change the input to the model to X, however, also replace the hidden states just after the attention layers with h4 obtained in Step 1. Call the output of the model in this setup as xx)
We observe that, the MSE between fp(-X,X) and X, averaged over 256 samples at masked positions is approximately 0.014 (this is comparable to optimal Lmosk), compared to the MSE between fx(-X,X) and -X being 0.8066. This demonstrates that the attention heads are causally relevant to the model output for missing entries.
G Autoregressive Setup
Training (Cross Entropy) Loss
Accuracy at Observed Entries
10-
5.5
08
0.6
4.5
35
3.0
Step (x100)
Step (x100)
Figure 11: GPT also shows abrupt learning for matrix completion. Training a 2-layer, 2- GPT model on x 7, rank-2 matrices in the autoregressive training setup. Here, the model is trained using cross-entropy loss in a next-token prediction setup over full input sequences of the form X11, X12. Ii,.... X77, [SEP), 11, X12...... X77 where X, X are partially observed and fully observed matrices, flattened and tokenized as the BERT experiments. We find that the sudden drop corresponds to the model learning copy the observed entries in the input matrix. While could not achieve performance comparable to BERT for missing entries, we believe it should be possible with some modifications to the training setup.
Step 200
Step 600
Head
Head
Head
Head 2
Layer 1
Layer 1
42
48
62
25
Layer
45
Layer 2
45
(a)
(b)
Figure 12: Attention heads demonstrate sudden change in matrix completion using GPT. We
find that even in the GPT case, the attention heads change from trivial (Layer 2, attending to the [SEP] token for all positions in the output) to those in Layer attending to the corresponding positions in the input (>identity maps). This corroborates with our finding about the model learning to copy observed entries after the sudden drop, in Fig. 11
17

---

H
Effect of Model Size
Training Models of Increasing Width
d=64
d=128
0.10
Mean Squared Error (MSE)
d=768
d-1024
0.08
0.05
0.02
Step (x1000)
Figure 13: Effect of model width. Training with different model widths (hidden state dimensionality) on 7 rank-2 inputs. The plot demonstrate that d 64 is small to obtain accurate matrix completion, and that the performance sub-optimal for d = 128. We scale the hidden layer width of the 2 layer MIPs 4d, as is done in practice.
Training on Models of Increasing Depth
Layers
Layers
8 Layers
Mean Squared Error (MSE)
Layers
0.08
0.04
0.00
10
30
50
Step (x1000)
Figure 14: Effect of model depth. Training with different model depths (number of layers) on 7 x 7 rank-2 inputs. The plot demonstrate that depth increases from to the sudden drop occurs earlier, but increasing depth beyond this (8, 12) has little effect. The final MSE obtained also follows the intuitive ordering (largest for decreasing with L upto 12; though the variation not significant.
Figure 15: Using attention head per layer. We find that training a 12-layer, 1 hehd BERT model on matrix completion leads similar loss (4e-3) and attention heads the 4 1-layer, 8-head model.
18

---

I
Effect of Matrix Size and Rank on Training
Training Mixture Different Sized Inputs
5x5
7x7
0.10
9
Mean Squared Error (MSE)
0.08
0.06
0.04
0.02
0.00
10
Step (x1000)
Figure 16: Learning order when trained on mixture of matrices. Training on uniform mixture of 5 x 5, 7 and 9 x 9 rank-1 matrices i.e., each step, 256 samples of size n, with chosen randomly from {5, 9. The plots show the test set MSE on separate 256 samples of the 3 different matrix sizes.
Training on Different Matrix Ranks
Rank-1
0.10
Rank-3
Mean Squared Errar (MSE)
0.08
0.06
0.04
0.02
0.00
10
30
50
Step (x1000)
Figure 17: Effect of problem complexity. Training layer, 12 head model on 10 x 10 matrices of rank = 1, 2, 3. There is a clear progression in terms of the training step where sudden drop occurs, and the final loss values scale roughly L e- 10", e 10 This also predicts that L ~ 0.02 for r = 4, i.c. the model does not solve matrix completion to low MSE, which is what we also observe in practice.
19

---

Effect of Input Distribution
J.1
Training
Training Inputs with Gaussian entries
0.10
Mean Squared Error (MSE)
0.08
0.06
0.04
0.2
0.00
Step (x1000)
Figure 18: Sudden drop is not limited to uniform distribution. Training on i.i.d. (0, 1) entries. We find that the sudden drop also occurs in this case, and the final loss value 5.6 x 10-3 similar to the value obtained for i.i.d. Uniform[- -1, 1] entries.
J.2 Inference
Nucies zor Norm
MSE
Nuclear Norm
Nudear Norm
30m
1:10
100
1x
31:20
109
400
4907
K00
(a) Rank
(b) No. of Rows
(c) No. of Columns
Figure 19: Model performs similar to nuclear norm minimization on OOD samples. OOD performance at inference for various values of rank, number of rows and columns of the input matrix. Except (c), the OOD performance of the model is close to the nuclear norm minimization solution for the same inputs. For (c), since we observed that positional embeddings depend on the column of the element, changing the number of columns adversely affects performance.
Matrix Distribution We also change the input distribution of the matrix entries to Gaussian and Laplace, and measure average MSE over 1024 samples of size and rank 2, evaluate the OOD performance of the trained model. We find that
for entries i..dd \/ 0.25),
10 and
for entries i..d. Laplace(O, 0.25), (parameterized by mean and scale), L 10
That is, the OOD performance on these distributions is similar to the MSE obtained for the in-training distribution (Uniform[-1) 1)).
20

---

K
Probing: Additional Results
MSE
Test MSE
Tram MSE
Maked slmeets
0.15
0.04
0.05
0.10
0.05
0.01
0.00
0.00
Layer
Layer
Figure 20: Hidden states encode the true value at missing entries. Probing for the corresponding element the fully observed matrix X. (Left) comparing the train and test MSE of the linear probe, to confirm that the probe does not overfit. (Right) Test MSE evaluated on the masked elements, that shows an interesting variation: the MSE goes down at a nonlinear rate, hinting that implicit laser-wise optimization could be taking place.
MSE
Avg. Absolute Cos-Sim Sim
Trais
0.145
Test
0.44
0.140
0.40
0.135
0.38
0.130
0.34
3
4
6
10
2
9 10 101111 12
Layar
Layer
Figure 21: Model does not encode singular vectors in hidden states. Probing for the first left singular vector at all positions; that is, for fully observed input matrix X E Raxn with SVD X = U.VT, we probe for u = [U11, U..)) E TR. (Left) Train and Test MSE across different layers of the model; (Right) Average Absolute Cosine similarity of the predicted with the actual u. Both evaluations are inconclusive, since the MSE is too large, and the cosine similarity is not much larger than the average absolute cosine similarity (+s 0.3) between 2 i.i.d. vectors (0, 17x7) R
L Ablating Groups of Structured Attention Heads
Group of Attention Heads Ablated
Ratio of & (w/ to w/o)
ahlaton
abbiloo
Wio dBlaton
WF ablatio
wilannatoo DLOTTY
(48)
0.0073
(4.3),
(2.2 (2.6)
(2.6) (3.) (SJ)
0.057
B.0643
2.1e-1
95-5
0.1885
0.0142
13.26
(2.5.
00012
8.2 5
11.6 6
0.016
2.231
0.0314
0.1038
21

---

M Attention Heads
M.1
Variation along training
Head
Head
Hoad
Head
Head
Hoad
Head
Hoad
48
Layer
Layer
48
48
Layer
(a) Step 4000
Head 1
Head 2
Head
Head 4
Head
Head
Head 7
Head 6
Layer
1:00
0:00
6:C
8:77
0000
020)
2.01
3270
Layer
(b) Step 14000
22

---

Head
Head
Hoad
Head
Head
Head
Head
Head
48
Layer
68
10
18
18.0
Layer
(c) Step 16000
Layer
41
41
48
Layer
48.0
Layer
40
(d) Step 20000
Figure 22: Attention heads across various training steps.
23

---

M.2 Effect of changing mask structure
Head
Head
42
Layer
45.0
(a)
Head
Head
Head
Head
Head 5
Head
Head
Head
45.0
0
Layer
Layer
Layer
(b)
Figure 23: Attention heads and corresponding masks; blue denotes masked position in the input matrix.
24

---

N
Attention Heads for larger inputs
Head
Head 2
Head
Head
Head 5
Head
Head
Head 8
Head
Head 10
Head 11
Head 12
Layer1
Layer2
Layera
60
65
48
46
46
45
Layer
Layer
40
10
40
40
41
Layer Layer G
40
40
40
Layer 8
Layer
46
45
45
byyrr 10
Layer 11
Layer 13
Figure 24: Attention heads in 12 layers, 12 heads model on 7 rank-2 input
25

---

Head
Head
Head
Head
Head
Head
Head
Head
Head
Head
Head 11
Head 12
Layer
143
143
143
243
143
143
143
143
143
143
343
143
143
Layer 2
Layer
147
145
145
145
145
145
145
345
145
145
Layor
143
Layers
143
163
143
153
143
143
149
14
Layer 8
163
161
149
Layer
143
143
Layer 8
363
143
143
143
Layer
Layer 10
143
163
143
145
143
143
143
143
143
143
Layer 11
Layer 12
Figure 25: Attention heads in 12 layers, 12 heads model on 12 x 12 rank input
26

---

Head
Head
Head
Head
Head
Head
Head 7
Head
Head
Head
Head 11
Head 12
Layer
224
224
224
224
224
224
224
224
224
224
224
224
224
Layer 2
Layer
220
224
224
224
214
224
224
224 2
214
204
224
214
Layor
Layer
220
224
224
224
224
224
234
224
224
224
214
Layer 8
224
224
224
224
234
Layer
224
224
Layer 8
224
224
224
224
224
224
Layer
220
Layer 10
224
224
224
214
224
224
224
214
234
214
214
Layer 11
Layer 12
Figure 26: Attention heads in 12 layers, 12 2eads model on 15 x 15 rank input
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have accurately listed our main claims in line with the actual contribution this work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed limitations of our work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
28

---

Justification:
There are no theoretical results in this work.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We describe our training experiment details Section 2.2 and other experi- ment details provided wherever relevant. Link to the code is available in Sec. 2.2
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
29

---

Answer: [Yes]
Justification: Link to the code is available in Sec. 2.2 and README file is included to help experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
httpol/liigggddde//ooooee httpg/udddd//oeeSuubni ssionpolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide all such experimental details in Sec. 2.2
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due to computational issues, we do not report error bars for our training experiments this time.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
30

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We discuss compute resource usage in Sec. 2.2
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our work conforms to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss societal impact of our work in the conclusion (Sec. 6
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
31

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system Iearns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Risk of misuse in case does not apply, since work involves interpretabil- ity analyses of a BERT model on a mathematical problem.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have mainly used the BERT model implementation from Huggingface (cited in Sec. 2.2 and CVXPY (cited in Sec. 3.2). We have not used other existing code (except libraries or frameworks available publicly). data or models.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
32

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Code released for this paper (Sec. 2.2 is accompanied by a README file that details how to reproduce empirical observations.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
33


Reviews:
[{"review_id": "spFbGztCVs", "reviewer": "Reviewer_LwY5", "summary": "The paper examine the ability of the transformer model to solve a low-rank matrix optimization problem. The numerical matrices with missing values are tokenized and flattened as sequences with masked values, and a BERT transformer is trained to predict the original sequence with both masked and unmasked values.  The paper make the following observations through extensive tests: \n+ The learning of transformer in matrix completion task exhibits two stages: copying masked values, and sudden transition to learning the underlaying patterns and drastically decreasing loss. \n+ During the second, different attention heads exhibits different behaviors correspond to different parts of the input: masks, observed values, position embeddings.\n+ The learned token embeddings exhibits clusters between elements from the same row in the original matrix, even when not row/column marker explicitly providing.\nThese findings aligns with prior studies focusing on masked-language-modeling(MLM) tasks and provides valuable insights into the learning mechanism of transformers.", "questions": "+ What is the purpose of training the model to predict all values instead of just the masked ones?\n+ What do the two panels in figure represent? The right panel caption appears incorrect and the loss surprisingly goes down as p_mask increases.\n+ Is there any consideration for tokenizing the matrix entries, besides mimicking language sequences? Since transformers are increasingly used to processed multimodal data whose input are already embeddings, it would be interesting to see if the same patterns hold for continuous inputs.", "limitations": "The authors have addressed societal impact and limitations in the conclusion section.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "+ The paper proposed on interesting combination by linking low rank matrix completion problem with MLM, potentially enabling the used of many NLP technique to solving the application involving the first problem.\n+ The paper performed extensive experiments and examine the learning behaviors of transformers thoroughly across learning stages and important components of transformer.\n+ The paper presented sufficient technical details for the experiments implemented such as the loss function and training parameters.", "weaknesses": "+ Though using matrix completion as an \"abstraction\" of MLM is interesting and novel, the author gave little concrete explanation or reference on why insight learned from such setting is transferable to language domain, and what is the benefit over directly analyzing language sequences in terms of gaining interpretation and insights.\n+ Despite adapting a new task, the main patterns the paper found in through experiments are similar to prior works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments, and are happy to know they found our work interesting!\n\n- *Though using matrix completion as an \"abstraction\" of MLM is interesting and novel … interpretation and insights.*\n\nWe 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.\n\nWe 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.\n\n[1] https://arxiv.org/abs/2309.07311\n\n\n- *Despite adapting a new task, the main patterns the paper found in through experiments are similar to prior works.*\n\nWe 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. \n\nFurthermore, 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. \n\n[1] https://arxiv.org/abs/2309.07311 \n\n\n- *What is the purpose of training the model to predict all values instead of just the masked ones?*\n\nWe 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.\n\n- *What do the two panels in figure represent? The right panel caption appears incorrect and the loss surprisingly goes down as p_mask increases.*\n\nWe 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.\n\n\n\n- *Is there any consideration for tokenizing the matrix entries, besides mimicking language sequences? Since transformers are … if the same patterns hold for continuous inputs.*\n\n\nIndeed, 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uGE128D4e1", "reviewer": "Reviewer_Kxvs", "summary": "The paper applies a BERT-style transformer encoder to do low-rank matrix completion. To frame it as a masked language modeling problem, they restrict the domain of the problem to smaller matrices and discretize the domain.\n\nThey find the model can solve the problem. The training dynamics are interesting. Initially the model just copies the input, but after a certain point, it solves the problem.\n\nThey analyze many components of before and after this shift: attention heads, activations, and embeddings. They find the behavior of the attention heads changes via visualization and activation patching.", "questions": "The authors do a good job of anticipating many questions.\n\nDoes this behavior generally to other mathematical tasks?\n\nCan we get any insight about the actually algorithm or implicit optimization procedure?\n\nWhat is the generalization to larger matrices?\n\nIs there a way to more rigorously verify the hypotheses from the **Attention Heads with Structured Mask** section?", "limitations": "The authors acknowledge the limitations in the discussion. This only works on small matrices and is not a generally useful matrix solver.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "It's an easy to read paper with good figures and sound methodology.\n\nThe authors anticipate many of the questions and acknowledge limitations.", "weaknesses": "The main weakness is how significant the contribution is. They have some interesting results which lead to some hypotheses. In the related work section, they frame the paper as mathematical capabilities of transformers, which is indeed an interesting area. But I don't feel like we've gain much insight into how the transformer does math. The paper feels like preliminary results that could lead to a more interesting paper down the road.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments, and are glad to know they found our results interesting!\n\n- *The main weakness is how significant the contribution is. They have some interesting results … how the transformer does math.* \n\nWe 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. \n\nWe 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. \n\nFurthermore, 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. \n\n\n[1] https://arxiv.org/abs/2309.07311\n\n- *Does this behavior generally to other mathematical tasks?*\n\nIt 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).\n\nWe find that our observations in matrix completion extend to this setup (Fig. O, P in attached PDF). Specifically, \nthere is a sudden drop in training loss, and corresponding increase in accuracy of the model in predicting c_i (Fig. O), and\nattention 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.\n\nThis 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.\n\nWe 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. \n\n\n[2] https://arxiv.org/abs/2402.03902\n[3] https://openreview.net/forum?id=L2a_bcarHcF\n\n- *Can we get any insight about the actually algorithm or implicit optimization procedure?*\n\nAs 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.\n\nOverall, 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. \n\n\n- *What is the generalization to larger matrices?*\n\nPlease see section “Larger matrices” in global response.\n\n- *Is there a way to more rigorously verify the hypotheses from the Attention Heads with Structured Mask section?*\n\nPlease 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).\n\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CEBQR9eBQE", "reviewer": "Reviewer_h7Ji", "summary": "The paper explores the behavior of BERT on matrix completion tasks. The authors show that the model's loss shows a phase transition, where the model switches from copying tokens for filling masked tokens to predicting the masked entries accurately. The authors also conduct probing studies to understand the structure of attention layers, and the hidden representations during the two phases. Overall, the paper takes an important step towards a mechanistic understanding of transformer models, and will be an interesting read for the wider community.", "questions": "Please check my questions above.\n\n[Post rebuttal]: I have increased my score. The authors' responses have resolved my primary concerns on architecture and matrix ranks. I think this is a timely analysis on training time phase transitions in transformer models in a synthetic setting. The preliminary observations on GPT-2 are very interesting. I hope the authors can include them in the next version.", "limitations": "The authors discuss limitations of their work in section 8, and also mention few interesting directions that can be explored by the community as future work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The strength of the paper lies in its simplistic exposition of motivation, experimental setup to showcase changing behavior of BERT during training, and the probing studies to explain internal behavior during phase transition. The authors build their motivation from Chen et al. [2024]'s observations on BERT's phase transitions during pretraining on language data, and show the characteristic of such transitions on synthetic matrix completion data. The authors compare the solution to nuclear norm optimization and show that the model outperforms  the candidate algorithm. Finally, with careful probing studies, the authors report the emergence of structure in model's internal representations, which help the model in the mask-token predictions. Overall, the paper takes an important step towards a mechanistic understanding of transformer models, and will be an interesting read for the wider community.", "weaknesses": "I have a few questions about the experimental setup and possible interesting followups that the authors may pursue.\n\n- How does convergence of BERT models change with increasing rank of the underlying model? Furthermore, how do results change if the authors mix in matrices of different sizes (say 3x3, 5x5, 7x7, 9x9) but fix rank of the underlying solution? Will the model perform equally well for each of them? \n\n- An interesting ablation would be to check the effect of the model's size on the convergence speed of the model. What happens if the number of attention heads is fixed to $1$ but the width and the depth are changed? More such ablations will help strengthen the understanding of the model.\n\n- In the experiments section \"Attention Heads with Structured Mask\", how does the removal of each group affect the behavior of the trained model? Are some groups more important than the rest?\n\n- In the \"Probing\" experiments, \"the model tracks information in its intermediate layers and uses it for computation\". Is this check only for the row or column number? If so, how does it \"only\" help the model computation? Is a more fine-grained probing possible, e.g. does the model compute the elements of the low rank decomposition in its embeddings?\n\n- Any discussion on how the results might change for auto-regressive training will be interesting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their extensive suggestions, and are happy to know that they found our work important and interesting! \n\n- *How does convergence of BERT …*\n\nThanks 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.\n\nFor 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.\n\nFinally, 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.\n\n(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.)\n\n- *An interesting ablation would be to check …*\n\nThanks 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.\n\n**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.\n\n**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.\n \n**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!\n\n\n- *In the experiments section ...*\n\nPlease see section “Structure mask” section in our global response for a response to this concern. \n\n- *In the \"Probing\" experiments ...*\nThis 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. \n\nFor probing results on singular vectors (i.e. low rank decomposition), please see section “Probing” in global response.\n\n- *Any discussion on how …*\n\nGreat 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "emCOOmhgzP", "reviewer": "Reviewer_1DQa", "summary": "The authors study how encoders perform at the task of matrix completion.\n\nThey generate synthetic data and train encoders of different sizes to predict masked out tokens.\n\nThey observe an interesting behaviour where initially the loss appears to be at a plateau but then drops.\n\nTheir hypothesis -which they investigate- is that in the first stage of training the model learns to copy, while in the second stage it learns to predict the masked out tokens.", "questions": "Did you consider other distributions?", "limitations": "See weaknesses.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I find the problem formulation interesting, as well as how the observed phenomenon is further investigated by the authors.", "weaknesses": "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)\n\nThe experimental setting is limited. Synthetic data generation is only carried out using only the uniform distribution. What about other distributions? What about OOD performance?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback, and are glad that they found our problem formulation and analysis interesting! \n\n- *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)*\n\nWe 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.\n\n[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).\n\nIn [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.\n\n[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.\n\n\n- *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?*\n\nThanks for this suggestion! We have now conducted new experiments to address your comment.\n\n**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).\n\n**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.\n\n**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.\n\n\n\n\n- *Did you consider other distributions?*\n\nPlease see our response to the previous question.\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kXv96sDX5F", "reviewer": "Reviewer_eRgv", "summary": "The paper explores the behavior of Transformer models in the context of low-rank matrix completion, treated as a masked language modeling (MLM) task. The authors train a BERT model to solve matrix completion and analyze its performance, particularly noting an algorithmic shift characterized by a sudden drop in loss, transitioning from copying input to accurately predicting masked entries. The paper aims to provide insights into the interpretability of Transformer models.", "questions": "1. Does the model operate through two distinct pathways for non-MASK and MASK positions, as suggested by the different behaviors observed during training?\n\n2. What would be the model's behavior if the attention matrices' parameters were frozen at the start of training to match those of a fully trained model? Would the model still exhibit the initial copying behavior, or would it directly perform matrix completion?\n\n3. How well does the model generalize to different types of low-rank matrices, especially matrices with different sparsity and rank levels? Can it accomplish the task of finding the lowest rank matrix?\n\n4. Can the approach be scaled to larger matrices, and what modifications would be necessary to achieve this?", "limitations": "See weaknesses.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents an interesting approach by framing matrix completion as an MLM task.\n\n2. The experiments are detailed, with clear descriptions of data preprocessing, training, and performance metrics.\n\n3. The identification of an algorithmic shift during training is a significant observation, providing insights into the model's learning dynamics.", "weaknesses": "1. The paper lacks a deep theoretical or mechanical analysis of why the model undergoes an algorithmic shift and how it learns the task.\n\n2. There is insufficient explanation of the internal mechanisms by which the model learns matrix decomposition, particularly in the second stage.\n\n3. The experiments are conducted on relatively small-scale matrices (up to 15x15), raising concerns about the scalability and generalizability of the findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments, and are excited to know that they found our approach interesting!\n\n- *The paper lacks a deep theoretical ...*\n\nWe 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.\n\nTo 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. \n\n- *The experiments are conducted ...*\n\nPlease see our global response on “Larger matrices’, where we discuss possible ways to scale up our approach to larger input matrices.\n\n- *Does the model operate …*\n\nOur 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. \n\n- *What would be the model's behavior …*\n\nPlease 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.  \n\nOur 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.\n\n**Train only Attention layers (labeled ‘att’)** – smaller loss plateau than “all” with sharp drop\n\n**Train only Embeddings (both positional and token, labeled ‘embed’ in the plot)** -  smaller loss plateau than “all” with sharp drop\n\n**Train only positional embeddings (labeled ‘pos’)** – this has the longest plateau of all settings\n\n**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.\n\nFrom 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.\n\n\n- *How well does the model generalize to different types …*\n\nWe 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.\n\n\n- *Can the approach be scaled …*\n\nPlease see section “Larger matrices” in our global response for addressing this concern. \n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their detailed feedback and are excited to see that they find our work interesting! To address raised questions, we have performed several new experiments, as described below. We will add these results to the final version of the paper. 

**Larger matrices.** We show that our setup extends to matrices of size 15x15, rank-4 with a 12–layer, 12–head model with similar training dynamics (please see Fig. E in PDF attachment) and final attention head maps (please see Fig. 14 in manuscript). Extending to very large matrices is limited by the context length of transformers and compute resources; though we are optimistic that such results should hold in general for larger matrices. We would also like to re–emphasize that our main focus is not solving matrix completion through transformers, but study transformers in a mathematical setup to (i) model the sudden drop in loss and (ii) interpret the model behavior to understand why this drop occurs.


**Probing.** We probe for the true matrix element at each position in each layer; i.e., map the 768-dimensional hidden states at position (i,j) for input X_mask, to the real value X_{ij} (Fig. H in attached PDF). Interestingly, we find that the MSE decreases at an exponential rate for masked positions. We conjecture that this might be related to the idea that each layer computation of the transformer implicitly corresponds to a step of some optimization algorithm (e.g. gradient descent [1, 2]), however we could not empirically verify that this is the case for our setup. We leave it to future work to confirm if this conjecture holds true in our case.  

We also attempted to probe the singular vectors of the ground truth matrix in the hidden states of the model. Concretely, we map the 768-dimensional hidden states at different layers through a linear model for a given input matrix X, to the 7-dimensional first left singular vector u of X (i.e. if X = U.S.V^T is the SVD of X, then u = first column of U). We found that the MSE is too large (~0.14), and the cosine similarity of the estimated vector and the true vector is near random (<0.3 i.e., less than the average cosine similarity of 7-dim vectors sampled from N(0, I)); please see Fig. (I) in the PDF attachment for these results for all layers. Hence, it is not immediately clear that the model has some information about the singular vectors of the input matrix stored in its hidden states, that is recoverable through linear probes. Similar finding holds for other singular vectors i.e second column of U, and first 2 columns of V.

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

[2] https://proceedings.mlr.press/v202/von-oswald23a.html

**Structured mask.** We find that uniform ablation on each category of attention heads in Section "Attention Heads with Structured Mask" leads to increase in L_mask to different extents. The maximum effect of ablations is in the observed heads category ((2,2–4), (2,6), (3,2), (3,3), (3,5), (3,1), (3,6)), quantified by measuring the ratio of MSE with ablations to MSE without ablations. The first layer attention heads that are claimed to possibly “process positional and token embeddings” in row 5 in the table, also affect the model output to a larger extent. This rigorously verifies the hypothesis that these attention head groups causally affect the model output. 


**Effect of input distribution** 

- Training distribution.
For training the model, we replace the distribution U[-1, 1] by a standard normal distribution (i.e. X_ij ~ N(0, 1) for input X). We can indeed recover the sudden drop in loss, where the MSE before and after the drop matches the values obtained for U[-1, 1]. Please see Fig. D in the attachment for reference.


- Test distribution.
For test time OOD performance, we find that replacing the distribution of input matrix entries from U[-1, 1] (on which the model is trained), with Gaussian (mean=0, stdev=0.5) and Laplace (location=0, scale=0.5) does not lead to any considerable change 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.

Please note that we already check OOD performance when varying p_mask (Fig 3 in manuscript). Further, we also check OOD performance by changing the number of rows / columns and rank of the input matrix – as observed in Fig. (A–C) in attached PDF, BERT 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 because positional embeddings in our model depend on the column index of the element, and hence (expectedly) changing the number of columns adversely affects model performance. The overall idea is that as long as the input matrix entries do not exceed beyond a certain threshold in magnitude, the model can solve matrix completion. This is because the model has learnt token embedding representations for a fixed range of input values (as defined by U[-1, 1]). It is unexpected to generalize to larger entries that the model does not see during training, since the model has no knowledge of how tokens for larger entries are represented in the embedding space. This also matches the observations in  Sec 5 [3], where the author shows that changing the variance of test distribution significantly affects OOD performance of the model.

[3] https://openreview.net/forum?id=L2a_bcarHcF

Meta-review: This submission has received scores of 6, 4, 7, 6, 5, indicating a weak acceptance of this submission. The reviewers generally find the paper interesting and well-presented, with a novel approach to studying the behavior of transformer models on a simplified matrix completion task. However, some concerns remain about the significance of the contributions and the generalizability of the findings. Still, the reviewers' scores and comments, along with the authors' detailed and responsive rebuttal, suggest that the paper's strengths outweigh its weaknesses. While there is room for further exploration and deeper mechanistic understanding, the paper makes a valuable contribution to the field and warrants acceptance.

The authors are encouraged to address the reviewers' concerns in the final version by expanding the discussion on the significance and generalizability of their findings, exploring the possibility of extending their analysis to larger-scale tasks and different model architectures, and providing a deeper theoretical or mechanistic understanding of the observed phenomena.

Strengths:
- Novelty and clarity: The paper presents a novel and well-executed approach to studying transformer behavior on a simplified task, making it easy to follow and understand.
- Interesting observations: The phase transition phenomenon and the analysis of attention heads and embeddings provide valuable insights into the model's learning dynamics.
- Thorough experimentation: The authors conduct extensive experiments and provide detailed descriptions of their methodology.

Areas to Improve:
- Significance of contributions: Some reviewers question the overall impact and generalizability of the findings, particularly regarding the transferability of insights to the language domain and larger-scale tasks.
- Depth of mechanistic understanding: While the paper offers valuable observations, a deeper theoretical or mechanistic analysis of the observed phenomena is desired.
- Generalizability: The experiments are primarily conducted on small matrices, raising concerns about the scalability and applicability of the findings to more complex scenarios.
Decision: Accept (poster)