lucidrains/x-transformers

x-transformers とは？

x-transformers は、PyTorchで多種多様なTransformerモデルを構築・実験できる、軽量で純粋なPythonライブラリです。単一の柔軟なビルディングブロック（XTransformer または TransformerWrapper）を提供し、以下のように設定可能です：

• エンコーダー・デコーダー（オリジナルの “Attention is All You Need” モデルのようなもの）
• デコーダーのみ（GPTスタイルの言語モデル）
• エンコーダーのみ（BERTスタイルのモデル）
• ビジョンTransformer（画像分類や画像からテキストへの変換タスク用）

このライブラリには、フラッシュアテンション、様々な正規化スキーム、メモリトークン、スパース/top-kアテンション、トーキングヘッド、GLUフィードフォワードなど、研究グレードの調整機能がバンドルされており、単一のフラグで有効にできます。

---

どのような人が使うべきか？

• 研究者：ボイラープレートコードを書かずに、新しいTransformerバリアントのプロトタイプを迅速に作成したい方。
• 実務者：小規模から中規模の実験（カスタム言語モデル、画像キャプション、マルチモーダルパイプラインなど）のために、コンパクトで設定可能なTransformerを探している方。
• 学生：ほとんどのオプションが単純な引数として公開されているため、アーキテクチャの選択がパフォーマンスにどのように影響するかを学びたい方。

---

主な機能（READMEより）

機能	内容	有効化方法
Flash Attention	PyTorch 2.0の scaled_dot_product_attention を使用し、メモリ効率と高速化を実現。	attn_flash=True
Persistent Memory KV	通常のアテンションの前に学習可能なキー/値スロットを追加。	attn_num_mem_kv=16
Memory Tokens	各層を通過する追加トークンを学習（安定性に寄与）。	num_memory_tokens=20
ScaleNorm / RMSNorm / Simple RMSNorm	収束を加速させる代替正規化層。	use_scalenorm=True, use_rmsnorm=True, use_simple_rmsnorm=True
GLU / Swish-GLU / Feed-forward without bias	精度やスループットを向上させることが示されているフィードフォワードブロックのバリアント。	ff_glu=True, ff_swish=True, ff_no_bias=True
ReLU² activation	GELUを上回ることもある単純な活性化関数。	ff_relu_squared=True
Sparse / Top-k attention	top-kのアテンションスコアのみを保持し、計算量を削減。	attn_sparse_topk=8 (または attn_hard=True でargmax)
Talking-heads	ソフトマックスの前後でヘッドを線形混合。	attn_pre_talking_heads=True, attn_post_talking_heads=True
One-KV-head / grouped KV heads	メモリ効率の良いデコードのためにキー/値ヘッドの数を削減。	attn_one_kv_head=True または attn_kv_heads=2
Dropout options	埋め込みドロップアウト、確率的深度、アテンションドロップアウト、フィードフォワードドロップアウト。	emb_dropout, layer_dropout, attn_dropout, ff_dropout

---

使用例（READMEより）

# エンコーダー・デコーダー (seq2seq) モデル
from x_transformers import XTransformer
model = XTransformer(
    dim=512, enc_num_tokens=256, enc_depth=6, enc_heads=8,
    dec_num_tokens=256, dec_depth=6, dec_heads=8,
    tie_token_emb=True)

# GPTスタイルのデコーダーのみ
from x_transformers import TransformerWrapper, Decoder
model = TransformerWrapper(
    num_tokens=20000, max_seq_len=1024,
    attn_layers=Decoder(dim=512, depth=12, heads=8))

# 画像分類用ビジョンTransformer
from x_transformers import ViTransformerWrapper, Encoder
model = ViTransformerWrapper(
    image_size=256, patch_size=32, num_classes=1000,
    attn_layers=Encoder(dim=512, depth=6, heads=8))

---

インストール

pip install x-transformers

---

まとめ

x-transformers は、モジュール式で研究に特化したPyTorchライブラリです。少数の引数で古典的および最先端のTransformerアーキテクチャ（言語、ビジョン、マルチモーダル）を構築し、最近の論文で提案されている幅広いアテンションや正規化のテクニックを試すことができます。