📌 项目地址:kyegomez/OpenMythos | ⭐ 13,441 颗星 | 🔧 Python | 📜 MIT
它是什么,又不是什么
OpenMythos 是一个针对传闻中 Claude Mythos 架构的开源理论性重建。它的核心目标是探索,而非复现。项目代码和说明基于公开的学术研究,与任何公司或闭源模型无关联。
这个项目的主要价值在于提供了一个可运行、可修改的平台,用于研究一种特定的、可能更高效的 Transformer 架构变体:循环深度 Transformer。
核心架构:循环替代深度
标准 Transformer 模型的能力由层数决定,层数是固定的。OpenMythos 实现的 RDT 架构通过一个可重复执行的模块来增加计算深度。
它的结构分为三部分:
- Prelude:前置的普通 Transformer 层,负责初始信息处理。
- Recurrent Block:核心循环模块。这个模块不是堆叠很多层,而是同一个模块循环执行多次(由参数
max_loop_iters控制)。每次循环都复用相同的权重,但处理不同的信息状态。 - Coda:后置的普通 Transformer 层,负责整合并输出最终结果。
这种设计的理论好处是,模型在处理不同复杂度的任务时,可以动态调整实际使用的计算量,而不是像传统模型那样对每个 token 都执行完全相同数量的计算。
可配置的模块:注意力与前馈
项目在实现细节上提供了高度的可配置性,方便研究者进行实验。
注意力机制有两种选择:
- GQA (Grouped-Query Attention):一种高效的注意力实现。如果环境中安装了
flash-attn >= 2.5.8,它会自动使用 Flash Attention 2 加速;否则回退到手动实现。 - MLA (Multi-Latent Attention):受 DeepSeek-V2 启发。它不缓存完整的 Key 和 Value 矩阵,而是缓存一个压缩后的潜在向量,以此减少内存占用。旋转位置编码在应用后无需重新计算。
前馈网络采用稀疏专家混合:
包含一组由输入 token 动态路由选择的专家,和一组始终被激活的共享专家。这种结构理论上可以更好地分配计算资源。
如何上手尝试
安装非常简单。
基础安装:
pip install open-mythos
如需启用 Flash Attention 2 以加速 GQA(需要 CUDA 和构建工具):
pip install open-mythos[flash]
以下是来自官方 README 的完整使用示例。它清晰地展示了如何配置不同注意力类型、控制循环次数,并检查模型状态。
import torch
from open_mythos.main import OpenMythos, MythosConfig
attn_type = "mla" # 或 "gqa"
base = {
"vocab_size": 1000,
"dim": 256,
"n_heads": 8,
"max_seq_len": 128,
"max_loop_iters": 4,
"prelude_layers": 1,
"coda_layers": 1,
"n_experts": 8,
"n_shared_experts": 1,
"n_experts_per_tok": 2,
"expert_dim": 64,
"lora_rank": 8,
"attn_type": attn_type,
}
if attn_type == "gqa":
cfg = MythosConfig(**base, n_kv_heads=2)
else:
cfg = MythosConfig(
**base,
n_kv_heads=8,
kv_lora_rank=32,
q_lora_rank=64,
qk_rope_head_dim=16,
qk_nope_head_dim=16,
v_head_dim=16,
)
model = OpenMythos(cfg)
total = sum(p.numel() for p in model.parameters())
print(f"n[{attn_type.upper()}] Parameters: {total:,}")
ids = torch.randint(0, cfg.vocab_size, (2, 16))
logits = model(ids, n_loops=4)
print(f"[{attn_type.upper()}] Logits shape: {logits.shape}")
out = model.generate(ids, max_new_tokens=8, n_loops=8)
print(f"[{attn_type.upper()}] Generated shape: {out.shape}")
A = model.recurrent.injection.get_A()
rho = torch.linalg.eigvals(A).abs().max().item()
print(
f"[{attn_type.upper()}] Spectral radius ρ(A) = {rho:.4f} (must be ~1.0 for stable dynamics)"
)
关键参数是 n_loops。在代码中,无论是调用 model(ids, n_loops=4) 还是 model.generate(ids, n_loops=8),它都直接决定了 Recurrent Block 的实际执行次数。
为什么这个项目值得关心
它提供了一条探索路径。标准 Transformer 的扩展方式通常是堆更多层,而 RDT 范式探索的是“更少参数,更多循环”。这可能导向更高效、或更具适应性的模型设计。
对于研究者和工程师,OpenMythos 是一个现成的实验台。你可以调整循环次数、切换注意力类型、修改专家数量,然后直接观察模型输出或内部动态(如检查注入矩阵的谱半径是否稳定)。这些操作在闭源模型上无法进行。
它也在尝试回答一个理论问题:通过动态调整计算深度,模型能否在组合泛化等任务上获得比固定架构更好的表现?项目的 n_loops 参数设计正是为了支持这类研究。