Abrupt Learning in Transformers: A Case Study on Matrix Completion

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.

Strengths

1. The paper presents an interesting approach by framing matrix completion as an MLM task. 2. The experiments are detailed, with clear descriptions of data preprocessing, training, and performance metrics. 3. 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. 2. There is insufficient explanation of the internal mechanisms by which the model learns matrix decomposition, particularly in the second stage. 3. The experiments are conducted on relatively small-scale matrices (up to 15x15), raising concerns about the scalability and generalizability of the findings.

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? 2. 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? 3. 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? 4. Can the approach be scaled to larger matrices, and what modifications would be necessary to achieve this?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

See weaknesses.

Summary

The authors study how encoders perform at the task of matrix completion. They generate synthetic data and train encoders of different sizes to predict masked out tokens. They observe an interesting behaviour where initially the loss appears to be at a plateau but then drops. Their 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.

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) 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?

Questions

Did you consider other distributions?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

See weaknesses.

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.

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. - 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? - 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. - 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? - 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? - Any discussion on how the results might change for auto-regressive training will be interesting.

Questions

Please check my questions above. [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.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

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.

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

Strengths

It's an easy to read paper with good figures and sound methodology. The 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.

Questions

The authors do a good job of anticipating many questions. Does this behavior generally to other mathematical tasks? Can we get any insight about the actually algorithm or implicit optimization procedure? What is the generalization to larger matrices? Is there a way to more rigorously verify the hypotheses from the **Attention Heads with Structured Mask** section?

Rating

4

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

The authors acknowledge the limitations in the discussion. This only works on small matrices and is not a generally useful matrix solver.

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: + 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. + During the second, different attention heads exhibits different behaviors correspond to different parts of the input: masks, observed values, position embeddings. + The learned token embeddings exhibits clusters between elements from the same row in the original matrix, even when not row/column marker explicitly providing. These findings aligns with prior studies focusing on masked-language-modeling(MLM) tasks and provides valuable insights into the learning mechanism of transformers.

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. + The paper performed extensive experiments and examine the learning behaviors of transformers thoroughly across learning stages and important components of transformer. + 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. + Despite adapting a new task, the main patterns the paper found in through experiments are similar to prior works.

Questions

+ What is the purpose of training the model to predict all values instead of just the masked ones? + What do the two panels in figure represent? The right panel caption appears incorrect and the loss surprisingly goes down as p_mask increases. + 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.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors have addressed societal impact and limitations in the conclusion section.