less than 1 minute read

Meta info.
  • Authors: Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve
  • Paper: https://arxiv.org/pdf/2404.19737
  • Affiliation: Meta AI
  • Published: April 30, 2024

TL; DR

한 번에 1개가 아닌 multi-token prediction을 학습하면 모델 성능이 더 좋다고. 4-token prediction을 학습한 LM이 배치가 큰 경우에도 최대 3배 추론 속도 향상 가능.

Untitled

Untitled

Untitled

Suggestion

  • (pic1) N개의 독립적인 output-head로 다음 N개의 토큰을 예측하도록 요청, inference 할 때는 첫 head만 사용하고 나머지 N-1개는 사용하지 않음 (speed-up 용)
  • (pic2) 각 head에서 backpropagation (d.grad)> shared trunk 에서 backpropagation(z.backward(gradient=d.grad)): 모든 head에 대해 loss gradient가 계산되면 shared trunk에 backpropagtion, 즉 모든 head의 정보 합쳐서(누적) update

Effects

  • 모든 logit과 gradient를 동시에 구현하지 않아도 되므로 (runtime에 영향끼치지 않고) peak memory 사용량을 O(nV + d) → O(V + d)로 줄인다고.
  • (pic3) 작은모델에서는 통하지 않는듯 하고, 모델 사이즈 커질수록 성능 향상.
  • HumanEval에서 12%, MBPP에서 17% 성능향상