TransformersコミュニティがModernBERTの半精度オーバーフローを報告、長文の平均プーリングでNaNが発生する可能性

Transformersコミュニティは11月1日、長文分類に関する数値上の問題を報告した。ModernBertForSequenceClassificationでclassifier_pooling="mean"を指定し、完全なfloat16で実行すると、エンコーダーの出力が正常でも分類スコアがNaNになる場合がある。報告された環境はTransformers 5.22.1、PyTorch 2.15.0で、Appleシリコン搭載コンピューターのCPUを使って検証された。問題報告

原因となるのは平均プーリングの計算順序だ。隠れ状態にアテンションマスクを掛け、トークン方向に合計してから、有効なトークン数で割る。報告者によると、中間の合計値がfloat16で表現できる範囲を超えると、その後に除算しても正常な数値には戻らない。公開されたテストでは、mmBERT-smallは2,048トークン、ModernBERT-baseは6,144トークンでNaNが発生した。これらの長さは特定の合成英語入力で得られたものであり、あらゆるテキストに共通する固定のしきい値ではない。再現コードと出力

この報告は、長文書を扱うワークフローに実務上の影響を及ぼす。公式ドキュメントによると、ModernBERTは最大8,192トークンに対応し、分類プーリングではclsまたはmeanを選べる。そのため、入力がモデルの長さ制限内であっても、プーリング段階の数値安定性を別途検証する必要がある。エンコーダー出力が有限値かどうかを確認するだけでは、分類ヘッドの前に起きる異常を見逃す可能性がある、というのが技術的な示唆だ。ModernBERTドキュメント

報告者が示した比較では、合計をfloat32で計算し、除算後に元の精度へ戻すと、テストで有限値が得られた。PyTorchのドキュメントにも、torch.sumのdtype引数で演算前に入力を変換でき、データ型のオーバーフローを防げると明記されている。これはプーリング部分の精度を局所的に調整する方法となるが、デプロイ前には性能と出力誤差を確認する必要がある。PyTorch APIドキュメント

確認時点では、議論に参加した別のユーザーが対応する意向を示しているものの、メンテナーによる確認や関連する修正の提出はない。報告者はCUDAでのテストを行っておらず、分類ヘッドも新規に初期化したものだ。そのため、この結果は既存の分類タスクの精度評価を示すものではない。今後は正式な回帰テストと、完全な半精度および自動混合精度の実行経路における違いを確認する必要がある。議論の状況