借一步网
作者:
在
在人工智能日新月异的时代,大型语言模型(LLMs)的迅猛发展让我们看到了通往通用人工智能(AGI)的曙光。然而,伴随着参数数量和模型能力的不断提升,随之而来的却是推理能耗的急剧攀升以及内存“绑架”问题。今天,我们要讲述的正是如何利用一种全新的架构——多头潜能注意力(MLA),借助数据高效的微调技术,让原本训练自多头注意力(MHA)的模型以极低的成本转变,达到大幅度压缩 KV 缓存并节省资源的目标。
想象一下,一个科研团队要不断扩展一座繁忙的城市,而这座城市中每一条街道都在高速传输信息。大型语言模型就像这座城市,信息在各个神经元之间穿梭不断。而其中,所谓的 Key-Value(KV)缓存就好比城市中无处不在的信号塔。这些信号塔(KV 缓存)的数量直接随着句子长度和模型规模线性增加。尽管 MHA 为模型提供了极高的表达能力和捕捉长程依赖的优势,但随之而来的巨额内存占用与显著的能源消耗也让研发者头疼不已。
此时,DeepSeek 提出了一种全新的注意力架构 —— 多头潜能注意力(MLA),它能够将庞大的 KV 缓存压缩为一个低秩潜在向量,从根源上减少推理时的内存开销并提高推理速度。本文将带你一起探秘从 MHA 到 MLA 的架构迁移,这个过程并非简单“换一个外壳”,而是需要巧妙利用原有模型参数,同时通过数据高效的微调技术,让节能与性能达到近乎完美的平衡。
在 Transformer 架构中,MHA(Multi-Head Attention)是其核心驱动力。对于每个输入序列中的标记(token),模型通过三个线性变换将其分别映射为查询(Query)、键(Key)和值(Value)向量。对每个注意力头来说,其计算过程可以抽象表示为:
其中 是输入序列中第 个标记的表示,、 和 分别是对应头的变换矩阵。随后,通过对查询与键的点积计算注意力得分,经过 Softmax 归一化后,再与值向量加权求和得到输出。在这一过程中,还融合了 Rotary Positional Encoding(RoPE),使得每个向量在传递中蕴含位置信息。
然而,在实际自动回归推理时,每接入一个新的标记,都需要将之前的所有键和值保存到 KV 缓存中,这就导致了内存消耗呈线性增长。更不用说当模型和输入变得非常庞大时,KV 缓存的存储成本就会成为瓶颈。
同时,为了应对长上下文等挑战,还涌现出类似 Grouped-Query Attention(GQA)、Multi-Query Attention(MQA)等变体,这些方法尝试通过共享或者减少部分参数来降低内存占用,但往往会带来模型性能上的妥协。
深度潜能注意力(Multi-Head Latent Attention, MLA)正是为解决上述问题而生。与传统 MHA 不同,MLA 将注意力模块拆分为两个相互补充的分支:
图解来说,就像要将一整箱水(全量 KV 数据)分离成“流动部分”和“存储部分”。位置感知部分就像需要保持流动、高速传输的水,而位置无关部分则可以将大部分水“打包”压缩后存储,从而大大降低总体容量。
在 MLA 中,这一分离不仅实现了内存的节省,还得益于低秩近似技术能够减少部分不重要参数的冗余,同时在不同层和头之间保持参数的重用。
从数学上来看,MLA 的注意力输出可表达为两个独立分支的加权求和,例如对于第 个标记、 个头,其输出计算公式为
这里,前半部分继续利用 RoPE 保持位置敏感,后半部分则借助低秩近似将大量数据压缩为更小的潜在向量,从而实现推理成本的经济化。
传统的 Transformer 模型中,RoPE 将每个查询和键向量分割成多个 2D 子空间,然后按照每个子空间特定的角度旋转,以编码绝对或相对位置。形式上,对于一个查询向量中的第 kkk 个 2D 子空间,其旋转矩阵为
其中 定义为某种频率参数,例如 。当我们试图将传统的全 RoPE 模型迁移至 MLA 架构时,一个重要的步骤便是如何有选择地“移除”部分 RoPE 编码,将其转换为无位置编码(NoPE)形式。
这一做法也称作部分 RoPE(Partial-RoPE)微调。直观上,这就像是给那些对注意力得分贡献较小的维度打包存储,使得它们能够兼容后续的低秩压缩技术。研究人员提出了多种策略来确定哪些频段需要保留 RoPE 信息,包括:
在实际微调过程中,研究表明基于 2 范数的重要性评分()方法往往效果最佳,因为它能够动态捕捉每个头的重要维度,使得被转换为 NoPE 的部分更加适合后续的低秩近似。
对于那些经过部分 RoPE 去除后形成的 NoPE 部分,如何进一步压缩 KV 缓存成为关键。这里研究者利用了矩阵的低秩性质:对于预训练中获得的键和值参数矩阵,通过奇异值分解 (SVD) 能够找出其主成分,并用较低的秩近似原矩阵。
具体而言,有两种策略:
这种方法分别对无 RoPE 部分的键和值参数矩阵进行独立的 SVD 分解,分别截断至所需的秩,再分别重构低秩近似的变换矩阵。公式上可以写为:
之后通过将对角矩阵的平方根与两侧矩阵组合,计算出下投影()和上投影()矩阵,从而将原始高维表示投影至低维潜在空间。
不同于将键和值分别分解,联合 SVD 方法将这两个矩阵拼接后共同分解,以充分捕捉它们之间的交互依赖。公式形式如下:
这种联合因式分解在实验中往往表现出更好的保留预训练知识的能力,从而能使得微调后的模型在推理性能上接近或超过原始 MHA 模型。
通过这样的低秩近似步骤,原本庞大的 KV 缓存可以压缩至原来的 10% 左右甚至更低。例如,在 Llama2-7B 模型上使用联合 SVD 后,KV 缓存的规模可减少 92.19%,而 LongBench 的性能下降却仅有约 0.5% 左右,这无疑为实际部署大模型提供了全新的技术路径。
研究人员将基于 MHA 的多个大语言模型(从 135M 参数到 7B 参数甚至更大)应用于 MHA2MLA 框架,并通过多种常识推理任务、长上下文生成任务进行细致实验。实验结果显示:
下表(模拟图表)直观展示了不同模型在常识推理任务(例如 MMLU、ARC、PIQA 等)上的表现对比:
此外,在长上下文处理任务(例如 LongBench 测试)中,结合 KV 缓存量化技术,MHA2MLA 显示了极高的内存节省和极低的性能损失,为众多应用场景提供了现实可行的解决方案。
尽管 MHA2MLA 架构迁移为大型语言模型在推理阶段带来显著的节能效应,但仍存在一些局限和未来改进的空间。例如:
总的来说,MHA2MLA 提供了一条切实可行的道路,使得已有的大型语言模型不必从零开始重训,就能享受到经济节能推理的红利。这对推动绿色 AI 发展有着积极的意义,也为各类实际应用(例如长文生成、多任务推理)提供了可靠的技术支持。
以上便是从多头注意力到多头潜能注意力(MLA)的迁移之路。通过部分 RoPE 微调与低秩 SVD 分解策略,我们不仅在推理时实现了 KV 缓存的近乎革命性的压缩,而且基本保持了原有模型令人惊叹的性能表现。正如本文所揭示的那样,这是一条通往高效、大规模语言模型部署的必由之路,也为绿色 AI 的未来提供了无限可能。
要发表评论,您必须先登录。
在人工智能日新月异的时代,大型语言模型(LLMs)的迅猛发展让我们看到了通往通用人工智能(AGI)的曙光。然而,伴随着参数数量和模型能力的不断提升,随之而来的却是推理能耗的急剧攀升以及内存“绑架”问题。今天,我们要讲述的正是如何利用一种全新的架构——多头潜能注意力(MLA),借助数据高效的微调技术,让原本训练自多头注意力(MHA)的模型以极低的成本转变,达到大幅度压缩 KV 缓存并节省资源的目标。
想象一下,一个科研团队要不断扩展一座繁忙的城市,而这座城市中每一条街道都在高速传输信息。大型语言模型就像这座城市,信息在各个神经元之间穿梭不断。而其中,所谓的 Key-Value(KV)缓存就好比城市中无处不在的信号塔。这些信号塔(KV 缓存)的数量直接随着句子长度和模型规模线性增加。尽管 MHA 为模型提供了极高的表达能力和捕捉长程依赖的优势,但随之而来的巨额内存占用与显著的能源消耗也让研发者头疼不已。
此时,DeepSeek 提出了一种全新的注意力架构 —— 多头潜能注意力(MLA),它能够将庞大的 KV 缓存压缩为一个低秩潜在向量,从根源上减少推理时的内存开销并提高推理速度。本文将带你一起探秘从 MHA 到 MLA 的架构迁移,这个过程并非简单“换一个外壳”,而是需要巧妙利用原有模型参数,同时通过数据高效的微调技术,让节能与性能达到近乎完美的平衡。
在 Transformer 架构中,MHA(Multi-Head Attention)是其核心驱动力。对于每个输入序列中的标记(token),模型通过三个线性变换将其分别映射为查询(Query)、键(Key)和值(Value)向量。对每个注意力头来说,其计算过程可以抽象表示为:
其中
是输入序列中第
个标记的表示,
、
和
分别是对应头的变换矩阵。随后,通过对查询与键的点积计算注意力得分,经过 Softmax 归一化后,再与值向量加权求和得到输出。在这一过程中,还融合了 Rotary Positional Encoding(RoPE),使得每个向量在传递中蕴含位置信息。
然而,在实际自动回归推理时,每接入一个新的标记,都需要将之前的所有键和值保存到 KV 缓存中,这就导致了内存消耗呈线性增长。更不用说当模型和输入变得非常庞大时,KV 缓存的存储成本就会成为瓶颈。
同时,为了应对长上下文等挑战,还涌现出类似 Grouped-Query Attention(GQA)、Multi-Query Attention(MQA)等变体,这些方法尝试通过共享或者减少部分参数来降低内存占用,但往往会带来模型性能上的妥协。
深度潜能注意力(Multi-Head Latent Attention, MLA)正是为解决上述问题而生。与传统 MHA 不同,MLA 将注意力模块拆分为两个相互补充的分支:
图解来说,就像要将一整箱水(全量 KV 数据)分离成“流动部分”和“存储部分”。位置感知部分就像需要保持流动、高速传输的水,而位置无关部分则可以将大部分水“打包”压缩后存储,从而大大降低总体容量。
在 MLA 中,这一分离不仅实现了内存的节省,还得益于低秩近似技术能够减少部分不重要参数的冗余,同时在不同层和头之间保持参数的重用。
从数学上来看,MLA 的注意力输出可表达为两个独立分支的加权求和,例如对于第
个标记、
个头,其输出计算公式为
这里,前半部分继续利用 RoPE 保持位置敏感,后半部分则借助低秩近似将大量数据压缩为更小的潜在向量,从而实现推理成本的经济化。
传统的 Transformer 模型中,RoPE 将每个查询和键向量分割成多个 2D 子空间,然后按照每个子空间特定的角度旋转,以编码绝对或相对位置。形式上,对于一个查询向量中的第 kkk 个 2D 子空间,其旋转矩阵为
其中
定义为某种频率参数,例如
。当我们试图将传统的全 RoPE 模型迁移至 MLA 架构时,一个重要的步骤便是如何有选择地“移除”部分 RoPE 编码,将其转换为无位置编码(NoPE)形式。
这一做法也称作部分 RoPE(Partial-RoPE)微调。直观上,这就像是给那些对注意力得分贡献较小的维度打包存储,使得它们能够兼容后续的低秩压缩技术。研究人员提出了多种策略来确定哪些频段需要保留 RoPE 信息,包括:
在实际微调过程中,研究表明基于 2 范数的重要性评分(
)方法往往效果最佳,因为它能够动态捕捉每个头的重要维度,使得被转换为 NoPE 的部分更加适合后续的低秩近似。
对于那些经过部分 RoPE 去除后形成的 NoPE 部分,如何进一步压缩 KV 缓存成为关键。这里研究者利用了矩阵的低秩性质:对于预训练中获得的键和值参数矩阵,通过奇异值分解 (SVD) 能够找出其主成分,并用较低的秩近似原矩阵。
具体而言,有两种策略:
1. 分离 SVD(SVD split)
这种方法分别对无 RoPE 部分的键和值参数矩阵进行独立的 SVD 分解,分别截断至所需的秩,再分别重构低秩近似的变换矩阵。公式上可以写为:
之后通过将对角矩阵的平方根与两侧矩阵组合,计算出下投影(
)和上投影(
)矩阵,从而将原始高维表示投影至低维潜在空间。
2. 联合 SVD(SVD joint)
不同于将键和值分别分解,联合 SVD 方法将这两个矩阵拼接后共同分解,以充分捕捉它们之间的交互依赖。公式形式如下:
这种联合因式分解在实验中往往表现出更好的保留预训练知识的能力,从而能使得微调后的模型在推理性能上接近或超过原始 MHA 模型。
通过这样的低秩近似步骤,原本庞大的 KV 缓存可以压缩至原来的 10% 左右甚至更低。例如,在 Llama2-7B 模型上使用联合 SVD 后,KV 缓存的规模可减少 92.19%,而 LongBench 的性能下降却仅有约 0.5% 左右,这无疑为实际部署大模型提供了全新的技术路径。
研究人员将基于 MHA 的多个大语言模型(从 135M 参数到 7B 参数甚至更大)应用于 MHA2MLA 框架,并通过多种常识推理任务、长上下文生成任务进行细致实验。实验结果显示:
下表(模拟图表)直观展示了不同模型在常识推理任务(例如 MMLU、ARC、PIQA 等)上的表现对比:
此外,在长上下文处理任务(例如 LongBench 测试)中,结合 KV 缓存量化技术,MHA2MLA 显示了极高的内存节省和极低的性能损失,为众多应用场景提供了现实可行的解决方案。
尽管 MHA2MLA 架构迁移为大型语言模型在推理阶段带来显著的节能效应,但仍存在一些局限和未来改进的空间。例如:
总的来说,MHA2MLA 提供了一条切实可行的道路,使得已有的大型语言模型不必从零开始重训,就能享受到经济节能推理的红利。这对推动绿色 AI 发展有着积极的意义,也为各类实际应用(例如长文生成、多任务推理)提供了可靠的技术支持。
以上便是从多头注意力到多头潜能注意力(MLA)的迁移之路。通过部分 RoPE 微调与低秩 SVD 分解策略,我们不仅在推理时实现了 KV 缓存的近乎革命性的压缩,而且基本保持了原有模型令人惊叹的性能表现。正如本文所揭示的那样,这是一条通往高效、大规模语言模型部署的必由之路,也为绿色 AI 的未来提供了无限可能。