대규모 언어 모델을 위한 멀티 토큰 예측으로 성능과 효율성 향상

대규모 언어 모델(LLM)들은 GPT, LLaMA 등이 있으며, 인간과 같은 텍스트를 이해하고 생성하는 능력으로 세계를 놀라게 했습니다. 그러나 이러한 모델들의 표준적인 훈련 방법인 “다음 토큰 예측”에는 일부 내재된 제한이 있습니다.

다음 토큰 예측에서, 모델은 이전 단어들이 주어진 시퀀스에서 다음 단어를 예측하도록 훈련됩니다. 이 접근 방식은 성공을 거두었지만, 모델이 장거리 의존성과 복잡한 추론 작업에 어려움을 겪을 수 있습니다. 또한, 교사 강제 훈련 체제와 추론 중에 자율 회귀 생성 프로세스 간의 불일치는 하위 최적의 성능으로 이어질 수 있습니다.

Meta AI의 Gloeckle et al.(2024)의 최근 연구 논문에서는 대규모 언어 모델을 강화하기 위한 새로운 훈련 패러다임인 “멀티 토큰 예측“을 소개합니다. 이 블로그 게시물에서는 이 혁신적인 연구의 핵심 개념, 기술 세부 사항 및 잠재적인 영향에 대해 깊이 있게 다룰 것입니다.

멀티 토큰 예측이란?

멀티 토큰 예측의 핵심 아이디어는 언어 모델이 단 하나의 다음 토큰을 예측하는 대신, 여러 미래 토큰을 동시에 예측하도록 훈련하는 것입니다. 구체적으로, 훈련 중에 모델은 훈련 코퍼스의 각 위치에서 다음 n 토큰을 예측하도록任务되며, 이는 n개의 독립적인 출력 헤드가 공유 모델 트렁크 위에서 작동합니다.

예를 들어, 4토큰 예측 설정에서 모델은 한 번에 다음 4토큰을 예측하도록 훈련됩니다. 이전 컨텍스트가 주어졌을 때, 이 접근 방식은 모델이 더 긴 범위의 의존성과 텍스트의 전반적인 구조 및 일관성을 더 잘 포착하도록鼓励합니다.

토이 예제

멀티 토큰 예측의 개념을 더 잘 이해하기 위해, 간단한 예제를 고려해 보겠습니다.

“빠른 갈색 여우가 게으른 개를 뛰어넘습니다.”

표준적인 다음 토큰 예측 접근 방식에서, 모델은 이전 컨텍스트가 주어졌을 때 다음 단어를 예측하도록 훈련됩니다. 예를 들어, 컨텍스트 “빠른 갈색 여우가 게으른”이 주어졌을 때, 모델은 다음 단어 “개를” 예측하도록任务됩니다.

그러나 멀티 토큰 예측에서, 모델은 여러 미래 토큰을 한 번에 예측하도록 훈련됩니다. 예를 들어, n=4를 설정하면 모델은 한 번에 다음 4토큰을 예측하도록 훈련됩니다. 동일한 컨텍스트 “빠른 갈색 여우가 게으른”이 주어졌을 때, 모델은 시퀀스 “개를 뛰어넘습니다”를 예측하도록任务됩니다.(개 뒤의 공백은 문장의 끝을 나타냅니다.)

모델이 여러 미래 토큰을 한 번에 예측하도록 훈련함으로써, 모델은 더 긴 범위의 의존성과 텍스트의 전반적인 구조 및 일관성을 더 잘 포착하도록鼓励합니다.

기술 세부 사항

저자들은 멀티 토큰 예측을 구현하기 위한 간단하지만 효과적인 아키텍처를 제안합니다. 모델은 공유 트랜스포머 트렁크로 구성되며, 이는 입력 컨텍스트의 잠재적인 표현을 생성합니다. 그리고 n개의 독립적인 트랜스포머 레이어(출력 헤드)가 각 미래 토큰을 예측합니다.

훈련 중에, 순방향 및 역방향 패스는 GPU 메모리 사용량을 최소화하기 위해 주의 깊게 조정됩니다. 공유 트렁크는 잠재적인 표현을 계산하고, 각 출력 헤드는 순차적으로 자신의 순방향 및 역방향 패스를 수행하며, 트렁크 수준에서 기울기를 축적합니다. 이 접근 방식은 모든 로짓 벡터와 그 기울기를 동시에 물리화하는 것을 피하여, 피크 GPU 메모리 사용량을 O(nV + d)에서 O(V + d)로 줄입니다. 여기서 V는 어휘 크기이고 d는 잠재적인 표현의 차원입니다.

메모리 효율적인 구현

멀티 토큰 예측기를 훈련하는 데 있어 하나의 도전은 GPU 메모리 사용량을 줄이는 것입니다. 어휘 크기(V)가 일반적으로 잠재적인 표현의 차원(d)보다 훨씬 크기 때문에, 로짓 벡터가 GPU 메모리 사용량의 병목 현상이 됩니다.

이 도전을 해결하기 위해, 저자들은 순방향 및 역방향 작업의 순서를 조심스럽게 조정하는 메모리 효율적인 구현을 제안합니다. 모든 로짓 벡터와 그 기울기를 동시에 물리화하는 대신, 각 독립적인 출력 헤드에 대해 순차적으로 순방향 및 역방향 패스를 수행하며, 트렁크 수준에서 기울기를 축적합니다.

이 접근 방식은 모든 로짓 벡터와 그 기울기를 동시에 메모리에 저장하는 것을 피하여, 피크 GPU 메모리 사용량을 O(nV + d)에서 O(V + d)로 줄입니다. 여기서 n은 예측되는 미래 토큰의 수입니다.

멀티 토큰 예측의优势

연구 논문은 멀티 토큰 예측을 사용하여 대규모 언어 모델을 훈련하는 몇 가지 매력적인优势를 제시합니다:

1. 개선된 샘플 효율성: 모델이 여러 미래 토큰을 한 번에 예측하도록鼓励함으로써, 멀티 토큰 예측은 모델의 샘플 효율성을 향상시킵니다. 저자들은 코드 이해 및 생성 작업에서 모델의 성능이 향상된 것을 보여주었으며, 최대 13B 파라미터를 가진 모델이 평균적으로 약 15% 더 많은 문제를 해결할 수 있음을 보여주었습니다.
2. 더 빠른 추론: 멀티 토큰 예측과 함께 훈련된 추가 출력 헤드는 자율 추론 디코딩을 위해 사용될 수 있으며, 이는 다양한 배치 크기에서 최대 3배 빠른 추론 시간을 허용합니다.
3. 장거리 의존성의 촉진: 멀티 토큰 예측은 모델이 더 긴 범위의 의존성과 패턴을 포착하도록鼓励합니다. 이는 더 큰 컨텍스트에서 이해와 추론이 필요한 작업에 특히 유용합니다.
4. 알고리즘적 추론: 저자들은 합성 작업에서 실험을 수행했으며, 멀티 토큰 예측 모델이 특히 작은 모델 크기에서 귀납적 헤드와 알고리즘적 추론 능력을 개발하는 데 더 우수함을 보여주었습니다.
5. 일관성과 일관성: 모델이 여러 미래 토큰을 한 번에 예측하도록 훈련함으로써, 멀티 토큰 예측은 모델이 더 일관성 있고 일관된 표현을 개발하도록鼓励합니다. 이는 더 긴 텍스트를 생성하는 작업에 특히 유용합니다. 예를 들어, 이야기, 창의적 글쓰기 또는 지침서 생성.
6. 개선된 일반화: 저자들의 실험은 멀티 토큰 예측 모델이 특히 분포 밖 환경에서 더 나은 일반화 능력을 보여주었다고 제안합니다. 이는 모델이 더 긴 범위의 패턴과 의존성을 포착할 수 있기 때문입니다.

예제와 직관

멀티 토큰 예측이 왜 그렇게 잘 작동하는지에 대한 몇 가지 예제를 고려해 보겠습니다.

1. 코드 생성: 코드 생성의 contexto에서 여러 토큰을 한 번에 예측하는 것은 모델이 더 복잡한 코드 구조를 이해하고 생성하는 데 도움이 될 수 있습니다. 예를 들어, 함수 정의를 생성할 때, 단 하나의 토큰을 예측하는 것은 모델이 함수 시그니처를 올바르게 생성하는 데 필요한 컨텍스트를 제공하지 않을 수 있습니다. 그러나 여러 토큰을 한 번에 예측하면, 모델은 함수 이름, 매개 변수 및 반환 유형 간의 의존성을 더 잘 포착할 수 있으며, 더 정확하고 일관된 코드 생성으로 이어질 수 있습니다.
2. 자연어 추론: 언어 모델이 여러 단계 또는 정보를 포함하는 질문에 대한 답변을 하는 시나리오를 고려해 보겠습니다. 여러 토큰을 한 번에 예측하면, 모델은 추론 프로세스의 다양한 구성 요소 간의 의존성을 더 잘 포착할 수 있으며, 더 일관성 있고 정확한 응답으로 이어질 수 있습니다.
3. 장형 텍스트 생성: 장형 텍스트를 생성할 때, 예를 들어 이야기, 기사 또는 보고서를 생성할 때, 모델은 다음 토큰 예측으로 훈련될 경우 더 긴 기간 동안 일관성을 유지하는 데 어려움을 겪을 수 있습니다. 멀티 토큰 예측은 모델이 텍스트의 전반적인 구조와 흐름을 더 잘 포착하도록鼓励합니다. 이는 더 일관성 있고 일관된 장형 생성으로 이어질 수 있습니다.

제한 및 미래 방향

논문에서 제시된 결과는 인상적이지만, 몇 가지 제한과 개방적인 질문이 있습니다.

1. 최적 토큰 수: 논문은 다양한 n 값(예측할 미래 토큰 수)을 탐색하며, n=4가 많은 작업에 잘 작동한다는 것을 발견했습니다. 그러나 최적의 n 값은 특정 작업, 데이터셋 및 모델 크기에 따라 다를 수 있습니다. 최적의 n 값을 결정하기 위한 원칙적인 방법을 개발하면 추가적인 성능 개선으로 이어질 수 있습니다.
2. 어휘 크기 및 토큰화: 저자들은 멀티 토큰 예측 모델에 대한 최적의 어휘 크기와 토큰화 전략이 다음 토큰 예측 모델과 다를 수 있다고 주목합니다. 이것을 탐색하면 압축된 시퀀스 길이와 계산 효율성 간의 더 나은 트레이드오프로 이어질 수 있습니다.
3. 보조 예측 손실: 저자들은 자신의 작업이 대규모 언어 모델을 위한 새로운 보조 예측 손실의 개발을 위한 관심을 불러일으킬 수 있다고 제안합니다. 대체 보조 손실 및 멀티 토큰 예측과 그들의 조합을 조사하는 것은 흥미로운 연구 방향입니다.
4. 이론적 이해: 논문은 멀티 토큰 예측이 왜 그렇게 잘 작동하는지에 대한 몇 가지 직관과 경험적 증거를 제공합니다. 그러나 이 접근 방식이 왜 그리고 어떻게 그렇게 잘 작동하는지에 대한 더 깊은 이론적 이해는 가치 있습니다.

결론

Gloeckle et al.(2024)의 연구 논문 “멀티 토큰 예측을 통한 더 나은 및 더 빠른 대규모 언어 모델”은 대규모 언어 모델의 성능과 능력을 크게 향상시킬 수 있는 새로운 훈련 패러다임을 소개합니다. 모델이 여러 미래 토큰을 한 번에 예측하도록 훈련함으로써, 멀티 토큰 예측은 장거리 의존성, 알고리즘적 추론 능력 및 더 나은 샘플 효율성을 개발하도록鼓励합니다.

저자들이 제안한 기술적 구현은 우아하고 계산적으로 효율적이며, 이를 대규모 언어 모델 훈련에 적용하는 것이 가능합니다. 또한, 자율 추론 디코딩을 위한 추가 출력 헤드를 사용할 수 있는 것은 실제적인 이점입니다.

まだ 개방적인 질문과 탐구할 영역이 있지만, 이 연구는 대규모 언어 모델 분야에서 흥미로운 발전입니다. 대규모 언어 모델의需求이 계속 증가함에 따라, 멀티 토큰 예측은 이러한 강력한 AI 시스템의 다음 세대에서 핵심 구성 요소가 될 수 있습니다.