DEVOCEAN

Part 1 에서는 Knowledge distillation 이 어떤 것인지, 어떤 방법론들이 존재하는지 살펴봤다.

(참고 : https://devocean.sk.com/search/techBoardDetail.do?ID=167285)

Part 2 에서는 distillation 을 이용해 LLM 추론 속도를 높이는 speculative decoding과, 가장 최신의 knowledge distillation 방법론들에 대해서 알아본다.

Another advantages of Distillation

사실 Distillation 의 장점은 student 모델의 성능 개선에만 있는 것이 아니다!

Speculative decoding

(Leviathan et al., 2023, Fast Inference from Transformers via Speculative Decoding)

Transformers는 학습 병렬화를 통해 거대 모델의 학습이 가능하게 했지만, 추론 할 때 엄청 느리다.

Input sequence 에 대한 logit 은 한번에 구할 수 있는데, token 을 생성할 때에는 sequential 하게 생성해야 하기 때문이다.

그래서 나온 Speculative decoding은 큰 모델의 답변과 비슷하게 답변하는 작은 모델을 통해 큰 모델의 속도를 높이는 방법이다.

• 작은 모델은 빠르게 생성할 수 있으니까 일단 token 여러개를 생성한다

• 큰 모델로 token 여러개에 대한 logit 을 한번에 뽑아서, 잘못 생성된 token 을 찾는다.

• 잘못 생성된 token 부터 다시 작은 모델로 토큰을 생성한다 -> 반복

이게 되려면 큰 모델과 작은 모델의 답변이 비슷해야 속도가 올라간다. 만약에 모든 토큰이 다 달라버리면, 괜히 작은모델을 태워서 리소스만 더 잡아먹는 것이다.

DistillSpec

(Zhou et al., 2024, DistillSpec: Improving Speculative Decoding via Knowledge Distillation)

그런데 같은 데이터셋으로 학습했다고 하더라도 작은 모델과 큰 모델의 답변이 비슷할거라는 보장이 없다. 따라서 Distillation 을 통해서 speculative decoding 성능을 극대화하는 방법이 제안되었다.

Distillation을 하면 student model 의 답변들이 teacher 와 비슷해지면서 10~45% 추론 속도가 증가했다.

이 방법론은 실제로 Google 검색 페이지에 적용되어 있다

Google에 적용된 speculative decoding 예시

Advanced knowledge distillation

SKD: speculative knowledge distillation

(Agarwal et al., 2025, On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes)

가장 최근에는 이 Speculative decoding을 응용한 하이브리드 방법론도 제안되었다.

GKD 가 성능이 좋기는 하지만, Student 의 on-policy response 가 썩 좋지 않은 상태에서 GKD 를 하는 게 괜찮을까에 대한 의문이 있다.

이런 상황에서는 off-policy 더라도 퀄리티 좋은 teacher 의 response 를 써서 Supervised KD 를 하는게 나을 수도 있다.이 두 방법론의 장점만을 취한 방법론이 SKD 다.

우선 Student 의 on-policy response 를 뽑는다. 그리고 각 token이 teacher 의 top-K token 에 포함되어 있는지 확인한다.

만약 모두 포함되어 있는 경우라면 on-policy 응답이 꽤 좋은 상황이므로 GKD를 수행한다.

포함되지 않는 토큰이 있는 경우라면 해당 토큰을 Teacher 의 token 으로 교체하여 퀄리티를 높이고, 이를 이용해 Supervised KD를 수행한다.

이 SDK 방법론을 적용했을 때 성능도 오르고 speculative decoding 속도도 올랐다고 한다.

왜냐면 Student 의 성능 관점에서는 on-policy distillation 이 낫고, Teacher 의 speculative decoding 관점에서는 teacher response 를 이용한 distillation 이 낫기 때문이다.

Conclusion

Knowledge distillation 의 trade-off 는 결국 리소스와 성능이다.

성능이 올라가는 방법론일수록 컴퓨팅 리소스 / 시간 리소스를 더 많이 필요로 한다.따라서 상황에 맞춰 적당한 distillation 방법론을 선택하도록 하자

• Online KD 의 장점

  • 데이터가 많이 필요하지 않음

  • Train-test distribution mismatch 가 작기 때문에 Agent 와 같은 long horizon task 에 적합함 (OOD 에 쉽게 빠지지 않음)

• Offline KD 의 장점

  • Teacher 모델로 한번만 response 생성해두거나 logit 을 저장해두면 계속 사용할 수 있음

    • logit을 미리 뽑아두면 teacher model 에게 필요한 GPU 리소스도 아낄 수 있음

  • Student on-policy generation 과정이 없기 때문에 학습 속도가 빠름