미루고 미뤄왔던 Deepseek tech report를 차근차근 읽어봤다. 정말 공부할게 많은 tech report 인것 같다.
크게 봐야할 지점이, attention 안에 있는 KV 캐시를 어떻게 효율적으로 바꿨는지, Feed Forward 부분에서 MOE를 어떻게 사용했는지(parameter를 조금만 쓰고 scaling 해서, 조금의 연산량으로 더 큰 모델을 만들수 있다). 그리고 Multi Token Prediction 이 세가지를 위주로 보면 될것 같다.
1) Multi Head Latent Attention
여기서는 Q,K,V projection을 변형해서, KV cache를 잘 compress 시킨다.
보통, token에서 q, k, v가 나오면 이걸로 attention을 하는데, (q,k,v 차원이 D라고 하자)
새로운 토큰이 들어오면, 매번 key 랑 value 를 저장시켜야한다. 그래서 새로운 token이 들어올때마다 memory가 2D만큼 계속 linear 하게 늘어난다.
Deepseek에서는 이 linear 하게 2D 증가하는 이 memory를 D만큼만 증가해도되게한다.
이걸 어떻게 했는지 살펴보자
먼저
d dimension을 갖는 토큰에서 시작해서, 이 토큰은 쿼리(Latent Ctq), 그리고 Key, Value(latent Ctkv) 이렇게 projected 된다, (이렇게 하면 k, v를 다 caching 하는 기존의 방법 대신, 저기 빨간색으로 표시해놓은 C tkv 부분만 캐싱해주면 된다. 기존에는 2D만큼 linear하게 증가했다면, 지금은 D만큼 증가하게 된다.)
그리고 이 Ctkv를 다시 Key, value로 upsampling해서 project 해준다. 즉 엄청 high dimension으로 upsample 해주고 number of head로 나눠준다. 그리 attention을 하기 전에 RoPE를 query, key 에 apply 하는데, 이 RoPE는 qk가 서로 자신의 relative position을 할수 있게해주는, positional encoding 같은거라고 생각해주면 된다.
Key에 더해지는 RoPE는 original 한 hiddenstate에서 다시 projection 되어서, RoPE embedding 이 나오는데, 이게 각 key 의 head에 concat 되는거다.
그러니까 사실상, KV 캐시를 기존에는 다 저장해서 2D만큼이 계속 저장되었다면, 여기서는 RoPE 임베딩이랑 Latent Ctkv 이거만 저장시키면 되는거다.
이제 쿼리쪽을 보자. 이부분은 캐시를 할 필요가 없으니까 사실 왜 한건지 잘 납득이 안간다. 왜 Q를 compress했을까. 이해할순 없지만 그냥 compress했다고 한다. 그리고 Key 랑 똑같은 방식으로 여기에도 RoPE가더해진다. 한가지 다른 점은 query 에서는 compress된 representation에서 RoPE 임베딩이 나왔는데, Key의 RoPE는 original hidden state에서 나온다.
Decoupled RoPE embedding을 가질 수 있어서, length extention이 더 쉬워져서 이렇게 했다고한다.
파이프라인이랑 대조해서, 차원놀이를 한번 해보면 이해가 더 잘 가는것 같다.
그리고 이렇게 나온 q, k , v가 드디어 attention 연산이 되는데, 이건 기존 attention 계산이랑 똑같다
2) MOE
MOE도 너무유명하니까, 기존 MOE에서 어떻게 달라졌는지를 위주로 살펴보자.
먼저 MOE를 간단하게 설명하자면,
이건 아주 간단한 형태의 Feed Forward Network 이다.
MOE는 이런 FFN이 여러개 있는 형태이다. FFN 각각이 각자 잘하는걸 담당하는 expert들인거고, token이 어떤 expert에 들어갈지, 선택되어서 들어간다.
이게 좋은게, FFN 하나가 n개의 active parameter 이 있다고 치자.
그럼 MOE로 설계를 해서 FFN이 5개가 되면 총 param 수는 5xn개 일거지만, 실제로 activate 한 param은 n개가 되니까, 같은 연산량으로 훨씬 훨씬 많은 정보를 담을 수 있다.
근데 이 MOE의 가장 고질적인 문제가, 예를 들어, 수학을 잘하는 expert, 물리를 잘하는 expert , 음악을 잘하는 expert가 있다고 치자. 만약 수학 데이터가 들어왔으면 이 토큰이 수학을 잘하는 expert한테 들어가야한다는걸 어떻게 알지? 에 대한 문제가 있다. 그리고 수학을 잘하는 expert한테 들어갔다고 치자. 근데 만약에 데이터가 수학 데이터밖에 없으면, 물리 expert랑 음악 expert는 영영 activate 가 안된다.
이런 문제를 해결하기 위해 Deepseek는 어떤 방법을 사용했을까?
차례대로 보자. (12)번 식에서 보면, 각 term의 역할이 각각 skip connection, shared experts, routed expert로 실제 MOE부분은 세번쨰 term이다. 이 세번째 term이 뭔지 각각 (13) (14) (15)식에 설명이 되어있다.
Nr개의 expert들이 있고, 모든 expert들을 더해준다. Active expert들만 더해주는게 아니라, 모든 expert를 더해주는게 핵심이다.
즉 모든 expert들은 router에서 나오는 weight(gi,t) - (13)번식 을 가지고 있을것이고, 인풋에서 나오는 각각의 expert output 이 있을것이다.
weight(gi,t) - (13)번식이 어떻게 나오는지 따라가보면 (15)번식을 보면된다. 먼저 input ut에 대해서, ei를 곱해줌으로써 차원을 각각 expert의 number 에 맞게 늘리고, 각각의 expert에 sigmoid를 씌워주는 느낌으로, 각 expert들이 얼마나 token ut에 자신있는지에 대한 weign를 뽑는다.
이걸 가지고, topK expert를 뽑고, 나머지는 0으로 만든다. 그러니까 물리 문제가 나오면, 수학 체육 과학 물리 음악 예술 음식 뭐 이런 expert들이 있다고 치면, top k를 뽑으면 수학 과학 물리 이런 애들의 weight만 그대로 가져가고 나머지는 0으로 만든다.
그리고 이거를 normalize시키는게 (13)번식.
Complementary Sequence-Wise Auxiliary Loss
extreme한 imbalance를 해결하기 위해 추가적인 loss가 필요하다. 이건 다른 MOE에서도 많이 쓰이는 전략인데 이게 그렇게 좋은 방법은 아니다. 왜냐하면 Lbal를 가장 optimize를 할 수 있는 경우를 생각해보면 그냥 모든 expert에 똑같이 optimize가 되면 되는데, 이게 우리의 최종 목표는 아니니까, 잘 optimize하지는 못한다.
그래서 이 Loss는 아주아주 조금 0.0001정다만 쓰고, Auciliary-Loss-Free 라는걸 사용한다.
Auxiliary-Loss-Free Load Balancing
모든 expert마다 sigmoit weight (Si,t) 가 있는데, 여기에 bias가 붙는다. bias가 클 수록 routing 될 확률이 높아지는것. 이걸 어떻게 optimize하냐면, 해당 expert가 overloaded 되면 bias를 줄이고, underloaded 되면 bias를 늘린다. 즉 특정 expert에 너무 많은토큰이 들어가면 bias를 줄이고, 너무 안들어가는가 싶으면 늘린다.
Multi Token Prediction
다음으로는 Multi Token Prediction을 보자. 아이디어만 먼저 보자면,
우리는 보통 Next Token Prediction을 하는데, 그 이유가 단순히 다음 token을 predict해서, 하나만 예측했다면 뒤에꺼까지 더 예측을 해서 속도를 빠르게 하기 위함이다. 근데 DeepSeek는 next token prediction을, richer information을 얻기 위해 하게 된다.