[Paper] MAR：通过模块感知架构细化实现高效大语言模型

Source: arXiv - 2601.21503v1

概述

大型语言模型（LLMs）已成为众多 AI 产品的核心，但它们的二次方时间注意力和密集前馈网络（FFNs）使得推理在计算和能耗上都非常昂贵。论文 “MAR: Efficient Large Language Models via Module‑aware Architecture Refinement” 提出了一种两阶段框架，用更高效的替代方案——用于序列处理的状态空间模型（SSMs）和稀疏激活的前馈网络——来替换 LLM 中昂贵的部分，同时保持（甚至提升）性能。

关键贡献

• Module‑aware Architecture Refinement (MAR)：一种系统化流水线，用线性时间 SSM 替代二次注意力，并在不对每层手动调优的情况下稀疏化 FFN 激活。
• Adaptive Ternary Multi‑step Neuron (ATMN)：一种新颖的脉冲神经元设计，弥合 SSM 与脉冲神经网络 (SNN) 之间的时间不匹配，使低信息密度信号能够高效处理。
• Spike‑aware Bidirectional Distillation Strategy (SBDS)：一种训练配方，联合从密集教师模型向 SSM 和 SNN 模块蒸馏知识，确保精炼后的架构恢复原始精度。
• 全面的能耗感知评估：作者在硬件上测量实际推理能耗，显示相比密集基线降低最高 45 %，同时匹配或超越其 BLEU/GLUE 分数。
• 可扩展至大模型规模：即使在参数量相当或更大的情况下，MAR 仍优于其他“高效”LLM 变体（如 LoRA‑剪枝、量化或稀疏化模型）。

Source: …

方法论

两阶段细化

1. 阶段 1 – 注意力替换：将每个 Transformer 自注意力块替换为线性时间 SSM（例如基于 HiPPO 的）。SSM 以 O(N) 时间处理 token 序列，消除 N² 的注意力矩阵。
2. 阶段 2 – FFN 稀疏化：使用基于激活的剪枝将密集 FFN 转换为稀疏版本。每个 token 只保留前 k 个激活，将密集矩阵乘法转变为成本更低的稀疏运算。

脉冲感知集成

• ATMN 将连续的 SSM 输出转换为三值脉冲（‑1、0、+1），跨多个时间步进行，既保留信息，又使下游 SNN 以极低能耗运行。
• SBDS 执行双向知识蒸馏：密集教师引导 SSM 的隐藏状态，同时 SNN 的脉冲动力学被正则化以模仿教师的 token 级 logits。此联合训练能够恢复因结构替换导致的性能损失。

训练流水线

• 预训练一个标准的密集 LLM（教师）。
• 冻结教师模型，使用 SBDS 训练 MAR‑细化的学生模型，交替进行仅 SSM 前向和 SSM+SNN 前向传播。
• 使用基于幅值的掩码对稀疏化的 FFN 进行微调，该掩码每隔几千步更新一次。

整个过程通过“模块感知”调度器自动化，该调度器根据层深度和 token 信息密度决定在何处使用 SSM 与传统注意力。

结果与发现

• 性能相当：MAR 在语言理解基准（GLUE、SuperGLUE）和生成任务（BLEU、ROUGE）上恢复了稠密模型 >99 % 的准确率。
• 能耗节省：在 NVIDIA A100 和低功耗 ARM CPU 上的实际测量显示，每标记推理能耗最高可降低 45 %。
• 可扩展性：当模型扩展到 13 B 参数时，MAR 仍然优于更大的稠密基线，表明效率提升随模型规模而叠加。

实际意义

• Edge‑AI 与设备端 LLM：线性时间 SSM 和脉冲组件使得在电池受限的设备（如智能手机、可穿戴设备）上运行复杂语言模型成为可能，且不牺牲质量。
• 云成本降低：对于推理密集型服务（聊天机器人、代码助手），MAR 可以削减 GPU 小时费用和碳足迹，直接转化为更低的运营成本。
• 简化部署流水线：由于 MAR 可作为注意力/前馈网络模块的直接替代，现有 Transformer 代码库只需最小的工程工作——替换模块并运行提供的 SBDS 训练脚本，即可完成改造。
• 兼容其他效率技巧：MAR 可与量化、模型并行或 LoRA 微调相结合，提供分层的优化方案。

限制与未来工作

• 硬件对 SSM 和脉冲的支持：虽然作者在 GPU/CPU 上测量了能耗，但最大的收益出现在能够原生处理三值脉冲的专用神经形态或 ASIC 加速器上；更广泛的硬件支持仍处于初期阶段。
• 训练开销：相较于普通密集模型，两阶段精炼加上双向蒸馏会额外增加约 30 % 的训练时间。
• 时间对齐敏感性：ATMN 设计假设了特定的序列长度；对于非常长的文档（>4 k 令牌），仍可能出现残余的时间错配。

未来方向 建议包括：

1. 为新兴 AI 芯片共同设计 SSM 核心。
2. 能够根据运行时延迟约束自适应稀疏度调度。
3. 将 MAR 扩展到多模态 Transformer（视觉‑语言、语音）。

结论：MAR 提供了一条务实的路径，使大型语言模型在不牺牲开发者依赖的面向用户的质量的前提下，更加绿色和更快。通过从线性时间动力学和脉冲稀疏性的视角重新思考核心注意力和前馈模块，该框架为真正可扩展的 LLM 部署打开了大门——从云集群到边缘设备。

作者

• Junhong Cai
• Guiqin Wang
• Kejie Zhao
• Jianxiong Tang
• Xiang Wang
• Luziwei Leng
• Ran Cheng
• Yuxin Ma
• Qinghai Guo

论文信息

• arXiv ID: 2601.21503v1
• 分类: cs.AI, cs.CL, cs.LG, cs.NE
• 出版日期: 2026年1月29日
• PDF: 下载 PDF