Transformersコミュニティ、Mamba2／Zamba2のキャッシュデコード差異を調査　時間刻み幅の制限が適用されていない可能性

Transformersに寄せられた状態空間モデルに関する報告を受け、プロンプトの事前計算とトークン単位の逐次デコードの整合性が調査対象となっている。問題は10月9日に報告され、翌日には別のコミュニティメンバーがMamba2とZamba2でlogitsの差異を再現したと報告し、ほかのモデルも関連するキャッシュ経路を共有していると指摘した。確認時点でissueはオープンのままで、ページには関連する修正PRは掲載されていない。コミュニティの報告

報告者が調べたコードでは、シーケンス全体をスキャンする処理は、時間刻み幅 dt にバイアスを加え、softplusを適用した後、time_step_limit の範囲に制限する。一方、キャッシュを使った単一トークンの更新は、更新関数がバイアスの加算とsoftplusを処理するものの、同じ範囲制限が抜けている。そのため、同じ入力でも、シーケンス全体を一括処理する場合と、プレフィックスを処理してキャッシュから続行する場合とで、異なる状態更新式が使われる可能性がある。元の検証はCPUおよびCUDA上の独立したPyTorch関数を対象としており、事前学習済みモデルの生成品質は測定していない。再現手順と修正案

公式ドキュメントでは、Mamba2の time_step_limit は時間刻み幅のクリップ範囲として定義され、既定値はゼロから無限大までとなっている。このため、有効な境界値が設定されているかどうかは、原因を調べるうえで重要な条件となる。キャッシュには状態空間の状態と畳み込み状態が保存され、後続トークンの処理に再利用される。ドキュメントは、カーネルや演算順序、低精度演算の違いによって数値差が生じる可能性も注意喚起しているため、テストでは通常の誤差と今回の制限適用漏れを区別する必要がある。Mamba2公式ドキュメント

エンジニアリングチームにとっての技術的リスクは、オフラインでの一括評価と、オンラインでのキャッシュ推論が一致しない可能性だ。これは報告された仕組みから導かれるリスクであり、回答品質が全般的に低下したことを示すものではない。重み、入力、精度、パディング方法を固定し、全体スキャンとキャッシュからの続行でlogitsと最終状態を比較することが推奨される。提案されているのは、単一トークン用カーネルの前にバイアス加算、softplus、クリッピングを一貫して適用し、カーネル側での重複処理を避けることだ。次の段階では、完全なモデルと高速化バックエンドで検証する必要がある。正式リリースが影響を受けるかどうかも、個別に確認する必要があり、開発ブランチでの事例だけから推定することはできない。issueでの検証範囲