📌 项目地址kyegomez/OpenMythos | ⭐ 13,441 颗星 | 🔧 Python | 📜 MIT

它是什么,又不是什么

OpenMythos 是一个针对传闻中 Claude Mythos 架构的开源理论性重建。它的核心目标是探索,而非复现。项目代码和说明基于公开的学术研究,与任何公司或闭源模型无关联。

这个项目的主要价值在于提供了一个可运行、可修改的平台,用于研究一种特定的、可能更高效的 Transformer 架构变体:循环深度 Transformer。

核心架构:循环替代深度

标准 Transformer 模型的能力由层数决定,层数是固定的。OpenMythos 实现的 RDT 架构通过一个可重复执行的模块来增加计算深度。

它的结构分为三部分:

  • Prelude:前置的普通 Transformer 层,负责初始信息处理。
  • Recurrent Block:核心循环模块。这个模块不是堆叠很多层,而是同一个模块循环执行多次(由参数 max_loop_iters 控制)。每次循环都复用相同的权重,但处理不同的信息状态。
  • Coda:后置的普通 Transformer 层,负责整合并输出最终结果。

这种设计的理论好处是,模型在处理不同复杂度的任务时,可以动态调整实际使用的计算量,而不是像传统模型那样对每个 token 都执行完全相同数量的计算。

可配置的模块:注意力与前馈

项目在实现细节上提供了高度的可配置性,方便研究者进行实验。

注意力机制有两种选择

  1. GQA (Grouped-Query Attention):一种高效的注意力实现。如果环境中安装了 flash-attn >= 2.5.8,它会自动使用 Flash Attention 2 加速;否则回退到手动实现。
  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 参数设计正是为了支持这类研究。

这篇文章对你有帮助吗?

发表回复