[DeepSeek] Multi-Head Latent Attention, MLA

DeepSeek-v3가 나오면서 세계가 놀라고 있습니다. 다양한 방법들이 적용된 것으로 보이는데요. 이번에는 딥시크에 적용된 멀티헤드 레이턴트 어텐션이 기존 멀티헤드 어텐션과 어떤 차이가 있는지를 알아보도록 하겠습니다.

 

기존 MHA(Multi-Head Attention)

$$
Q = XW_Q \ K = XW_K \ V = XW_V
$$

기존 멀티헤드 어텐션은 위의 가중치를 통해 어텐션을 계산합니다. 위에서 구한 Q, K, V를 사용하여 아래와 같은 공식을 통해 어텐션을 계산하며, 여러 어텐션 헤드를 연결하여 최종적으로 멀티헤드를 구현하면 저희가 알고 있는 멀티헤드 어텐션이 나오게 됩니다.

$$
Attention = Softmax({{QK^T}\over{\sqrt{d_k}}})V
$$

 

MLA(Multi-Head Latent Attention)

딥시크에서 설명하는 MLA는 Key와 Value 차원을 저차원 잠재 벡터로 압축하여 연산량 및 메모리 사용량을 대폭 감소시키는 방법입니다.

사용되는 각 공식을 살펴보겠습니다.

1. $C_t^{KV}=W^{DKV}h_t$ : 입력 벡터 $h_t$를 $c_t^{KV}$로 압축합니다.
   여기서 $W^{DKV} \in \mathbb{R}^{d_c^\mathbb{R} \times d}$ 는 다운프로젝션 행렬이고, $d_c$는 압축된 저차원을 의미합니다.
2. $k_t^C = W^{UK}c_t^{KV}, v_t^C = W^{UV}c_t^{KV}$ : 저차원 벡터 $c_t^{KV}$를 업프로젝션하여 키와 밸류를 생성합니다.
3. $k_t^R = RoPE(W^{KR}h_t), q_t^R = RoPE(W^{QR}h_t)$ : 쿼리 및 키 벡터에 RoPE를 적용하여 위치 정보를 통합합니다. 여기서 $W^{KR}, W^{QR}$은 RoPE 적용을 위한 벡터입니다.
4. $c_t^Q=W^{DQ}h_t, q_t^C=W^{UQ}c_t^Q$ : 쿼리 벡터도 저차원 압축을 진행합니다. D, U는 위와 동일하게 다운프로젝션, 업프로젝션을 의미합니다.
5. $k_t = [k_t^C;k_t^R], q_t = [q_t^C;q_t^R]$ : 기존 벡터와 RoPE를 적용한 벡터를 concat하여 최종적으로 사용할 쿼리 및 키 벡터를 준비합니다.
6. 최종적으로 다음과 같은 연산을 통해 어텐션을 계산하고, 이후에는 멀티헤드어텐션과 동일합니다. $o_t=\sum^t_{j=1}Softmax_j({{q^T_tk_j}\over{\sqrt{d_h+d_h^R}}})v_j^C$

위의 그림에서 빗금 친 부분은 실제 인퍼런스 중에 캐시되는 부분인데요. 기존 MHA는 모든 K, V를 캐시해두어야 했지만 MLA는 일부만 캐시하면 된다는 장점을 가지고 있습니다.

 

정리

정리하면 MLA는 Key, Value 차원을 저차원으로 매핑시킴으로써 메모리 사용량을 크게 줄이면서도 동일한 수식 구조와 성능을 유지한 방법입니다. LoRA의 방법 중 저차원 매핑을 여기다 적용한 것으로 보입니다.

이 방법은 DeepSeek-V2에서부터 적용한 것으로 보이는데, 다음에는 V3에서 처음 적용한 것으로 보이는 DeepSeekMoE 에 대해 정리해보겠습니다.