Paper: The Closeness of In-Context Learning and Weight Shifting for Softmax Regression (ICLR-2024)

Reviewer comment (Theory):
Definition 1.3: I miss the motivation for why this would advance the understanding of in-context learning for the Transformer. It seems that the problem solved by Self-Attention would involve the matrix A quadratically in the exponential, while here A appears linearly in the exponential. Could you elaborate on why this intermediate step is useful for analyzing what would happen in the Transformer?

Author response:
The motivation behind Definition 1.3 (Softmax Regression) is to bridge the gap between the simple linear models studied in prior work [1] and the more complex softmax-based models used in actual LLMs​​. While in a full Transformer model, matrix 'A' appears quadratically in the exponential, our setting of 'A' in the exponential allows for a tractable analysis while still capturing essential characteristics of the softmax function. This formulation makes our analysis more relevant to actual Transformer models in LLMs than purely linear models, yet more manageable than full quadratic models.

Reviewer comment (Theory):
Upon reading the text a few times, I still do not understand why the bounds of **Thm 5.1** and **Thm 5.2** would imply that the transformation induced by the layer would approximate the gradient step, or if I understood **Oswald et. al** correctly, that at least there is a choice of layer parameters that would make it approximate the gradient step.

Author response:
Thanks for your careful reading. Let us emphasize the high level idea of our results here. A self-attention layer update can be treated as an update to the tokenized document $A$ [1]. The gradient descent can be treated as an update to the model parameter $x$ here. In the softmax regression loss function, these  two update can be treated as the change of the value $b$ here. Hence, by showing the similarity between the change of $b$, we show the closeness between them. In [1], they show the change of some specified parameters by gradient descent approximate the ability of in-context learning. This does not contradict the results stated in [1] but rather adds a new dimension to understanding the interplay between layer transformations and gradient steps.

Reviewer comment (Novelty):
The significance of the findings, such as the upper bounds of data transformation, is somewhat understated, rendering them supplementary to prior research.

Author response:
We appreciate the feedback and will take steps to highlight the importance and novelty of these findings more clearly in our revised manuscript. Our work provides a significant extension to the existing body of research, and we aim to ensure that this is adequately reflected and understood by our readers.

Reviewer comment (Presentation):
Certain mathematical proofs and deductions may benefit from a more concise presentation, potentially relocating them to the appendix, while experimental results could find a more prominent place in the main paper.

Author response:
Thanks for the feedback on the presentation of the paper. We will definitely consider your suggestions to improve the clarity and flow of the paper, allowing readers to more easily grasp the theoretical framework before delving into the specific results and their implications.

Reviewer comment (Novelty):
More comparison with the work from Deng et al. (2023b) is needed, which seemingly to be the work most closely related to this work.

Author response:
We believe our contributions are distinct and significant compared to Deng et al. (2023b). Concretely, Deng et al. (2023b) focuses primarily on softmax regression in a broader context, whereas our work specifically studies the in-context learning capabilities of transformer-based large language models from a softmax regression perspective. Our main results in Section 5, including the Lipschitz bounds and the bounded shift for in-context learning (Theorems 5.1 and 5.2)​​, are unique contributions that extend beyond the scope of Deng et al. (2023b). We will add a more detailed comparison in the revised paper to clarify these differences.

Reviewer comment (Presentation):
Organization of the content is preferred to be more focused on what to be proven. i.e. "why bounding the single step of $x$ and $A$ relates to clarifying the relationship between in-context learning and softmax weight shift."

Author response:
We appreciate the feedback on the organization of our paper. In the revision, we will restructure the content to focus more on our primary contributions.

Reviewer comment (Writing):
Typo: Lipschtiz → Lipschitz

Author response:
Thank you for pointing out the typo. We will correct "Lipschtiz" to "Lipschitz" in the revised paper.