Swarmformer Small Ef

🚀 スワームフォーマーの強化: 効率的なシーケンスモデリングの賢いアプローチ

私たちは、効率的なシーケンスモデリングアーキテクチャの改良に取り組んでおり、SwarmFormerにいくつかの更新を加え、その性能、拡張性、および安定性を向上させました。新しい設計では、階層的アテンション、動的クラスタリング、およびゲート付きフィードバックメカニズムを導入しています。これにより、モデルは長いシーケンスをより効果的に処理し、計算オーバーヘッドを削減することができます。

🚀 クイックスタート

この新しいSwarmFormerアーキテクチャは、元のモデルの制限を克服し、性能と計算効率を向上させます。以下では、改良の理由と主要な強化点について説明します。

✨ 主な機能

なぜSwarmFormerを改良したのか？

元のSwarmFormerはトークン-クラスタ相互作用モデルを導入しました。このモデルでは、トークンが自組織的にクラスタを形成し、高レベルで情報を交換し、その後、改良された表現を伝播させます。このアプローチは長距離依存関係を効率的に処理しましたが、いくつかの制限がありました。 ❌ 固定クラスタ割り当てにより、トークンのグループ化が硬直的になりました。 ❌ ローリングシフトによるローカルアテンションは、細粒度の依存関係を捉えるのに最適ではありませんでした。 ❌ クラスタからトークンへの更新にゲートがなかったため、ノイズの多い更新が発生しました。 ❌ アテンションレイヤで重み共有がなかったため、パラメータ数が増加しました。 これらの問題を解決するために、モデルの表現力を向上させつつ、計算効率を維持する一連の重要な強化点を導入しました。

新しいSwarmFormerアーキテクチャの主要な強化点

1. ローリングシフトの代わりにローカルウィンドウアテンション

ローリングシフトアテンションの代わりにローカルウィンドウアテンション（Sliding TransformersやConv Filtersに似たもの）を採用しました。これにより、冗長なシフトなしでより効率的なローカル特徴抽出が可能になり、局所性モデリングが向上します。

2. クラスタに対するマルチヘッドアテンション

クラスタに対して単一のアテンションメカニズムを使用する代わりに、マルチヘッド自己アテンション（MHA） を適用しました。これにより、各アテンションヘッドが異なるクラスタ-トークン関係を学習できるようになり、文脈表現が向上します。

3. 均一なチャンキングの代わりにトークンからクラスタへのゲート付き機構

以前は、トークンが均一にクラスタに割り当てられていましたが、これは柔軟性を制限していました。現在では、アテンションベースの動的ルーティングメカニズムを使用しており、トークンが適応的にクラスタを選択できるようになっています。これにより、クラスタ形成における意味的な一貫性が向上します。

4. 安定したトークン更新のためのゲート付きフィードバックメカニズム

クラスタから直接トークン埋め込みを更新する代わりに、残差MLPゲートメカニズムを導入しました。これにより、ノイズの多いクラスタ更新がフィルタリングされ、トークンに関連する情報のみが伝播されます。

5. すべてのMLPとアテンションブロックの前にレイヤー正規化

すべてのフィードフォワードレイヤーとアテンションレイヤーの前にLayerNormを追加することで、学習が大幅に安定し、勾配の流れと収束が改善されることがわかりました。

6. クラスタアテンションにおける線形射影の重み共有

モデルサイズを削減しつつ表現力を維持するために、GlobalClusterAttentionモジュールのクエリ、キー、およびバリュー射影で重みを共有するようにしました。この最適化により、学習可能なパラメータ数が減少し、性能が維持されます。

7. ピラミッド構造による階層的クラスタリング

すべてのレイヤーで固定クラスタサイズを使用する代わりに、階層的ピラミッドを実装しました。 ✅ 下層は細粒度の局所相互作用に焦点を当てます（より多くのクラスタ）。 ✅ 上層は抽象的な粗粒度の表現を処理します（より少ないクラスタ）。 このマルチスケールクラスタ形成により、モデルは高レベルの情報を効率的に伝播できるようになり、局所的な詳細を失うことがなくなります。

8. 微分可能なクラスタリングのためのガンベル-ソフトマックス

クラスタ割り当ての学習可能性を向上させるために、ガンベル-ソフトマックスサンプリングを実装しました。これにより、モデルは誤差逆伝播によってクラスタ割り当てを学習できるようになり、強化信号（クラスタの一貫性など）が最適化を導くことができます。

🔧 技術詳細

計算量の比較

新しいSwarmFormerアーキテクチャは、元の実装よりも計算コストが低くなっています。以下は、この主張を裏付ける数学的な比較です。

元のSwarmFormerの計算量

トークンからクラスタへのアテンション

元のSwarmFormerでは、各トークンがすべてのクラスタにアテンションを向けるため、計算量は次のようになります。 $$ O(NCd)$$ ここで、

• N = シーケンス長
• C = クラスタ数
• d = 隠れ次元

クラスタからトークンへのブロードキャスト

各クラスタがすべてのトークンを更新するため、さらに次の計算量が発生します。 $$O(NCd)$$

総計算量（元のSwarmFormer）

$$O(NCd)+O(NCd)=O(2NCd)$$

新しいSwarmFormerの計算量

ローリングシフトアテンションの代わりにローカルウィンドウアテンション

すべてのトークンに対するグローバルアテンションの代わりに、各トークンはサイズwのローカルウィンドウにのみアテンションを向けます（通常、w≪N）。 $$O(Nwd)$$ これにより、ローリングシフト操作が置き換えられ、計算コストが大幅に削減されます。

重み共有によるマルチヘッドクラスタアテンション

元のバージョンでは、クエリ、キー、およびバリュー射影に別々の重みがありました。現在では、これらの射影で重みを共有することで、クラスタアテンションレイヤーのパラメータ数とFLOP数が削減されます。アテンションの計算量は次のようになります。 $$O(NCd)$$ ただし、行列乗算の回数が減少します。

トークンからクラスタへのゲート付き機構

すべてのトークンがすべてのクラスタを更新する代わりに、トークンは学習されたルーティングに基づいてクラスタを選択的に更新します。これにより、すべてのトークンからすべてのクラスタへの更新数が、参加するトークンの一部pに減少します。 $$O(pNCd)$$ ここで、p<1です。通常、pは0.5以下であるため、計算量が大幅に削減されます。

ゲート付きフィードバックメカニズム（MLPフィルタリング）

クラスタからトークンへの更新を完全に伝播する代わりに、更新をブロードキャストする前にゲート付きの残差MLPを適用します。MLPの計算量は次のようになります。 $$O(Nd^2)$$ ただし、不必要な更新が防止され、後続のレイヤーでの実効的な計算量が削減されます。

ピラミッド構造による階層的クラスタリング

すべてのレイヤーで固定数のクラスタを使用する代わりに、深くなるにつれてクラスタ数を徐々に減少させます。

• 下層: C個のクラスタ
• 中層: C/2個のクラスタ
• 上層: C/4個のクラスタ これにより、クラスタリングの計算量が効果的に削減されます。 $$O(NCd+NC/2d+NC/4d+…)$$ これは幾何級数を形成し、総計算コストが削減されます。

最終的な計算量比較

• w≪N（ウィンドウアテンションによるコスト削減）
• p<1（より少ないクラスタ更新）
• $$d^2$$項は小さなMLPにのみ存在し、完全なアテンションレイヤーではない
• 階層的クラスタリングにより、総クラスタ相互作用が減少する 次の不等式が成り立ちます。 $$O(2NCd) > O(Nwd + pNCd + Nd^2)$$ これは、新しいアーキテクチャが計算コストが低いことを示しています。

結論: 新しいSwarmFormerはより効率的です

• ✅ ウィンドウアテンションと階層的クラスタリングにより、FLOP数が減少します。
• ✅ ゲート付きフィードバックとトークンからクラスタへのゲート付き機構により、冗長な更新が減少します。
• ✅ 重み共有により、パラメータ数がさらに削減されます。 要約すると、🚀 新しいSwarmFormerアーキテクチャは、性能を維持または向上させながら、より高速な学習と推論を実現します！

階層的アテンションと適応的クラスタリングについてどう思いますか？コメントで議論しましょう！ 🎯

参考文献

@article{legg2025swarmformer,
  title={SwarmFormer: Local-Global Hierarchical Attention via Swarming Token Representations},
  author={Legg, Jordan and Sturmanis, Mikus and {Takara.ai}},
  journal={Takara.ai Research},
  year={2025},
  url={https://takara.ai/papers/SwarmFormer-Local-Global-Hierarchical-Attention-via-Swarming-Token-Representations.pdf}
}

📄 ライセンス

このプロジェクトはApache-2.0ライセンスの下でライセンスされています。