logits 대신 마지막 레이어의 hidden states를 캐싱하기: 전체 어휘(full-vocab) logits는...
요약
DeepSeek의 기술 보고서는 최고 수준의 엔지니어링 기법을 보여주며, 특히 v4 모델에서 10개 이상의 교사(teacher) 모델을 사용한 온-정책 증류 과정을 다루고 있습니다. 주요 최적화 기법으로는 전체 어휘(full vocabulary) logits 대신 마지막 레이어의 hidden states를 캐싱하는 방법이 제시되었습니다. 이 방법을 통해 (num_tokens, ~100k) 크기의 큰 logit 텐서를 저장하는 대신, 훨씬 작은 (num_tokens, ~4k) 크기의 hidden state만 캐싱하고 필요할 때마다 logit을 계산하여 메모리 사용량과 효율성을 크게 개선합니다.
핵심 포인트
- 여러 교사 모델(teacher models)을 사용하는 온-정책 증류(on-policy distillation) 관리 방법론 제시
- 교사 인덱스에 따라 미니 배치(mini-batch)를 정렬하여 여러 교사 모델의 헤드 로딩 횟수를 최소화함
- 전체 어휘 logits 대신 마지막 레이어 hidden states를 캐싱하는 최적화 기법을 적용함
- 큰 logit 텐서 (num_tokens, ~100k) 대신 작은 hidden state (num_tokens, ~4k)만 캐싱하여 메모리 효율성을 극대화함
DeepSeek의 기술 보고서(tech reports)는 최고 수준의 엔지니어링 (engineering)을 통해 학습 실행 (training runs)을 완수해내기 때문에 항상 읽을 만한 가치가 있습니다.
v4 보고서에는 전체 어휘 (full vocabulary)를 증류 (distilling)하면서 10개 이상의 교사 (teachers) 모델을 사용하여 온-정책 증류 (on-policy distillation)를 어떻게 관리했는지에 대한 숨겨진 보석 같은 정보가 포함되어 있습니다.
교사 (teacher) 모델에 따라 예측 (predictions) 순서를 정하는 것입니다.
그들은 미니 배치 (mini-batch)당 교사 모델의 헤드 (head)를 여러 번 로드하는 것을 피하기 위해 교사 인덱스 (teacher index)에 따라 배치를 정렬합니다. 그런 방식으로 한 명의 교사 모델에 대한 모든 예측을 생성한 다음 다음 모델로 넘어갑니다.
이는 데이터 측면에서 잠재적으로 간단한 변경 사항이 될 수 있습니다.
logits 대신 마지막 레이어의 hidden states를 캐싱하는 것입니다.
전체 어휘 (full-vocab) logits는 (num_tokens, ~100k)인 반면, hidden states는 (num_tokens, ~4k)입니다.
커다란 logit 텐서 (tensor)를 캐싱하는 대신, 더 작은 hidden state를 캐싱하고 logit = prediction_head(hidden_state)를 즉석에서 (on the fly) 계산합니다.
이 정도 규모 (scale)로 학습할 때 다른 팀들이 겪는 엔지니어링 과제 (engineering challenges)에 대해 읽는 것은 언제나 즐거운 일입니다.
다른 세부 사항을 확인하고 싶다면:
AI 자동 생성 콘텐츠
본 콘텐츠는 X @cmpatino_ (자동 발견)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기