技術的概要：FAR（IMC 対応推論のための機能保存型アテンション置換）

1. 問題定義

自己アテンションに依存するトランスフォーマーアーキテクチャは、ReRAM 交叉配列などの新興**メモリ内計算（IMC）**ハードウェアに展開される際、重大な効率性のボトルネックに直面します。自己アテンションは、グローバルなトークン統合を可能にすることで並列デジタルプロセッサ（GPU/TPU）上で卓越した性能を発揮しますが、その計算特性は IMC の実行モデルと根本的に矛盾します：

• 活性化同士の乗算： 自己アテンションは、$QK^\top$ や softmax 正規化のような操作に大きく依存しており、これらは重み固定のベクトル行列乗算（VMM）を最適化する IMC ハードウェアではなく、動的で非局所的な活性化の混合を必要とします。
• メモリオーバーヘッド： トークン相互作用の二次的な複雑さ（$O(T^2)$）と、空間的に分散した活性化を繰り返し読み取る必要性は、断片化されたアナログ操作、広範なアナログ - デジタル/デジタル - アナログ（ADC/DAC）変換、および交叉配列の利用率の低下を招きます。
• 現在の限界： 既存の IMC アクセラレータは、しばしばアテンションをデジタルユニットにオフロードして二次的なボトルネックを未解決のまま放置するか、あるいは本質的なデータ移動コストに対処することなく、アテンションを直接アナログ配列にマッピングしようとします。

同時に、最近の分析は、推論中にアテンション層が著しい冗長性を示し、アテンションマップが深層を通じて滑らかになり圧縮可能になることを示唆しています。これは、アテンションの機能的役割が、IMC の局所性と重み固定性に適合する逐次モジュールによって近似可能であるという仮説を動機付けます。

2. 手法：FAR フレームワーク

著者は、事前学習されたビジョントランスフォーマー（特に DeiT ファミリー）のすべての自己アテンションブロックを逐次モジュール、具体的には**マルチヘッド双方向 LSTM（BiLSTM）に置換するフレームワークであるFAR（機能保存型アテンション置換）**を提案します。このアプローチは、ゼロから訓練する必要のない「ドロップイン」型のアーキテクチャ変換として設計されています。

2.1 置換戦略

• ブロック単位置換： 事前学習されたトランスフォーマーの各アテンションブロック $A_l$ は、学習可能な代替モジュール $A'_l$（マルチヘッド BiLSTM）に置換されます。他のすべてのコンポーネント（MLP ブロック、パッチ埋め込み、レイヤー正規化）は凍結されたままです。
• 蒸留による機能的同等性： 置換はレイヤーごとの蒸留を通じて監督されます。各レイヤー $l$ において、類似度損失（$L_{sim}$）は、元のアテンションヘッドの出力と新しい BiLSTM ヘッドの出力との差を最小化します：
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  これにより、代替モジュールが教師の表現構造を複製することが保証されます。
• 訓練プロトコル： このプロセスは 2 つの段階を含みます：
  1. 蒸留フェーズ： モデルの残りを凍結したまま、代替 BiLSTM モジュールのみを、類似度損失とタスク損失（分類）の組み合わせを用いて訓練します。
  2. 微調整フェーズ： 残りの精度低下を回復させるために、モデル全体を共同で微調整します。

2.2 マルチヘッド BiLSTM アーキテクチャ

マルチヘッドアテンションの帰納的バイアスを保存するために、代替モジュールはマルチヘッド BiLSTMとして構造化されます：

• 投影と分割： 入力埋め込みは投影され、$N$ つの部分空間に分割されます（教師のアテンションヘッドの数に一致）。
• 双方向処理： 各部分空間は、前方および後方のトークン依存性の両方を捉える個別の BiLSTM ヘッドによって処理されます。
• 結合と投影： 出力は結合され、元の次元に投影されます。
• IMC 互換性： アテンションとは異なり、この構造は再帰的な重みの再利用と逐次的なデータフローに依存しており、すべてのペアの活性化相互作用を排除します。これは、重みが固定されデータが逐次的に流れる ReRAM 交叉配列に自然にマッピングされます。

2.3 IMC 制約のための構造化プルーニング

リソース制約のある IMC 配列へのさらなる適応のため、著者は BiLSTM ヘッドの隠れ次元に構造化プルーニングを適用します：

• DeepHoyer の拡張： LSTM 隠れユニットを標的とした DeepHoyer フレームワークを拡張します。LSTM ゲート（入力、忘却、セル、出力）は隠れユニットを共有するため、単一のユニットをプルーニングするには、関連するすべての重み行列（$W_{ih}$、$W_{hh}$、および投影重み）にわたって調整された削除が必要です。
• グループスパース性： 各隠れユニットに対応する重みのグループに対して微分可能な Hoyer スパース性尺度を適用し、ネットワークがどのユニットが冗長であるかを学習することを促します。
• レイヤー適応型圧縮： プルーニングはヘッドごと、方向ごとに独立して実行され、モデルが各レイヤーの特定の冗長性に合わせて容量を適応させることを可能にします。その結果、IMC 配列サイズに適合するコンパクトなモジュールが生成されます。

3. 主要な結果

このフレームワークは、ImageNet およびさまざまな下流タスク（CIFAR-10/100、Cars、Flowers、iNaturalist）における DeiT ファミリー（Tiny、Small、Base）で評価されました。

• 精度の維持：
  • DeiT-Tiny： FAR-Tiny は、ImageNet 上でベースラインの DeiT-Tiny を +1.2% 上回る（73.3% 対 72.2%）性能を示し、逐次代替モジュールが小規模モデルにとって有益な帰納的バイアスとして機能することを示唆しています。
  • 大規模モデル： DeiT-Small および Base については、FAR は教師と比較して約 1〜2% のわずかな差で競争力のある精度を維持し、アテンションの機能的役割が逐次的再帰によってほぼ再現可能であることを実証しました。
  • 下流タスクへの転移： FAR モデルは下流の分類タスクへの強力な転移性を維持しており、学習されたシーケンス間マッピングが効果的に一般化することを示しています。
• プルーニングへの頑健性：
  • 構造化プルーニングを適用すると、平均隠れ次元が 40〜60% 削減され、精度の低下は追加で 1〜2% のみでした。
  • プルーニングパターンは不均一で、中間レイヤーでは保持率が高く、ネットワークの端ではより攻撃的なプルーニングが行われており、表現の冗長性の内部の不均一性を反映していました。
• ヘッドごとの特化： トークン相互作用マップの可視化は、FAR ヘッドが元のアテンションヘッドの明確な意味的役割を成功裡に継承し、冗長なパターンに収束するのではなく、特定の物体領域に集中していることを示しました。

ハードウェア効率分析

分析用 ReRAM IMC モデルを用いて、著者は遅延とエネルギーを推定しました：

• 遅延： ReRAM IMC 上の FAR は、二次的なアテンションオーバーヘッドを回避するため、ReRAM IMC 上の DeiT より約 18 倍高速であり、GPU 上の DeiT より約 400 倍高速と推定されます。
• エネルギー： ReRAM IMC 上の FAR は、ReRAM IMC 上の DeiT より約 3 倍エネルギー効率が高く、GPU 上の DeiT より約 150 倍効率的と推定されます。
• メカニズム： これらの向上は、活性化同士の乗算の排除、ADC/DAC 変換の削減、および重み固定実行の実現に起因します。著者は、softmax コスト（IMC 上のアテンションにとって重大なボトルネック）を除外しているため、これらの推定値は FAR に対して保守的であると指摘しています。

4. 意義と主張

本論文は、FAR が、ゼロからモデルを訓練することなく、IMC ベースのエッジアクセラレータに事前学習されたトランスフォーマーを展開するための実用的でハードウェア効率の高い代替手段を提供すると主張しています。

• 完全置換の実現可能性： この研究は、実用的な精度を維持しながら推論時にアテンションをエンドツーエンドで置換できることを実証し、高性能なビジョンタスクにとってすべてのペアのトークン相互作用が厳密に必要であるという概念に挑戦しています。
• ハードウェアとの整合性： アテンションを逐次的で重み固定のモジュールに置換することで、FAR は計算データフローを ReRAM 交叉配列の物理的な強みと整合させ、メモリトラフィックと遅延を大幅に削減します。
• 事前学習知識の保存： このフレームワークは、大規模な事前学習中に学習された豊かな意味的表現が、ブロック単位の蒸留を通じて構造的に異なり IMC 親和的なアーキテクチャに保存され、転送可能であることを証明しています。

要約すると、FAR は、アテンションベースモデルの支配性と新興のメモリ中心ハードウェアの制約の間のギャップを埋めるために、IMC 親和的な逐次モジュールを中心にトランスフォーマー推論を再構築する道筋を提供します。