微信扫一扫,关注公众号

  • 科技行者

  • 算力行者

见证连接与计算的「力量」

首页 OmniDraft:高通AI研究院让小模型秒变万能助手,一个68M模型竟能给多个大模型当"草稿员"

OmniDraft:高通AI研究院让小模型秒变万能助手,一个68M模型竟能给多个大模型当"草稿员"

2025-07-11 09:50
分享至:
----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.-
2025-07-11 09:50 科技行者

这项由高通AI研究院的Ramchalam Kinattinkara Ramakrishnan、Zhaocong Yuan等七位研究人员共同完成的研究,发表于2025年7月3日,论文编号为arXiv:2507.02659v1。感兴趣的读者可以通过arXiv平台获取完整论文内容。这项研究解决了一个在人工智能领域颇为头疼的问题:如何让一个小巧的AI模型为各种不同的大型AI模型提供高效的"草稿服务"。

当我们使用ChatGPT、Claude这样的大型语言模型时,它们需要逐字逐句地生成回答,就像一个作家在稿纸上一个字一个字地写作。这个过程很慢,特别是在手机或其他移动设备上使用时更是如此。为了解决这个问题,研究人员想出了一个巧妙的办法:让一个小而快的"草稿员"模型先快速写出初稿,然后让大模型来检查和修正这个初稿。这就像是让一个速记员先快速记录,然后让专业编辑来润色一样。

然而,现实中存在一个棘手的问题。不同的AI模型就像来自不同国家的人,它们使用着不同的"词汇表"。一个专门为Llama模型训练的草稿员,无法直接为Qwen或其他模型提供草稿服务,因为它们对同一个词汇的理解方式不同。这就好比一个习惯了美式英语的速记员,突然要为一个只懂英式英语的编辑工作,两者之间的词汇差异会造成很多误解。

高通AI研究院的团队提出了一个名为OmniDraft的解决方案,它的核心思想是创建一个"万能翻译官",让同一个小型草稿模型能够为任何大型目标模型提供服务。这个方案包含了三个巧妙的创新。

一、跨词汇表的智能翻译系统

研究团队首先解决的是不同模型之间的"语言障碍"问题。他们设计了一个叫做"n-gram缓存"的翻译系统,这个系统就像是一个智能词典,能够记住不同模型之间的词汇对应关系。

传统的做法是只处理两个模型词汇表中完全相同的词汇,这就像两个人只能用共同认识的词汇交流,大大限制了交流的丰富性。而OmniDraft的n-gram缓存更加聪明,它能够处理更复杂的对应关系。比如,草稿模型可能将"snowflake"(雪花)分解为"snow"、"f"、"la"、"ke"四个部分,而目标模型可能将其识别为"snow"和"flake"两个部分。n-gram缓存能够学会这种对应关系,将草稿模型的四个片段正确地组合成目标模型能理解的两个词汇。

这个过程就像一个经验丰富的翻译官,不仅能翻译单个词汇,还能理解不同语言中词汇组合的方式。当草稿模型提出一系列词汇片段时,翻译系统会查看缓存,看看这些片段是否能组合成目标模型更喜欢的形式。如果找到了匹配的组合,就会将多个小片段合并成一个完整的词汇,大大提高了被目标模型接受的可能性。

更重要的是,这个缓存系统是动态学习的。每当系统遇到新的词汇对应关系时,都会将其记录下来,供将来使用。这就像一个翻译官在工作中不断积累经验,遇到的对应关系越多,翻译能力就越强。

二、在线混合蒸馏训练

解决了翻译问题后,研究团队面临的第二个挑战是如何让草稿模型更好地理解目标模型的"思维方式"。他们开发了一种叫做"在线混合蒸馏"的训练方法。

这个过程可以比作师傅带徒弟的学习方式。草稿模型(徒弟)在实际工作中观察目标模型(师傅)的表现,然后调整自己的行为来更好地配合师傅。具体来说,当目标模型接受了草稿模型的建议时,草稿模型会记住这次成功的经验;当目标模型拒绝建议并给出修正时,草稿模型也会从这次"纠错"中学习。

传统的训练方法通常是离线进行的,就像学生在考试前突击复习一样。而OmniDraft采用的是在线学习方式,更像是边工作边学习的学徒制。这种方法的优势在于,草稿模型能够根据具体的使用场景和用户数据不断调整自己,而不是一成不变地使用固定的知识。

混合蒸馏的"混合"体现在训练方法的灵活性上。对于可以直接对应的词汇,系统使用一种叫做"逆向KL散度"的方法来对齐两个模型的概率分布,这就像让徒弟学习师傅对同一个问题的判断方式。对于需要通过n-gram缓存翻译的词汇,系统则使用"最大似然估计"的方法,重点提高这些词汇被正确预测的概率。

研究团队还引入了一个动态权重参数λ,用来平衡这两种训练方式的重要性。这个参数可以根据实际情况调整,比如当遇到的翻译词汇较多时,可以增加翻译相关训练的权重;当直接对应的词汇较多时,则增加概率对齐训练的权重。

三、自适应草稿长度调整

OmniDraft的第三个创新是智能的草稿长度调整机制。这个机制就像一个经验丰富的秘书,能够根据不同情况调整汇报的详细程度。

在实际应用中,草稿模型需要决定每次应该提供多少个词汇建议。提供太少的建议可能无法充分利用加速的潜力,而提供太多的建议则可能导致大部分被拒绝,反而浪费计算资源。传统的做法是使用固定的草稿长度,但这显然不够灵活。

OmniDraft引入了一个"接受率预测头",这个小型神经网络能够预测每个词汇建议被目标模型接受的可能性。基于这些预测,系统会动态计算继续提供更多建议的风险。如果预测显示后续建议被拒绝的概率很高,系统就会提前停止,避免浪费计算资源。

这个预测机制使用了一种叫做"sigmoid函数"的数学工具来估计接受概率,然后计算所有建议中至少有一个被拒绝的总体概率。当这个概率超过预设的阈值时,系统就会停止生成更多建议。这就像一个精明的销售员,能够判断客户的兴趣程度,在合适的时候结束推销。

在在线学习环境中,这个预测头面临着一个特殊的挑战:随着草稿模型不断改进,词汇被接受的概率也在变化,这意味着预测头需要同步调整。研究团队提出了两种解决方案。

第一种是"联合训练"方法,让草稿模型和预测头同时更新。这种方法简单直接,但可能因为两个组件的学习速度不同而产生不稳定性。第二种是"交替训练"方法,为预测头维护一个更大的数据缓冲区,包含历史数据,这样可以提供更稳定的训练环境。实验表明,交替训练方法通常能获得更好的性能。

四、实验验证与性能表现

为了验证OmniDraft的有效性,研究团队进行了大量的实验测试。他们选择了一个仅有68M参数的Llama小模型作为草稿员,并测试了它与多个不同大型模型的配合效果,包括Llama3-8B、Qwen2-7B和Vicuna-7B。

实验涵盖了四个不同的任务领域。在数学推理任务中,他们使用了GSM8K数据集,这个数据集包含了各种小学数学应用题。在编程任务中,他们结合了MBPP和HumanEval两个代码生成数据集。在文本生成方面,他们使用了Alpaca指令跟随数据集。在文本摘要任务中,他们采用了XSum新闻摘要数据集。

实验结果令人印象深刻。在跨词汇表的场景中,传统的直接映射方法(SpDDM)几乎无法工作,接受率通常在0.1左右,加速比甚至低于1,这意味着不仅没有加速反而变慢了。而使用OmniDraft的LDM(直接映射训练)方法能将接受率提升到0.2-0.4的范围,加速比达到1.2-1.6倍。

当加入n-gram损失项后,性能进一步提升。LDM + λLN-gram方法在所有任务上都表现出色,接受率通常能达到0.2-0.4,加速比在1.2-1.7倍之间。特别值得注意的是,在GSM8K数学推理任务上,无论是Llama3-8B还是Qwen2-7B作为目标模型,都能获得最大的加速效果,这可能是因为数学推理任务具有更强的结构性和可预测性。

研究团队还测试了使用LoRA(Low-Rank Adaptation)技术的效果。LoRA是一种参数高效的微调方法,只需要更新模型的一小部分参数。实验显示,即使使用LoRA这种"轻量级"的训练方式,OmniDraft仍然能够获得显著的性能提升,虽然效果略低于全参数微调,但对于资源受限的边缘设备来说,这种方案提供了很好的性能和效率平衡。

在自适应草稿长度调整的实验中,研究团队发现了一些有趣的现象。联合训练方法虽然能够获得更高的接受率,但在某些任务上的加速比反而不如交替训练方法。这表明高接受率不一定直接转化为更好的加速效果,可能是因为联合训练方法容易低估接受概率,导致过早停止生成建议。

五、技术细节与实现要点

OmniDraft的实现涉及许多精巧的技术细节,这些细节的处理直接影响系统的最终性能。

在跨词汇表映射方面,系统需要处理一个微妙的概率分布问题。当多个草稿词汇片段被合并成一个目标词汇时,如何计算这个合并词汇的概率分布是一个关键问题。研究团队提出了一个近似解决方案:对于合并后的n-gram词汇,使用所有组成片段概率的乘积作为其概率;对于原本的前缀片段,从其原始概率中减去n-gram的概率,这样可以避免概率质量的重复计算。

这种处理方式虽然不是理论上的完美解决方案,但在实践中证明是有效的。它避免了重新运行草稿模型来计算完整分布的计算开销,同时保证了概率分布的基本合理性。

在在线学习的实现中,研究团队采用了小批量更新的策略。系统不会在每个样本后立即更新模型,而是积累一定数量的样本后再进行批量更新。这种做法既保证了学习的及时性,又避免了过于频繁的模型更新带来的计算开销。

对于n-gram缓存的管理,系统实现了一个高效的查找和更新机制。缓存使用哈希表结构存储,键是草稿词汇序列,值是对应的目标词汇。为了控制内存使用,系统还可以实现缓存淘汰策略,比如删除使用频率最低的条目。

在自适应草稿调整中,停止阈值γ的选择对性能有重要影响。研究团队发现,不同任务需要不同的阈值设置。对于结构化程度较高的任务(如数学推理和代码生成),可以使用较为激进的阈值(γ=0.7),而对于更加开放性的任务(如文本生成),则需要更保守的阈值(γ=0.3)。

六、实验深度分析与发现

通过深入的消融实验,研究团队揭示了OmniDraft各个组件的具体贡献和相互作用。

n-gram缓存的有效性分析显示了一个有趣的现象:即使在没有专门训练的情况下,仅仅使用n-gram缓存作为后处理步骤,就能显著改善跨词汇表场景下的性能。在GSM8K任务的子集测试中,基础的SpDDM方法接受率只有0.16,而加入n-gram后处理后,接受率提升到0.20,加速比从1.04x提升到1.16x。这证明了词汇翻译本身就是一个重要的瓶颈。

当进一步加入针对n-gram的专门训练后,性能得到更大幅度的提升。LDM + λLN-gram方法能够将接受率提升到0.46,加速比达到1.66x。这表明仅仅解决词汇映射问题是不够的,还需要通过训练来对齐概率分布。

研究团队还分析了n-gram缓存中词汇的分布特征。他们发现,虽然大部分n-gram的出现频率很低,但确实存在一些高频的有用模式。在编程任务中,常见的编程语法结构(如"def "、"return "、"import "等)会被频繁地从多个片段合并为单一词汇。在数学任务中,数学符号和操作符的组合也表现出类似的模式。

关于不同损失函数的比较实验揭示了另一个重要发现。研究团队测试了多种训练目标的组合,包括纯n-gram训练、直接映射KL散度训练、以及各种混合方案。结果显示,纯n-gram训练往往不稳定,这可能是因为n-gram词汇在数据中的占比相对较小,缺乏足够的监督信号。而混合训练方案不仅性能更好,训练过程也更加稳定。

在LoRA微调的秩选择实验中,研究团队测试了从8到128的不同秩值。结果显示,随着秩的增加,性能确实有所提升,但在秩达到32之后,改善变得很小。考虑到参数效率和部署便利性,秋值32提供了最佳的性能-效率权衡。

七、应用前景与实际意义

OmniDraft技术的意义远超出了技术层面的创新,它为人工智能的实际应用开辟了新的可能性。

在边缘计算场景中,OmniDraft特别有价值。手机、平板电脑等移动设备的计算能力有限,无法高效运行大型语言模型。通过OmniDraft,用户可以在设备上运行一个小型的通用草稿模型,然后将其与云端的各种大型模型配合使用。这种架构既保证了响应速度,又提供了模型选择的灵活性。

对于企业级应用,OmniDraft解决了模型部署和维护的复杂性问题。传统上,如果一个组织想要使用多种不同的大型模型,就需要为每种模型单独训练和维护对应的草稿模型。有了OmniDraft,只需要维护一个通用的草稿模型即可,大大降低了运维成本。

在个性化服务方面,OmniDraft的在线学习能力使得系统能够根据特定用户或应用场景进行定制化优化。比如,一个专门处理法律文件的应用可以让草稿模型逐渐学习法律术语的使用模式,提高在该领域的效率。

研究团队还展示了LoRA适配器切换的可能性。通过为不同的任务和目标模型训练不同的LoRA适配器,单个草稿模型可以快速切换到不同的"专业模式"。这就像一个多面手,能够根据需要快速转换角色。

八、技术挑战与未来方向

尽管OmniDraft取得了显著的成果,但研究团队也诚实地指出了当前方案的一些局限性和未来需要解决的问题。

首先是稳定性问题。由于系统只在数据流的单次迭代中进行学习,对于分布差异较大的新数据,仍然可能出现性能不稳定的情况。未来的改进可能需要引入更sophisticated的适应策略,比如元学习或少样本学习技术。

内存管理是另一个实际部署中需要考虑的问题。虽然n-gram缓存通常不会太大,但在资源极度受限的边缘设备上,仍然需要精心设计缓存淘汰策略。研究团队建议可以基于使用频率、最近访问时间或任务相关性来设计智能的缓存管理算法。

特殊词汇的处理也是一个挑战。当前的方案主要针对常规文本词汇设计,对于特殊符号、多语言混合文本或者多模态内容(如包含图像的输入),还需要额外的处理机制。

在跨词汇表场景中明确结合自适应草稿长度调整也是一个有趣的未来方向。当前的实验主要在同词汇表场景中测试了自适应调整,而在跨词汇表场景中,由于n-gram合并的存在,自适应调整的策略可能需要相应的修改。

九、技术实现的工程细节

对于想要复现或应用这项技术的开发者,研究团队提供了详细的实现指导。

在硬件要求方面,实验使用了NVIDIA A100 GPU,但研究团队指出,对于实际部署,特别是边缘设备部署,硬件要求可以大大降低。草稿模型本身只有68M参数,即使在CPU上也能获得可接受的推理速度。

训练超参数的选择经过了仔细的调优。学习率设置为1e-4到2e-5之间,使用AdamW优化器,权重衰减根据是否使用LoRA进行调整。批大小设置为8,这在内存使用和训练稳定性之间提供了良好的平衡。

在数据预处理方面,所有实验都使用了极低的温度设置(0.01),这确保了生成过程的确定性,便于分析和调试。对于实际应用,可能需要根据具体场景调整温度参数。

混合损失中的λ参数统一设置为0.2,这个值在所有测试任务中都表现良好。研究团队指出,虽然针对特定任务可能可以找到更优的λ值,但0.2提供了很好的通用性。

代码实现使用了PyTorch 2.1.0框架,支持混合精度训练(FP16)以提高训练效率。整个系统的模块化设计使得各个组件可以独立测试和优化。

说到底,OmniDraft代表了人工智能系统设计中的一个重要理念转变:从专用化转向通用化,从静态优化转向动态适应。这项由高通AI研究院开发的技术不仅解决了一个具体的技术难题,更为整个AI加速领域提供了新的思路。

通过巧妙的跨词汇表翻译、在线蒸馏学习和自适应草稿调整,OmniDraft证明了一个小型模型确实可以成为多个大型模型的"万能助手"。实验结果显示,仅有68M参数的草稿模型就能为各种7B-8B参数的大型模型提供1.5到2倍的加速效果,这在边缘计算和资源受限环境中具有重要的实用价值。

更重要的是,这项技术的开源友好设计和详细的实现指导,为整个社区提供了一个可以直接应用和进一步改进的基础。随着移动设备算力的不断提升和AI应用的日益普及,OmniDraft这样的通用加速技术将在让AI更加普惠和易用方面发挥重要作用。对于任何对AI加速技术感兴趣的研究人员或开发者,这项研究都提供了宝贵的参考和启发。感兴趣的读者可以通过arXiv:2507.02659v1获取完整的技术细节和实现代码。

Q&A

Q1:OmniDraft是什么?它解决了什么问题? A:OmniDraft是高通AI研究院开发的通用AI加速框架,它让一个小型"草稿员"模型能够为多种不同的大型AI模型提供加速服务。主要解决了传统方案中草稿模型只能配合特定目标模型使用的局限性,实现了"一个草稿员服务所有大模型"的目标。

Q2:OmniDraft会不会替代现有的AI加速方案? A:不会完全替代,但会显著改善现有方案的灵活性和效率。它特别适合需要在多种AI模型间切换的应用场景,以及资源受限的边缘设备。对于单一模型的专用加速,传统方案仍然有其价值。

Q3:普通开发者能使用OmniDraft吗?有什么技术要求? A:可以使用。研究团队在论文中提供了详细的实现指导,代码基于PyTorch框架。虽然实验使用了高端GPU,但实际部署时硬件要求不高,甚至可以在CPU上运行小型草稿模型。开发者可以通过arXiv获取完整技术文档。

分享至
0赞

好文章,需要你的鼓励

推荐文章
----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.-