lucidrains/x-transformers

什么是 x-transformers？

x-transformers 是一个轻量级、纯 Python 的库，让您可以在 PyTorch 中构建并实验多种 Transformer 模型。它提供了一个单一且灵活的构建模块（XTransformer 或 TransformerWrapper），可以配置为：

• 编码器-解码器（类似原始的 “Attention is All You Need” 模型）
• 仅解码器（GPT 风格的语言模型）
• 仅编码器（BERT 风格的模型）
• 视觉 Transformer，用于图像分类或图像转文字任务

该库还捆绑了一系列研究级的调整功能——flash-attention、各种归一化方案、内存标记、稀疏/top-k 注意力、talking-heads、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	在 softmax 前后对头部进行线性混合。	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	嵌入 dropout、随机深度、注意力 dropout、前馈 dropout。	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 架构（语言、视觉、多模态），并尝试最近论文中广泛的注意力与归一化技巧。