近日,华中科技大学软件学院2023级本科生黄佳诺同学以共同第一作者身份参与完成的学术论文“Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inference”被自然语言处理领域国际旗舰会议Annual Meeting of the Association for Computational Linguistics(ACL)录用。该论文由黄佳诺、张尧杰等共同完成,并在上海交通大学张林峰助理教授指导下开展研究。
论文围绕扩散大语言模型高效推理问题展开,提出了一种面向扩散语言模型的免训练缓存淘汰框架MaskKV,为提升长上下文场景下扩散大语言模型的推理速度和显存利用效率提供了新的解决思路。论文被ACL录用,体现了该研究在扩散大语言模型高效推理方向上的创新性与学术价值。

ACL是自然语言处理与计算语言学领域具有重要影响力的国际学术会议之一,在全球范围内享有极高学术声誉。ACL长期汇聚该领域前沿研究成果,是学术界和工业界广泛关注的重要交流平台。
论文题目:Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inference
研究方向:扩散大语言模型、高效推理、KV Cache压缩、长上下文推理
近年来,大语言模型在自然语言处理、智能问答、代码生成和复杂推理等任务中展现出强大的能力。当前主流大语言模型多采用自回归生成方式,即按照从左到右的顺序逐个生成文本,推理过程存在较强的顺序依赖。扩散大语言模型作为一种新兴的非自回归生成范式,通过迭代去噪和并行预测的方式生成文本,有望突破传统自回归模型在推理速度上的限制。然而,扩散大语言模型在推理过程中需要反复更新并缓存整段序列的表示,包括输入提示、已生成内容和掩码标记,使其在长上下文场景下面临较高的计算开销和显存压力。

针对上述问题,论文提出面向扩散大语言模型的高效缓存淘汰框架MaskKV。不同于传统面向自回归模型的缓存压缩方法,MaskKV充分利用扩散语言模型中掩码标记的特殊作用,通过掩码查询引导的评分机制识别对文本重建更为关键的上下文信息,并据此保留重要的提示词缓存。同时,论文进一步设计了自适应缓存预算分配策略,根据不同网络层和不同注意力头的功能差异进行更细粒度的缓存资源分配,从而在降低显存占用的同时保持模型生成性能。

实验结果表明,MaskKV在LLaDA和Dream等代表性扩散大语言模型上均取得了良好效果。在LLaDA模型上,当KV缓存压缩至256对时,MaskKV仍能保持约94%的完整缓存性能;在32K长上下文条件下,该方法最高可实现31倍推理加速,并显著降低峰值显存占用。相关实验表明,MaskKV能够在模型性能、推理速度和显存效率之间取得较好平衡,为扩散大语言模型在资源受限和长上下文应用场景中的部署提供了有效支持。
该研究从扩散语言模型的生成机制出发,重新审视了掩码标记在推理阶段的作用。论文指出,掩码标记不仅是待预测位置的占位符,也可以作为感知上下文重要性的“探针”,帮助模型识别关键缓存并指导缓存资源分配。这一思路区别于传统自回归模型中基于历史注意力累积的缓存淘汰方法,为扩散大语言模型的高效推理优化提供了新的研究视角。

黄佳诺同学参与完成的该项研究聚焦大语言模型高效推理这一人工智能前沿问题,围绕扩散语言模型的结构特性和推理机制开展深入分析,提出了具有实际部署潜力的优化方法,体现了软件学院学生在人工智能、自然语言处理与软件系统优化交叉方向上的科研探索能力和创新意识。
近年来,软件学院坚持面向软件技术前沿和国家战略需求,持续推进人工智能、软件工程等方向的人才培养与科研训练,引导学生在真实科研问题中提升创新能力、工程能力和学术表达能力。未来,学院将继续坚持以高质量人才培养为核心,支持更多青年学子面向国际学术前沿大胆探索、勇于创新,为服务高水平科技自立自强和数字中国建设贡献青春力量。