【自然语言处理】【大模型】DeepSeek-V2论文解析

论文地址：https://arxiv.org/pdf/2405.04434

相关博客
【自然语言处理】【大模型】语言模型物理学 第3.3部分：知识容量Scaling Laws
【自然语言处理】Transformer中的一种线性特征
【自然语言处理】【大模型】DeepSeek-V2论文解析
【自然语言处理】【大模型】BitNet：用1-bit Transformer训练LLM
【自然语言处理】BitNet b1.58：1bit LLM时代
【自然语言处理】【长文本处理】RMT：能处理长度超过一百万token的Transformer
【自然语言处理】【大模型】MPT模型结构源码解析(单机版)
【自然语言处理】【大模型】ChatGLM-6B模型结构代码解析(单机版)
【自然语言处理】【大模型】BLOOM模型结构源码解析(单机版)

一、简介

• DeepSeek-V2是一个总参数为236B的MoE模型，每个token仅激活21B的参数，并支持128K的上下文长度。
• 提出了Multi-head Latent Attention(MLA)，通过压缩kv cache至隐向量，从而保证高效推理。
• 相比于DeepSeek 67B，DeepSeek-V2实现了更好的表现，节约了42.5%的训练成本，降低了93.3%的kv cache，提升最大吞吐5.76倍。
• 预训练语料包含了8.1T tokens并进一步进行SFT和RL。

二、模型结构

1. MLA(Multi-Head Latent Attention)

​ 传统Transformer采用MHA(Multi-Head Attention)，但是kv cache会成为推理瓶颈。MQA(Multi-Query Attention)和GQA(Grouped-Query Attention)可以一定程度减少kv cache，但效果上不如MHA。DeepSeek-V2设计了一种称为MLA(Multi-Head Latent Attention)的注意力机制。MLA通过低秩key-value联合压缩，实现了比MHA更好的效果并且需要的kv cache要小很多。

1.1 标准MHA

​ 令$d$为embedding维度，$n_h$是注意力头的数量，$d_h$是每个头的维度，${\text{h}}_t\in {\mathbb{R}}^d$是注意力层中第$t$个token的输入。标准MHA通过三个矩阵$W^Q,W^K,W^V\in {\mathbb{R}}^{d_h{n}_h\times d}$来产生${\text{q}}_t,{\text{k}}_t,{\text{v}}_t\in {\mathbb{R}}^{d_h{n}_h}$。
$$\begin{array}{rl}{\text{q}}_t& =W^Q{\text{h}}_t\\ {\text{k}}_t& =W^K{\text{h}}_t\\ {\text{v}}_t& =W^V{\text{h}}_t\end{array}\text{\hspace{1em}(1)(2)(3)}$$
在MHA中${\text{q}}_t,{\text{k}}_t,{\text{v}}_t$会被划分为$n_h$个头：
$$\begin{array}{rl}& [{\text{q}}_{t,1};{\text{q}}_{t,2};\dots ,{\text{q}}_{t,n_h}]={\text{q}}_t\\ & [{\text{k}}_{t,1};{\text{k}}_{t,2};\dots ,{\text{k}}_{t,n_h}]={\text{k}}_t\\ & [{\text{v}}_{t,1};{\text{v}}_{t,2};\dots ,{\text{v}}_{t,n_h}]={\text{v}}_t\\ & {\text{o}}_{t,i}=\sum _{j=1}^t\text{Softmax}(\frac{{\text{q}}_{t,i}^{\top }{\text{k}}_{j,i}}{\sqrt{d_h}}){\text{v}}_{j,i}\\ & {\text{u}}_t=W^O[{\text{o}}_{t,1};{\text{o}}_{t,2};\dots ,{\text{o}}_{t,n_h}]\end{array}\text{\hspace{1em}(4)(5)(6)(7)(8)}$$
其中${\text{q}}_{t,i},{\text{k}}_{t,i},{\text{v}}_{t,i}\in {\mathbb{R}}^{d_h}$是第$i$个注意力头的query、key和value，$W^O\in {\mathbb{R}}^{d\times d_h{n}_h}$是输出投影矩阵。在推理时，所有的key和value都会被缓存来加速推理。对于每个token，MHA需要缓存$2n_h{d}_{h}l$个元素。

1.2 低秩Key-Value联合压缩

​ MLA通过低秩联合压缩key和value来减少kv cache：
$$\begin{array}{r}{\text{c}}_t^{KV}\\ {\text{k}}_t\end{array}$$