Stick-breaking Attention

要約

セルフ アテンション メカニズムは伝統的にソフトマックス オペレーターに依存しており、トークンの順序を考慮した RoPE や位置バイアスのような位置埋め込みが必要です。
しかし、長さを使用する現在の方法では、依然として一般化の課題に直面しています。
スティックブレイクプロセスに基づいた代替アテンションメカニズムを提案します。現在のトークンの前の各トークンに対して、現在のトークンに割り当てる残りのスティックの割合を表すブレークポイント $\beta_{i,j}$ を決定します。
。
スティックが完全に割り当てられるまでこのプロセスを繰り返し、一連の注意の重みが得られます。
このプロセスには当然、文法解析に対する言語的動機を持つ最新性バイアスが組み込まれています (Shen et. al., 2017)。
我々は、従来のソフトマックスベースの注意メカニズムをスティックブレイキング注意に置き換えることの意味を研究します。
次に、数値的に安定したスティックブレイク アテンションの実装について説明し、このメカニズムに対応するようにフラッシュ アテンションを適応させます。
現在のsoftmax+RoPEアテンションシステムのドロップイン代替品として使用すると、スティックブレイキングアテンションは長さの一般化と下流のタスクに関して現在の方法と競合して機能することがわかりました。
スティックブレイクは長さの一般化でもうまく機能し、$2^{11}$ コンテキスト ウィンドウでトレーニングされたモデルが $2^{14}$ でうまく機能し、複雑さが改善されます。

要約(オリジナル)

The self-attention mechanism traditionally relies on the softmax operator, necessitating positional embeddings like RoPE, or position biases to account for token order. But current methods using still face length generalisation challenges. We propose an alternative attention mechanism based on the stick-breaking process: For each token before the current, we determine a break point $\beta_{i,j}$, which represents the proportion of the remaining stick to allocate to the current token. We repeat the process until the stick is fully allocated, resulting in a sequence of attention weights. This process naturally incorporates recency bias, which has linguistic motivations for grammar parsing (Shen et. al., 2017). We study the implications of replacing the conventional softmax-based attention mechanism with stick-breaking attention. We then discuss implementation of numerically stable stick-breaking attention and adapt Flash Attention to accommodate this mechanism. When used as a drop-in replacement for current softmax+RoPE attention systems, we find that stick-breaking attention performs competitively with current methods on length generalisation and downstream tasks. Stick-breaking also performs well at length generalisation, allowing a model trained with $2^{11}$ context window to perform well at $2^{14}$ with perplexity improvements.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google