注意力が盲目になるとき：ALiBi位置エンコーディングにおける数値的障害

要旨

我々は、ALiBi位置エンコーディングのこれまで見落とされていた故障モードを特定する。すなわち、その線形バイアススケーリングが浮動小数点精度のアンダーフローを引き起こし、アテンション重みの大部分をゼロにして、影響を受けたアテンションヘッドを部分的に盲目にする。我々はこの故障モードを分析し、その影響を特徴づけ、4つの緩和戦略を検討する。さらに、ALiBiに基づく最先端の事前学習モデルにおいてこの故障モードが発生することを実証する。148Mパラメータのデコーダモデルを用いた包括的な事前学習実験により、その影響をコンテキスト外劣化から切り離して理解することができる。我々は、ALiBiの故障モードが標準的なデコーダベンチマークにわずかな影響しか与えない一方で、トークン検索を著しく損なう可能性があることを見出す。我々は4つの学習時緩和戦略を提案し、それらを個別および組み合わせて評価したところ、対数スケールの距離がパスキー検索において最も一貫した改善をもたらすことが分かった。この問題にもかかわらず、デフォルトのALiBiスロープは、特にneedle-in-a-haystack検索において、驚くほど強力なベースラインであり続ける。これらの知見に基づき、ALiBiを用いてモデルを学習する方法について具体的な推奨事項を提供する。

English

We identify a previously overlooked failure mode of ALiBi positional encoding: its linear bias scaling underflows floating-point precision, which zeroes out a large fraction of attention weights and renders the affected attention heads partially blind. We analyze this failure mode, characterize its impact, and examine four mitigation strategies. We further demonstrate its occurrence in state-of-the-art pretrained models based on ALiBi. Comprehensive pretraining experiments with 148M-parameter decoder models help us to disentangle its effects from out-of-context degradation. We find that ALiBi's failure mode can substantially impair token retrieval while having only a minor effect on standard decoder benchmarks. We propose four training-time mitigation strategies and evaluate them individually and in combinations, finding that log-scaled distances yield the most consistent improvements in passkey retrieval. Despite this problem, default ALiBi slopes remain a surprisingly strong baseline, particularly for needle-in-a-haystack retrieval. Based on these findings we provide concrete recommendations on how to train models with ALiBi.