GRAM-MLX：在 Apple Silicon 上实现随机多轨迹递归推理

GRAM-MLX 将 KAIST/NYU/Mila 的 GRAM 论文移植到 Apple Silicon 的 MLX 框架，结合递归深度与随机宽度，实现多轨迹并行推理，在推理效率上实现突破。

• 原作者/维护者：DeadByDawn101
• 来源平台：github
• 原始标题：GRAM-MLX
• 原始链接：https://github.com/DeadByDawn101/GRAM-MLX
• 来源发布时间/更新时间：2026-06-06T19:38:09Z

原作者与来源

• 原作者/维护者：DeadByDawn101
• 来源平台：github
• 原始标题：GRAM-MLX
• 原始链接：https://github.com/DeadByDawn101/GRAM-MLX
• 来源发布时间/更新时间：2026-06-06T19:38:09Z 原作者与来源\n\n- 原作者/维护者: DeadByDawn101 / RavenX LLC\n- 来源平台: GitHub\n- 原始标题: GRAM-MLX: Generative Recursive Reasoning on Apple Silicon\n- 原始链接: https://github.com/DeadByDawn101/GRAM-MLX\n- 论文来源: Baek et al., KAIST/NYU/Mila, Yoshua Bengio 团队 (arXiv:2605.19376)\n- 发布时间: 2026-06-06\n\n背景：递归推理的局限\n\n大型语言模型在复杂推理任务上的能力不断提升，但传统的递归推理模型（如 OpenMythos/TRM）存在一个根本局限：它们是确定性的。给定相同的输入，模型总是沿着相同的轨迹进行推理，最终得到相同的答案。这种单一路径的推理方式限制了模型探索多种可能性的能力，容易陷入局部最优。\n\n人类在解决复杂问题时，往往会尝试多种不同的思路，比较各种路径的优劣，最终选择最佳方案。如何让 AI 模型也能具备这种"多轨迹并行探索"的能力，是提升推理质量的关键方向。\n\nGRAM 的核心创新\n\nGRAM（Generative Recursive reAsoning Models）由 KAIST、NYU 和 Mila 的研究者提出，引入了随机引导机制，使模型能够同时探索多条推理路径，并通过 Latent Process Reward Model（LPRM）选择最优结果。\n\n双轴扩展架构\n\nGRAM 的创新在于同时扩展两个维度：\n\n深度（Depth）：增加递归推理的步数。训练时使用 16 步，测试时可扩展到 64 步，实现更深层次的推理。\n\n宽度（Width）：引入随机性，并行采样多条轨迹。通过 20-100 个样本的并行探索，获得多样化的解决方案。\n\n这种"深度 × 宽度"的组合带来了惊人的效果：GRAM 使用 20 个样本、16 步推理，准确率达到了 97.0%，而确定性模型需要 320 步才能达到 90.5% 的准确率。效率提升了 20 倍，效果反而更好。\n\n技术架构解析\n\nGRAM 的核心架构包含以下关键组件：\n\n编码器与序曲（Prelude）：输入首先经过编码器处理，生成初始的隐藏状态表示。\n\n递归推理循环：在每个递归步骤中，模型执行三个操作：\n1. 低层细化（fL × K）：对当前状态进行多轮精细调整\n2. 高层提议（fH）：基于输入生成高层语义表示\n3. 随机引导：从条件高斯分布中采样噪声，注入到隐藏状态：ε ~ N(μ(ht,x), σ(ht,x))，这是 GRAM 的关键创新\n\n尾声与输出（Coda）：经过 T 轮递归后，生成最终输出。\n\n训练与推理机制\n\n训练阶段采用摊销变分推断：\n- 后验分布 q(ε|x,y) 同时观察输入和正确答案，指导训练过程\n- 先验分布 p(ε|x) 仅观察输入，用于推理阶段\n- KL 散度 KL(q||p) 作为正则项，控制探索的多样性\n\n推理阶段通过 LPRM（Latent Process Reward Model）评估所有轨迹的质量，选择得分最高的作为最终答案。\n\nMLX 移植的意义\n\nMLX 是 Apple 专为 Apple Silicon 设计的机器学习框架，充分利用了统一内存架构和神经引擎。GRAM-MLX 的移植意义重大：\n\n本地高效推理：Apple Silicon 用户可以在本地运行复杂的递归推理模型，无需依赖云端 API。\n\n统一内存优势：MLX 的统一内存架构允许模型高效处理多轨迹并行的内存需求，避免传统 GPU 的显存瓶颈。\n\n与 OpenMythos 生态整合：GRAM-MLX 扩展了 OpenMythos 的 Reasoning-as-a-Service 流水线，从"单轨迹蒸馏"升级为"多轨迹选择后蒸馏"，有望训练出更强大的生产级模型。\n\n模型配置与使用\n\nGRAM-MLX 提供了三种规模的预配置：\n\n- gram_small：约 2000 万参数，256 维，4 头注意力\n- gram_base：约 5000 万参数，512 维，8 头注意力\n- gram_large：约 2 亿参数，768 维，12 头注意力\n\n使用示例简洁直观：\n\npython\nimport mlx.core as mx\nfrom gram_mlx import GRAM, gram_small\n\nmodel = gram_small()\nids = mx.random.randint(0, 32000, (1, 32))\n\n训练模式：单轨迹\nlogits, kl_loss = model(ids, n_steps=16)\n\n推理模式：多轨迹采样\nbest_logits, scores = model.sample_trajectories(\n ids, n_steps=16, n_samples=20\n)\n\n\n模型无关的通用包装器\n\nGRAM-MLX 还提供了一个模型无关的包装器，可以为任何现有模型添加随机多轨迹推理能力。这意味着开发者不需要从头训练 GRAM 架构，而是可以将这一推理增强技术应用到自己的模型上。\n\n实际应用价值\n\nGRAM-MLX 的技术方案在以下场景具有显著价值：\n\n数学与逻辑推理：复杂数学问题往往需要尝试多种解法，GRAM 的并行探索能力可以显著提高解题成功率。\n\n代码生成与调试：不同的代码实现路径可能各有优劣，多轨迹探索有助于找到更优雅、更高效的解决方案。\n\n科学假设生成：在科学研究中，同时考虑多种假设并评估其合理性是常态，GRAM 的架构天然适合这类任务。\n\n决策支持系统：面对复杂的决策问题，并行评估多个方案的风险和收益，能够提供更全面的决策依据。\n\n总结与展望\n\nGRAM-MLX 代表了递归推理领域的重要进展。通过引入随机宽度维度，它突破了确定性递归模型的局限，在保持高效推理的同时显著提升了准确率。Apple Silicon 的本地部署能力让这一技术更加普及，为开发者和研究者提供了强大的推理工具。\n\n随着多轨迹推理技术的成熟，我们可以期待未来出现更多能够"深思熟虑"的 AI 系统，它们不再满足于第一个想到的答案，而是像人类一样，耐心探索、比较权衡，最终给出更可靠、更智慧的回应。