NeuIPS 2024 | YOCO的高效解码器-解码器架构

该研究提出了一种新的大模型架构,名为YOCO(You Only Cache Once),其目的是解决长序列语言模型推理中的内存瓶颈。YOCO通过解码器-解码器结构的创新设计,显著减少推理时的显存占用并提升了长序列的处理效率。

现有大模型的挑战:长序列缓存的瓶颈

当前的大规模语言模型通常采用基于解码器的Transformer架构,通过缓存前序生成的键值对(KV缓存)来提升推理效率。然而,随着输入序列的增长,KV缓存占据的显存也在急剧增加。例如,一个65B参数量的语言模型,在处理512K长度的序列时,需要约86GB显存,超过了一张H100-80GB GPU的容量。此外,预填充(prefilling)长输入序列的延迟也非常高,限制了模型的实际部署。

YOCO架构:解码器-解码器的新范式

为了应对上述问题,YOCO架构应运而生。它由两个部分组成:

  1. 解码器 (Self-Decoder):用于生成全局KV缓存。

  2. 交叉 解码器 (Cross-Decoder):通过交叉注意力机制复用自解码器生成的KV缓存。

整个模型从外部看起来依然像传统的解码器模型,能够自然适应自回归生成任务。YOCO的独特之处在于,它只需要在自解码器中缓存一次KV对,然后在交叉解码器中多次复用。这种设计不仅大幅降低了显存消耗,还保留了传统模型的全局注意力能力。

YOCO的三大核心优势

1. 显存 节省,支持更长序列

YOCO通过只缓存一次全局KV对,将显存需求显著降低。与传统Transformer相比,YOCO在处理长序列时的显存占用得到了大幅优化。实验数据显示,在65B参数量的模型中,YOCO将KV缓存的显存需求减少了80倍,使得在处理1M长度的输入时,YOCO仅需12.4GB显存,而Transformer则需占用超过100GB显存。

2. 预填充加速,优化用户体验

在推理阶段,YOCO利用交叉解码器复用自解码器的输出,使得模型在预填充阶段可以提前退出,节省计算资源。例如,在处理512K长度的序列时,YOCO的预填充时间仅为6秒,相较于传统Transformer的180秒,实现了约30倍的加速。

3. 提升推理 吞吐量 ,降低部署成本

YOCO显著提升了推理的吞吐量,使每秒生成的token数量大幅增加。在512K长度下,YOCO的吞吐量为每秒43个token,而传统Transformer仅为4.5个,达到了近10倍的提升。这意味着,YOCO能够以更少的硬件资源完成更多的推理任务,从而降低了部署成本。

YOCO的实验表现

实验结果表明,YOCO在多个语言建模任务中展现了极强的竞争力。与同等规模的Transformer模型相比,YOCO在处理长序列任务时,不仅性能更优,而且显存占用和推理延迟也明显更低。

在"针尖查找(Needle-in-a-Haystack)"测试中,YOCO将上下文长度扩展至1M,依然保持了高精度的查询结果。这表明,YOCO在长序列任务中的记忆与处理能力非常出色,适合处理如长文档摘要、代码理解等需要长上下文的任务。

此外,下表展示了YOCO与其他模型在语言建模任务中的性能对比。实验数据显示,YOCO的性能与现有最优模型相当,甚至在某些任务上有所超越,进一步验证了YOCO架构的有效性。

技术实现细节:高效的自解码器模块

YOCO的自解码器模块采用了两种高效的注意力机制来减少计算复杂度:

  1. 滑动窗口注意力(Sliding-Window Attention):通过限制注意力的窗口范围,将KV缓存的复杂度从线性降低到常量级,极大减少了显存的消耗。

  2. 门控保留机制(Gated Retention):结合数据驱动的门控机制,提升了模型的计算效率和生成效果。

YOCO在书籍和代码数据集上的实验结果表明,其累积平均负对数似然(NLL)随着上下文长度的增加而不断降低,显示了其在长序列任务中的优越性能。

总结:未来前景

YOCO通过创新性的解码器-解码器架构,突破了长序列语言模型在推理阶段的内存和延迟瓶颈。在保持模型性能的前提下,显著提升了显存利用率和推理速度,使其成为未来大模型架构的有力候选。未来,YOCO的设计理念或许还可以拓展到更多任务场景中,如长文本生成、对话系统等需要高效长序列处理的领域。

论文链接:

https://openreview.net/forum?id=25Ioxw576r

相关推荐
蹦蹦跳跳真可爱5892 小时前
Python----深度学习(基于深度学习Pytroch簇分类,圆环分类,月牙分类)
人工智能·pytorch·python·深度学习·分类
蚂蚁20143 小时前
卷积神经网络(二)
人工智能·计算机视觉
z_mazin5 小时前
反爬虫机制中的验证码识别:类型、技术难点与应对策略
人工智能·计算机视觉·目标跟踪
lixy5796 小时前
深度学习3.7 softmax回归的简洁实现
人工智能·深度学习·回归
youhebuke2256 小时前
利用deepseek快速生成甘特图
人工智能·甘特图·deepseek
訾博ZiBo6 小时前
AI日报 - 2025年04月26日
人工智能
郭不耐6 小时前
DeepSeek智能时空数据分析(三):专业级地理数据可视化赏析-《杭州市国土空间总体规划(2021-2035年)》
人工智能·信息可视化·数据分析·毕业设计·数据可视化·城市规划
AI军哥6 小时前
MySQL8的安装方法
人工智能·mysql·yolo·机器学习·deepseek
余弦的倒数7 小时前
知识蒸馏和迁移学习的区别
人工智能·机器学习·迁移学习
Allen Bright7 小时前
【机器学习-线性回归-2】理解线性回归中的连续值与离散值
人工智能·机器学习·线性回归