lucidrains/x-transformers

何を解決するか

研究論文から得た多様な実験的アーキテクチャ変種と最適化を、簡単に統合できる簡潔で機能豊富なTransformerアーキテクチャの実装を提供します。研究者や開発者が、論文から得た実験的機能を自前で実装せずに、モデルに組み込むことが可能になります。

どう動くか

XTransformer、TransformerWrapper、Encoder、Decoder などの柔軟なラッパーとモジュールを提供しており、シンプルなキーワード引数で構成可能です。完全なエンコーダデコーダ、デコーダ専用（GPT風）、エンコーダ専用（BERT風）、およびビジョンTransformer（ViT）の設定をサポートしています。メモリ効率と高速化を実現するFlash Attentionなどの高度な最適化を統合し、画像キャプションや言語・ビジョンモデル（例：PaLI）などのマルチモーダルタスクもサポートしています。

対象ユーザー

異なるTransformer変種、正規化手法、アテンション機構を実験して、モデルのパフォーマンスや効率を向上させたいAI研究者や機械学習エンジニア。

特徴

• 多様なアーキテクチャ対応: エンコーダ専用、デコーダ専用、完全なエンコーダデコーダ構成をサポート。
• 実験的アテンション機能: 永続的メモリKV、メモリトークン（レジスタトークン）、トークヘッドアテンション、スパースtop-kアテンションを含む。
• 高度な正規化: RMSNorm、SimpleRMSNorm、ScaleNorm、L2正規化された埋め込みを実装。
• 効率性最適化: Flash Attentionおよびマルチクエリ/グループクエリアテンション（1つのKVヘッド）の統合サポート。
• フィードフォワードの強化: GLU変種（GELU、Swish）、ReLU²、バイアスなしフィードフォワード層をサポート。
• ビジョン統合: 画像分類およびマルチモーダルタスク用に ViTransformerWrapper を含む。