DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

URL

• paper: https://arxiv.org/pdf/2405.04434
• code: https://github.com/deepseek-ai/DeepSeek-V3

TL;DR

• DeepSeek 系列是国内一家名叫深度求索的科技公司推出的一种混合专家语言模型，这家公司背后是幻方量化。
• DeepSeek-V2 的核心思想是 Multi-head Latent Attention (MLA) 和 DeepSeekMoE 两个模块，这两个模块分别替换了传统 Transformer 层的 Attention Module 和 FFN，是本博客想要探讨的重点。

Architecture

Multi-head Latent Attention (MLA)

MLA 总体计算流程

• 核心思想是将 hidden feature 通过 MLP 映射到 latent space，降低计算（attention）和存储（kv cache）的复杂度。
• 和主流的 kv cache 方案不同，MLA 只需要 cache $c_t^{KV}$ 和 $k_t^R$ 两部分。

MLA 详细计算流程

• 输入：$h_t \in \mathbb R^d$ ，表示第 t 个 token 在某个 attention layer 层的输入特征
• 计算 key / value：
  • $c_t^{KV}=W^{DKV}h_t$ ，其中 $W^{DKV} \in \mathbb R^{d_c\times d}$，DKV 表示 down-projection key value
  • $k_t^C=W^{UK}c_t^{KV}$ ，其中 $W^{UK} \in \mathbb R^{d_hn_h\times d_c}$，UK 表示 up-projection key，$d_c \ll d_hn_h$
  • $v_t^C=W^{UV}c_t^{KV}$ ，其中 $W^{UV} \in \mathbb R^{d_hn_h\times d_c}$，UV 表示 up-projection value，$d_c \ll d_hn_h$
• 计算 query：
  • $c_t^Q=W^{DQ}h_t$ ，其中 $W^{DQ} \in \mathbb R^{d_c'\times d}$，DQ 表示 down-projection query
  • $q_t^C=W^{UQ}c_t^Q$ ，其中 $W^{UQ} \in \mathbb R^{d_hn_h\times d_c'}$，UQ 表示 up-projection query，$d_c' \ll d_hn_h$
• 计算 RoPE：
  • RoPE 是一种在输入 attention layer 之前，对 query 和 key 做 position encoding 的方法
  • RoPE 和 MLA 在设计上是冲突的，因此 MLA 对 RoPE 做了一些修改，主要是：额外计算 multi-head query rotary 和 shared key rotary
  • $[q^R_{t,1},q^R_{t,2},...,q^R_{t,n_h}]=q_t^R=RoPE(W^{QR}c_t^Q)$ ，其中 $W^{QR} \in \mathbb R^{d_h^Rn_h\times d_c'}$，$QR$ 表示 query rotary
  • $k_t^R=RoPE(W^{KR}k_t^C)$ ，其中 $W^{KR} \in \mathbb R^{d_h^R\times d}$，$KR$ 表示 key rotary
  • $q_t,i=[q^C_{t,i},q^R_{t,i}]$ ，将 query 和 query rotary 拼接，得到正式的 query
  • $k_t,i=[k^C_{t,i},k^R_{t}]$ ，将 key 和 key rotary 拼接，得到正式的 key
• 计算 attention：
  • $o_{t,i}=\sum^t_{j=1}Softmax_j(\frac{q^T_{t,i}k_{j,i}}{\sqrt{d_h+d_h^R}})v^C_{j,i}$
• 计算 output：
  • $u_t=W^O[o_{t,1},o_{t,2},...,o_{t,n_h}]$ ，其中 $W^O \in \mathbb R^{d\times d_hn_h}$

MLA 优势

• 大幅降低了 kv cache 的存储空间
  

DeepSeekMoE

• 见 DeepSeekMoE 详解

Thoughts

• MLA 算是对 self-attention 比较深度的改进，兼顾了 kv cache 和 RoPE 的设计，比 MQA 和 GQA 设计看上去更用心。
• DeepSeekMoE 感觉在 GShard 上的改进并不大，主要是在 Expert 的基础上加入了 Sub-Expert 的概念，以及常开的 Isolated Shared Expert 设计。
• 这二者结合，给了 DeepSeek-V2 一个很好的性能提升，可以在相同计算复杂度下，塞下更多的参数量，提高模型表现。