Summary
This work studies the mechanism for relevant token selection in the attention model by drawing connections with the implicit bias literature and max-margin SVM formulation. The authors consider the prompt attention model $f(X)=v^TX^T\text{softmax}(XW^Tp)$, with tokenized input $X$, value weights $v$, key-query weights $W$, and tunable token/prompt $p$. They show that running gradient descent (GD) on $p$ and $W$ (for fixed $W$ and $p$, respectively) is equivalent, so they consider optimizing only $p$ for fixed $W$.
They consider a data setting where the quality of $t^{th}$ token of input $X$ is determined by its scores $Yv^Tx_t$, with the globally optimal tokens being the ones with the highest scores.
First, they consider optimizing $p$ when $v$ is fixed using a loss that is decreasing and smooth. In this setting, they show the following results:
- Under an assumption on the token scores, $p$ converges (in direction) to the global max-margin solution, which separates the globally optimal tokens from the rest.
- They also develop a regularization path analysis to show global convergence by relaxing the said assumption.
- The main result shows that with an appropriate initialization and a small enough step size, the GD iterates of $p$ converge (in direction) to a local max-margin solution that separates the locally-optimal tokens from the rest. Here, locally optimal tokens are defined as the ones which have higher scores than their SVM neighbors.
Next, they consider the joint optimization of $p$ and $v$ using logistic loss. In this case, for a given $p$, if the resulting features are separable, $v$ has an implicit bias to converge to the max-margin solution as the problem is linear in $v$. Here, optimal tokens as the ones that maximize the downstream label margin. They show that:
- When the attention features (for the max-margin $p$) are all support vectors (for the respective max-margin $v$), both $v$ and $p$ converge to their respective max-margin solutions.
- When this is not the case (i.e. the attention features are not all support vectors), $p$ asymptotically selects one token per input, and it suffices to select tokens with the highest scores while also mixing other tokens. This does not impact the margin of $v$, which still converges to the max-margin solution.
Throughout, the authors illustrate these results through numerical experiments.
Strengths
1. This work gives interesting theoretical insights into the mechanism for relevant token selection in the attention model.
2. It lays the groundwork to analyze attention mechanism using the lens of implicit bias and motivates several interesting directions for future work.
3. Overall, it is an interesting paper, with clear exposition and well-connected ideas. It makes several meaningful contributions (as listed in the summary).
Weaknesses
1. The experimental results seem limited and the paper would benefit from the inclusion of more experiments.
2. Some points of discussion can enhance intuition, and certain aspects regarding figures need some clarification.
Questions
1. Experiments:
- The numerical experiments illustrate the theoretical results well. However, it would be good to include additional experiments on semi-synthetic or real datasets, such as those considered in [1], [2].
- Fig. 1(b) illustrates the local convergence of $p$ in the red or blue direction, depending on the initialization. It would be helpful to see if there are cases when the GD iterates do not converge to either of the two solutions shown in Fig. 1(b). This would give some insight into the gap between the empirical observations and the theoretical result for this case.
2. Discussion/clarification:
- Points of discussion:
- Generally, margin maximization refers to separating samples/features from two classes, whereas in this work, the main contribution is to show that the tunable prompt converges to a solution that separates the globally/locally optimal tokens from non-optimal ones. I suggest including more discussion on this part in the introduction. Relatedly, it would help to clarify lines 36-37.
- In Section 2, choice of $v$ determines the scores, and hence the solution learned by $p$. Some discussion on this would be nice.
- The data setting is interesting, but since there are a lot of cases, it would be helpful to show some connections between (some of) the synthetic data setting(s) and some datasets that are commonly used in practice.
- In the discussion on convergence to the local max-margin soution, and the description of Fig. 1(b) in Section 2.2, it would be helpful to clarify that depending on the initialization, the GD iterates will converge to either the global max-margin (when initialized in the cone associated with that solution) or the local max-margin.
- Figures:
- In Fig. 1(a), it is unclear what the role of the blue line (local max-margin solution) is, since both the non-opt tokens take the same value.
- Fig. 1(c) description needs some clarification. The separating hyperplane in. the figure looks like it's the max-margin solution for the teal datapoint, but not the max-margin across all three colors (margin with green is small).
- In Figs. 2(b) and (c), it would be helpful to have the legend for red and blue curves associated with $p$ and $v$, respectively.
- Other:
- There are some typos/inconsistencies in the proof of Lemma 1 that should be corrected.
- Theorem 2 shows global convergence of $p$ via regularization path analysis. It can be moved above assumption B to improve flow, as it is more general.
- In Section 2.2, it would be helpful to use some specific notation for the solution that separates the locally optimal tokens from the rest ($p^{mm}(\alpha)$), such as $p^{mm}_l$.
- Fig. 2(a) comparing the transient dynamics for correlation and logistic loss is interesting. However, it needs a minor clarification in the description. It is stated that when $p$ selects the optimal token, the gradient norm $\propto \gamma_i$ for correlation loss, and $\propto \gamma_ie^{-\gamma_i}$ for logistic loss, and we can compare loss for tokens with different scores. However, if $p$ selects the optimal token, the score would be fixed. Maybe, it can be rephrased to “if $p$ selects the token with score $\gamma_i$”.
- In Section 3.1, label margin is defined as $\frac{1}{||v_{mm}||}$, but in the example, label margin $\gamma$ is defined as $||v_*||$. It would be helpful to mention what $v_*$ is for clarity.
References:
[1] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
[2] Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015, 2023.
Rating
8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.