中科院联手美团推出MoTE:如何优化大模型内存?

发布时间:2025-06-24 23:35  浏览量:2

这项由中科院计算技术研究所与美团联合发布的研究成果于2025年6月发表在机器学习顶级期刊上(论文编号:arXiv:2506.14435v1),有兴趣深入了解的读者可以通过该编号在arXiv平台访问完整论文。

当我们谈论人工智能的时候,经常会听到"大模型"这个词。可以把大模型想象成一个超级聪明的万事通,它能看图识物、能与人对话,甚至能帮你写作业。但这样的万事通有个致命缺点——它们就像超级能吃的大胃王,需要消耗巨大的内存才能运行。

以目前主流的多模态大模型为例,它们需要的内存动辄几十GB,这对于普通电脑或手机来说简直是天文数字。就好比你想在家里养一头大象,但你的房间只有卧室那么大——根本装不下。正因为这个问题,很多强大的AI功能只能在云端运行,普通用户无法在自己的设备上直接使用。

中科院计算技术研究所的王红宇、徐佳宇等研究人员联合美团的技术团队,提出了一个巧妙的解决方案——MoTE(Mixture of Ternary Experts,三值专家混合模型)。这个名字听起来很复杂,但其实原理很简单:与其让一个超级专家处理所有问题,不如培养很多个"小而精"的专家,每个专家只负责自己擅长的领域,而且这些专家还特别"节食",内存消耗极低。

一、从"大胃王"到"节食专家"的华丽转身

要理解MoTE的巧妙之处,我们先来看看传统AI模型面临的困境。目前的大模型就像一个什么都懂的全科医生,无论你是眼科问题、心脏病还是骨折,都得找这一个医生。这个医生确实很厉害,但问题是他的"大脑"(内存)需要装下所有医学知识,导致他需要一个超级大的办公室(服务器)才能工作。

研究团队发现,当他们尝试用现有的专家混合(MoE)技术来解决这个问题时——也就是把一个全科医生替换成多个专科医生——虽然计算效率提高了,但内存消耗反而更大了。这就像你原本雇一个全科医生,现在要雇16个专科医生,每个医生都需要自己的办公室,总的办公空间需求反而增加了好几倍。

具体来说,当研究人员将一个3B参数的模型扩展为包含16个专家的MoE模型时,内存需求从5.2GB猛增到73.2GB——增长了14倍之多。这对于消费级设备来说完全不现实。

MoTE的解决思路相当巧妙。研究团队没有选择训练更少但"营养丰富"的专家,而是反其道而行之——训练更多但"节食"的专家。这些专家使用三值量化技术,即每个参数只能取-1、0、1三个值中的一个。听起来很限制?实际上这就像用黑白灰三种颜色画画,虽然颜色选择少了,但熟练的画家依然能创作出精美的作品。

二、巧妙的"新老搭配"策略

MoTE最聪明的地方在于其"新老搭配"的策略。研究团队发现,如果直接把原来训练好的全精度网络层替换成三值专家,性能会大幅下降。这就像你把一个经验丰富的老医生直接换成刚毕业的实习生,虽然实习生很节约成本,但医疗水平肯定不如老医生。

为了解决这个问题,研究团队设计了一个巧妙的架构:保留原来的全精度前馈网络作为"共享专家",同时添加多个三值"路由专家"。这就像医院里既保留了一个经验丰富的全科医生作为总顾问,同时又招聘了多个年轻但高效的专科医生。每个病人来看病时,总顾问都会参与诊断,同时根据病情特点,还会有一个最合适的专科医生协助治疗。

在技术实现上,MoTE使用top-1路由机制,即每次只激活一个三值专家,再加上始终激活的共享专家。这样设计既保证了性能,又控制了计算开销。整个架构的计算过程可以用一个简单的数学公式表示:输出 = 共享专家的结果 + 最合适的三值专家的结果。

三、三阶段训练让模型"循序渐进"

MoTE的训练过程就像培养一个优秀的医疗团队,需要三个阶段的精心安排。

第一阶段类似于"基础培训"。研究团队首先训练一个两层的多层感知机连接器,让视觉编码器和语言模型能够"对话"。这就像让放射科医生学会如何向内科医生描述X光片的内容——两个专业领域需要建立共同的交流语言。

第二阶段是"专业提升"。研究团队使用更复杂的视觉-语言指令数据对模型进行微调,提升语言模型和连接器的能力。这个阶段就像让医生团队学习如何处理更复杂的综合性病例,不仅要看懂检查结果,还要能给出准确的诊断建议。

第三阶段是关键的"专家扩展"期。在这个阶段,研究团队将第二阶段的密集模型扩展为MoTE架构。有趣的是,只有新添加的三值专家需要训练,而继承自密集模型的层(共享专家和注意力层)都保持冻结状态。这就像在一个成熟的医疗团队基础上,只培训新来的专科医生,而经验丰富的老医生和护士团队保持现有的工作模式。

训练过程中,研究团队采用了量化感知训练技术。权重和激活值会被实时量化为三值和8位整数,但梯度和优化器状态仍保持全精度。这种做法确保了训练的稳定性,就像在培训过程中严格控制新医生的诊疗范围,但保持详细的培训记录用于持续改进。

四、内存消耗大幅缩减的秘密

MoTE在内存优化方面的表现令人印象深刻。以3B参数模型为例,传统的MoE-LLaVA需要18.1GB的专家内存,而MoTE只需要6.8GB——节省了超过60%的内存。这种巨大的节省来自于三值量化技术的威力。

三值量化的工作原理相当精妙。每个权重参数都通过一个特殊的量化函数处理,该函数会计算权重矩阵的平均绝对值作为缩放因子,然后将每个权重值除以这个缩放因子并四舍五入到最接近的三个值:-1、0、1。虽然从理论上说,三值需要1.58位来表示(log?3),但在实际的GPU系统中,研究团队使用INT2格式来存储和处理这些权重,进一步提高了效率。

更令人惊喜的是,MoTE还能与后训练量化技术完美兼容。研究团队发现,由于MoE层已经是三值的,只需要对共享专家进行额外的量化处理即可进一步压缩模型。在相同的3.4GB专家内存预算下,MoTE比MoE-LLaVA的平均准确率高出4.3%,这意味着在严格的内存限制下,MoTE能够以更少的资源实现更好的性能。

五、性能表现超出预期

研究团队在多个标准数据集上对MoTE进行了全面测试,涵盖了数学推理、图像理解、文档分析等多个任务。测试结果显示,MoTE展现出了令人瞩目的规模效应。

在0.5B参数规模上,MoTE的表现略逊于全精度的MoE-LLaVA,平均准确率相差0.6%。但随着模型规模增大,这个差距迅速缩小。当模型达到1.5B参数时,MoTE开始反超MoE-LLaVA,平均准确率高出1.7%。在3B参数规模上,MoTE的优势进一步扩大,在数学推理任务MathVista上的表现比MoE-LLaVA高出3.7%,在图表问答任务ChartQA上高出2.6%。

这种规模效应的出现并非偶然。研究团队分析认为,随着模型规模增大,三值量化对性能的负面影响会被更大的模型容量所补偿,而内存效率的优势则变得更加明显。这就像一个大型医院,虽然每个专科医生的个人能力可能不如全科专家,但通过合理的分工协作,整体效率反而更高。

特别值得一提的是,在与其他同规模模型的比较中,MoTE表现出了显著的竞争优势。例如,在与Qwen2-VL-2B的对比中,尽管MoTE只用了21.6B的训练tokens(远少于Qwen2-VL的1.4T tokens),但在平均性能上仍有2.0%的提升。在数学推理任务上,MoTE比MiniCPM-V-2.0-3B高出11.1%,比Phi-3-Vision-4B高出5.3%。

六、专家路由的智能分工

研究团队对MoTE的内部工作机制进行了深入分析,发现了一个有趣的现象:不同的专家会自发地形成模态特化的分工模式。通过可视化分析,研究人员发现文本tokens和图像tokens在路由时展现出了明显不同的模式。

例如,专家1在处理图像tokens时,在第一层和最后五层都表现出很高的激活率,而在处理文本tokens时则相对较少被选择。这种现象类似于医院里的专科医生自然分工——眼科医生主要处理视觉相关问题,心内科医生主要处理心血管问题,虽然没有明确规定,但通过长期实践形成了自然的专业化分工。

更有趣的是,这种专业化分工在不同任务间保持了相当的一致性。无论是AI2D的图形理解任务、MMBench的综合评测,还是ChartQA的图表分析,专家们的路由模式都展现出了相似的特点。这表明MoTE学到的不是任务特定的特征,而是更加通用的模态处理能力。

研究团队还使用主成分分析(PCA)技术提取了top-10的路由路径,发现文本和图像tokens确实遵循着不同的"高速公路"。这种路径分化进一步证实了专家们根据输入内容的不同特点进行了自适应的专业化分工。

七、训练策略的精妙设计

在训练策略的设计上,研究团队做了多项巧妙的优化。他们发现,直接从全精度训练开始再逐步过渡到三值训练并不能带来性能提升,反而会增加训练复杂度。因此,MoTE采用了端到端的三值训练策略,从一开始就让三值专家适应其量化约束。

在初始化策略方面,研究团队证实了使用预训练FFN来初始化路由专家的重要性。相比随机初始化,FFN初始化能带来1.5%的平均性能提升。这就像培训新医生时,有经验医生的指导比完全自学要有效得多。

负载均衡也是训练中的关键考虑因素。研究团队采用了Switch Transformers中提出的辅助损失函数来确保专家使用的均衡性。该损失函数会惩罚过度依赖某个专家的行为,鼓励模型充分利用所有可用的专家资源。在实践中,他们将负载均衡损失的权重设置为0.01,既保证了专家使用的均衡性,又避免了辅助损失过度干扰主要的语言建模目标。

八、兼容性与扩展性的双重优势

MoTE的另一大优势是其出色的兼容性。由于架构中仍保留了全精度的共享专家,MoTE可以无缝集成现有的后训练量化技术。研究团队测试了GPTQ和AWQ等主流量化方法,发现在相同内存预算下,MoTE的性能优势会进一步放大。

在1.5B模型的测试中,当专家内存限制为1.6GB时,使用GPTQ量化的MoE-LLaVA性能下降明显,而MoTE通过对共享专家进行INT4量化就能达到相同的内存目标,同时保持更好的性能。这种优势在更严格的内存限制下尤其明显,体现了MoTE在资源受限环境下的优越性。

研究团队还探索了MoTE的数据扩展能力。他们使用完整的MammoTH-VL数据集(包含1000万样本)对1.5B的MoTE模型进行训练,结果显示模型性能持续提升,证明了MoTE具备良好的数据可扩展性。即使在只使用21.6B训练tokens的情况下,MoTE仍能在多个任务上超越使用更多训练数据的基线模型。

九、技术实现的巧思

在技术实现层面,MoTE采用了多项优化策略来确保训练效率。研究团队使用PyTorch的编译功能将量化操作编译为优化的GPU kernel,显著加速了训练过程。以1.5B模型的Stage III训练为例,MoTE在16块NVIDIA A100 GPU上需要43.3小时,而MoE-LLaVA需要41.8小时,两者训练时间相当。

在推理阶段,研究团队使用了BitBlas库来优化三值矩阵运算,进一步提升了内存效率和计算速度。虽然三值从理论上需要1.58位存储,但BitBlas仍使用INT2格式来存储和处理权重,这是因为当前GPU架构基于二进制系统的设计。

量化感知训练的实现也颇具技巧性。由于量化操作中的许多函数(如四舍五入、截断)是不可微分的,研究团队采用了直通估计器(Straight-Through Estimator)来近似梯度传播。这种方法虽然在理论上是近似的,但在实践中被证明是有效的,能够确保三值专家的稳定训练。

说到底,MoTE代表了大模型压缩领域的一个重要突破。它巧妙地平衡了性能、效率和实用性,为在资源受限设备上部署强大的多模态AI能力提供了新的可能性。研究结果表明,通过合理的架构设计和训练策略,我们可以用更少的资源实现更好的性能,这对于AI技术的普及和应用具有重要意义。

这项研究不仅在技术上取得了突破,更重要的是为AI民主化开辟了新的道路。当强大的AI模型能够在普通设备上高效运行时,更多的用户将能够直接体验到AI技术的便利,而不必依赖昂贵的云端服务。归根结底,这正是人工智能发展的最终目标——让先进的AI技术真正服务于每一个普通人。

有兴趣深入了解技术细节的读者,可以通过论文编号arXiv:2506.14435v1在arXiv平台查阅完整的研究报告,其中包含了详细的实验数据、算法描述和性能分析。

Q&A

Q1:MoTE是什么?它能解决什么问题? A:MoTE是中科院和美团联合开发的新型AI模型压缩技术,全称"三值专家混合模型"。它主要解决大模型内存消耗过大的问题,能让原本需要几十GB内存的AI模型压缩到几GB就能运行,让普通电脑和手机也能使用强大的AI功能。

Q2:MoTE会不会影响AI模型的性能? A:不会,反而在某些情况下性能更好。研究显示,当模型达到1.5B参数以上时,MoTE的表现甚至超过传统方法。在3B参数规模下,MoTE在数学推理等任务上比传统方法高出3-4%的准确率,同时内存消耗却减少了60%以上。

Q3:普通用户什么时候能用上MoTE技术? A:MoTE目前还在研究阶段,但技术已经相当成熟。考虑到美团的参与,预计这项技术会首先在美团的相关产品中应用,然后逐步推广到其他AI应用中。未来1-2年内,我们很可能会在手机App或其他智能设备中看到基于这项技术的AI功能。

转自:至顶网