Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality

要約

マルチタスク線形回帰のインコンテキスト学習のためのマルチヘッド ソフトマックス アテンション モデルをトレーニングするための勾配フローのダイナミクスを研究します。
適切な初期化の選択の下で、勾配流の全体的な収束を確立します。
さらに、勾配流れダイナミクス中に、各注意ヘッドがマルチタスク モデルの 1 つのタスクを解決することに集中する、興味深い「タスク割り当て」現象が現れることを証明します。
具体的には、勾配流のダイナミクスが 3 つのフェーズに分割できることを証明します。損失がかなりゆっくりと減少し、注意を向けるヘッドが徐々に個別のタスクへの傾向を強めるウォームアップ フェーズと、各ヘッドが単一のタスクを選択する創発フェーズです。
タスクと損失が急速に減少し、注意パラメータが限界まで収束する収束フェーズになります。
さらに、勾配流によって学習された制限モデルが、定数係数までは可能な限り最良のマルチヘッド ソフトマックス アテンション モデルと同等であるという意味で、勾配流の最適性を証明します。
私たちの分析では、ICL の予測精度に関して、単頭注意モデルと多頭注意モデルの厳密な区別も明らかにしています。
収束解析の重要な手法は、パラメータ空間の勾配流れダイナミクスをスペクトル領域の一連の常微分方程式にマッピングすることです。ここで、注意の重みの半特異値の相対的な大きさがタスクの割り当てを決定します。
私たちの知る限り、私たちの研究はマルチヘッド ソフトマックス アテンション モデルの最初の収束結果を提供します。

要約(オリジナル)

We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We establish the global convergence of gradient flow under suitable choices of initialization. In addition, we prove that an interesting ‘task allocation’ phenomenon emerges during the gradient flow dynamics, where each attention head focuses on solving a single task of the multi-task model. Specifically, we prove that the gradient flow dynamics can be split into three phases — a warm-up phase where the loss decreases rather slowly and the attention heads gradually build up their inclination towards individual tasks, an emergence phase where each head selects a single task and the loss rapidly decreases, and a convergence phase where the attention parameters converge to a limit. Furthermore, we prove the optimality of gradient flow in the sense that the limiting model learned by gradient flow is on par with the best possible multi-head softmax attention model up to a constant factor. Our analysis also delineates a strict separation in terms of the prediction accuracy of ICL between single-head and multi-head attention models. The key technique for our convergence analysis is to map the gradient flow dynamics in the parameter space to a set of ordinary differential equations in the spectral domain, where the relative magnitudes of the semi-singular values of the attention weights determines task allocation. To our best knowledge, our work provides the first convergence result for the multi-head softmax attention model.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google