Paper: The mechanistic basis of data dependence and abrupt learning in an in-context classification task (ICLR-2024)

Reviewer comment (Theory):
Why do the authors switch from the p_C = 0 case to the p_C > 0 case halfway through the paper? I find this confusing as it makes the relevance of the mechanistic analysis to the p_C = 0 regime unclear.

Author response:
We use p_C > 0 to ensure that the network robustly and quickly learns the ICL solution. This was simply so that we could run each experiment for a shorter period and run more experiments overall. We now clarify this in the text. Note that we focus on ICL dynamics and not the ICL vs IWL tradeoff in this section and subsequent ones.

Reviewer comment (Theory):
What are the authors' intuition as to why does \xi undergo a phase transition driving ICL?

Author response:
Thank you for the kind comments. Our model suggests that \xi changes abruptly due to the sequence of three nested logits that make up an induction head (see the Abrupt transitions during ICL section of the Discussion). We can potentially accelerate learning of \xi (or similar metrics for other tasks) by using a curriculum. We have added a few sentences in the discussion about the use of curricula in the Implications for LLMs section.

Reviewer comment (Experiments):
More consideration of what will happen in a larger model will be useful for the field. Clearly, things will not work just in the way they do in these reductions when the task is learned in a larger transformer. While fuller characterization of that will be a task for future work, noting this issue as a limitation of the present effort and considering how these results inform us about what is happening in LLMs will be valuable.

Author response:
Thank you for the kind comments and valuable suggestions. Here is a summary of the changes we have made in response to the suggestions:

2. Figure 8 has been moved to the Appendix (now Figure A.4) to provide more space for an expanded discussion. We contextualize our work, comment on implications for LLMs and briefly discuss limitations.

Reviewer comment (Theory):
There is a lack of intuitive understanding of why the loss landscape of the 3-parameter model does not have a saddle point at the point where all three parameters are equal to 0. Perhaps this understanding is latent in the equations and could be worked out with effort, but it would be useful if the authors provided such an intuitive explanation to help understand why the behaviors of the 2- and 3-parameter reduced models might or might not be applicable to the full model.

Author response:
Thank you for the kind comments and valuable suggestions. Here is a summary of the changes we have made in response to the suggestions:

3. A few sentences have been added to provide some intuition (last two paragraphs before the discussion). Note that the origin is not a critical point (and thus not a saddle) of the loss except when L = N. This is because the network gradually aligns the regression vectors with their corresponding labels (increasing $\xi$) when learning to randomly pick one of the contextual labels in the slow phase.

Reviewer comment (Writing):
The paper is harder to read than it should be due to failure to account for the difficulty of extended chains of arbitrary associative bindings requiring long-distance leaps across context. These bindings are at the heart of the mechanisms investigated but are hard for human readers when arbitrary, as they often are in this paper.

Author response:
Thank you for the kind comments and valuable suggestions. Here is a summary of the changes we have made in response to the suggestions:

4. We have removed the arbitrary labels for progress measures and hypotheses, shrunk/split long sentences and moved the figures. The text has been revised for clarity. The current placement of the figures may still not be ideal (esp fig 4), but the page limit makes it hard to move figure 4 around.

Reviewer comment (Presentation):
The use of terms like "the former" vs "the latter" as referring expressions, arbitrary labels (a–d) for key phenomena, random ordering of the assignments of these labels to lines in graphs, arbitrary labels for hypotheses (I–V), and placement of figures (especially Figure 4) far from their discussion points makes the paper less accessible. Although conclusions are ultimately stated in conceptual terms, there should be engagement with this conceptual structure in the referential expressions used. For example, H3 could be abbreviated as sCLA → ILA+TILA (slow-learned context-label attention → Item-label attention and Target-item-label association). Arbitrary label schemes (a, b, c, d and I–V) should be dropped, H4 and H5 should be expressed directly, and the order of defining symbols should correspond to their order of appearance in the hypotheses.

Author response:
Thank you for the kind comments and valuable suggestions. Here is a summary of the changes we have made in response to the suggestions:

4. We have removed the arbitrary labels for progress measures and hypotheses, shrunk/split long sentences and moved the figures. The text has been revised for clarity. The current placement of the figures may still not be ideal (esp fig 4), but the page limit makes it hard to move figure 4 around.