Transformersコミュニティ、Mambaのキャッシュ継続計算における異常を再現　複数トークン入力で再帰状態が失われる可能性

Transformersでは10月31日、Mamba-1の正確性に関する報告が寄せられた。モデルにテキストを処理させてキャッシュを保持し、その後複数の新しいトークンをまとめて入力すると、出力がシーケンス全体を一度に処理した場合と異なる可能性がある。同日、別のコミュニティ参加者がメインブランチで再現したと報告した。確認時点で、このissueはオープンのままで、修正PRへのリンクもなかった。公開issue

これはキャッシュインターフェースの実際の意味に関わる問題だ。公式ドキュメントによると、cache_paramsには状態空間モデルと畳み込みの状態が含まれ、これを渡すと新しい入力は既存のコンテキストを引き継ぎ、古いトークンを再度入力せずに処理できる。したがって、エンジニアがこのインターフェースを使って分割プリフィルや会話の継続計算を実装する場合は、分割の前後で数値が一致するか確認する必要がある。Mambaドキュメント

報告者は、Transformers 5.20.2、CPU、float32を使用し、2つのネイティブ高速化パッケージはインストールしていなかった。mamba-130m-hfで、全入力を一度に処理する方法、1トークンずつ継続処理する方法、分割して継続処理する方法を比較した。逐次処理でのlogitsの最大差は約0.000263だった一方、分割処理では約115.897に達した。後者では、最も高いスコアのトークンが一致する割合は71.4%だった。これらは小規模な再現例での数値であり、実タスクの精度低下にそのまま換算できるものではない。再現コードと結果

公開コードを見ると、再帰状態が存在し、かつシーケンス長が1の場合に限って状態更新の分岐に入る。複数トークンの入力は全スキャンの経路に進み、その呼び出しでは既存の再帰状態が渡されていない。参照用のスキャン実装は状態をゼロで初期化し、処理後にキャッシュを上書きする。この点は報告された根本原因を裏付けるが、すべてのバックエンドで実測済みという意味ではない。Mamba実装

デプロイ側では、同じ入力を全入力・分割・逐次の各経路で処理し、リクエストが成功したかだけでなくlogitsを比較するチェックを追加するとよい。今後は、修正によって既存状態が引き継がれるか、回帰テストが追加されるか、高速化カーネルがどう扱われるかを追う必要がある。コミュニティからは、同様のコードを共有する他のモデルも調べるべきだとの声が出ているが、影響範囲はまだ検証されていない。その後の議論