BitLM 기술 요약: 비트 연속 확산을 통한 멀티 토큰 언어 생성 해방

문제 제기

현재의 대규모 언어 모델 (LLM) 은 어휘 수준의 소프트맥스를 통해 한 번에 하나의 토큰을 생성하는 "다음 토큰" 패러다임 하에서 작동합니다. 이러한 자기회귀 (AR) 방식은 상당한 성공을 거두었지만, 추론 시 근본적인 순차적 병목 현상을 부과합니다. 자연어는 본질적으로 멀티 토큰 단위 (구, n-그램) 로 구조화되어 있지만, 표준 인터페이스는 원자적이고 범주적인 결정을 하나씩 강요합니다.

예측 디코딩이나 블록 단위 병렬 디코딩과 같은 기존 병목 현상 완화 시도들은 종종 근본적인 어휘-소프트맥스 인터페이스를 intact 하게 유지하며, 병렬성을 내재적 속성이 아닌 외부 최적화로 취급합니다. 반면, 비자기회귀 또는 순수 확산 모델들은 종종 효과적인 언어 모델링에 필수적인 인과적 구조를 희생합니다. 본 논문은 어휘 소프트맥스가 역사적 인터페이스 선택일 뿐 근본적인 요구사항이 아니며, 출력 공간의 기하학을 변경하면 인과적 의존성을 훼손하지 않고 더 효율적인 사전 학습과 빠른 추론을 가능하게 할 것이라고 주장합니다.

방법론: BitLM

BitLM 은 표준 어휘 소프트맥스를 고정 길이 이진 공간에서의 조건부 탈잡음으로 대체하는 새로운 아키텍처를 제안합니다. 이 모델은 경량 확산 헤드가 결합된 인과적 LLM 백본을 사용하여 블록 내에서 여러 토큰을 병렬로 실현합니다.

1. 이진 토큰 인터페이스

토큰 ID 를 직접 예측하는 대신, BitLM 은 크기 $V$인 어휘의 각 토큰 $y_i$를 길이 $B = \lceil \log_2 V \rceil$인 고정 길이 이진 코드로 매핑합니다.

• 매핑: 토큰 ID 는 "아날로그 비트"와 유사하게 $0 \to -1$ 및 $1 \to 1$인 이진 벡터 $\{-1, 1\}^B$로 변환됩니다.
• 임베딩: 이러한 이진 코드는 위치별 MLP 를 통해 차원 $B$에서 백본의 은닉 차원 $d$로 승격되어 표준 토큰 임베딩을 대체합니다.

2. 블록-인과적 분해

모델은 크기 $m$인 블록 단위로 텍스트를 생성합니다. 이는 블록-인과적 어텐션 마스크를 사용합니다:

• 동일한 블록 내의 토큰들은 서로 어텐션할 수 있습니다 (공동 실현 허용).
• 블록 $n$의 토큰들은 블록 $1$부터$n-1$까지의 토큰들에만 어텐션할 수 있습니다 (블록 간 인과적 의존성 보존).
• 이를 통해 모델은 표준 AR 모델의 인과적 귀납 편향을 유지하면서 병렬성 정도 ($m$) 를 조절할 수 있습니다.

3. 조건부 확산 헤드

각 목표 블록 $n$에 대해, 모델은 이전 블록 ($n-1$) 의 은닉 상태에 조건부인 탈잡음 문제로 생성을 처리합니다:

• 잡음 과정: 깨끗한 이진 블록 $A_0$이 선형 보간을 통해 가우시안 잡음으로 오염되어 시간 단계 $t \in [0, 1]$에서 잡음 상태 $A_t$가 생성됩니다.
• 탈잡음: 경량 확산 헤드는 잡음 입력 $A_t$, 시간 단계 $t$, 그리고 문맥 조건 $C_{n-1}$(백본의 은닉 상태) 을 기반으로 깨끗한 블록 $\hat{A}_0$을 예측합니다.
• 학습 목적: 모델은 예측된 깨끗한 블록과 실제 깨끗한 블록 ($x_0$-예측) 간의 $L_2$ 손실을 최소화하도록 학습됩니다.
• 추론: 생성은 가우시안 잡음으로 블록을 초기화하고 $K$단계에 걸쳐 반복적으로 탈잡음하는 방식으로 진행됩니다. 마지막으로 하드 시그널 프로젝션 ($\text{sign}(\cdot)$) 이 연속 출력을 이산 이진 초입방체로 매핑한 후, 이를 다시 토큰 ID 로 디코딩합니다.

4. 아키텍처 세부 사항

• 백본: Qwen-3(인과적 트랜스포머) 의 아키텍처를 따릅니다.
• 확산 헤드: BitDance(조건부 적응 레이어 정규화를 사용하는 경량) 의 아키텍처를 따릅니다.
• 학습: FineWeb(350B 토큰) 에서 사전 학습되고 XSum 에서 파인튜닝되었습니다.

주요 기여

1. 이진 코드 공식화: 어휘 수준 예측을 컴팩트한 이진 공간에서의 탈잡음으로 대체하는 언어 모델링 프레임워크를 도입하여, 토큰 생성을 범주적 선택이 아닌 반복적 헌신으로 재구성합니다.
2. 내재적 블록-인과적 병렬성: 블록 간 인과적 의존성을 보존하면서 블록 내에서 멀티 토큰 공동 실현을 자연스럽게 가능하게 하는 아키텍처를 제안합니다. 예측 디코딩과 달리, 이 병렬성은 사후 디코딩 트릭이 아닌 생성 인터페이스의 내재적 결과입니다.
3. 기하학적 관점: 심볼 출력 공간의 기하학 (단순형 대 이진 초입방체) 이 디코딩 체제를 근본적으로 변화시킨다는 관점을 발전시켜, 대규모 어휘 소프트맥스가 언어 모델링을 위한 유일한 실행 가능한 종점이 아님을 시사합니다.

실험 결과

저자들은 0.6B 에서 8B 파라미터까지 다양한 크기의 BitLM 모델을 평가했습니다.

• 확장성: 모델 크기가 증가함에 따라 FineWeb 의 사전 학습 손실이 일관되게 감소하여, 이진 확산 인터페이스가 효과적으로 확장됨을 입증했습니다.
• 추론 설정: XSum 요약 작업에서 최적의 성능은 15 단계의 탈잡음과 9.0 의 분류기 없는 가이드 스케일로 달성되었습니다.
• XSum 성능:
  • 파인튜닝 후, 확산 헤드가 있는 8B BitLM 은 ROUGE-1/2/L 점수 26.05 / 6.44 / 20.12를 달성했습니다.
  • 이 성능은 Lead-3 기준선과 사전 학습된 BitLM 을 능가하여, 모델이 비자명한 요약 매핑을 학습했음을 나타냅니다.
  • 그러나 결과는 강력한 포인터-생성기 기준선 (예: PTGEN+COV) 보다 낮습니다. 저자들은 이러한 결과를 "유망하지만 초기 단계"라고 규정하며, 이진 공식화가 높은 정밀도의 어휘적 실현이 필요한 작업을 위해 아직 완전히 최적화되지 않았음을 지적합니다.

중요성과 주장

본 논문은 BitLM 이 한 번에 하나의 토큰이라는 패러다임이 근본적인 요구사항이 아닌 인터페이스 선택임을 입증한다고 주장합니다. 인과적 추론 (LLM 백본이 처리) 과 심볼적 실현 (확산 헤드가 처리) 을 분리함으로써, BitLM 은 새로운 설계 차원인 심볼 출력 공간의 기하학을 제공합니다.

저자들은 BitLM 이 모든 지표에서 최첨단 소프트맥스 기반 모델을 능가한다고 주장하지 않습니다. 대신 다음과 같이 주장합니다:

1. 출력 기하학을 변경하면 인과적 구조를 희생하지 않고 내재적 병렬 생성을 가능하게 합니다.
2. 성공적인 사전 학습과 파인튜닝을 통해 이진 확산 인터페이스가 다운스트림 텍스트 생성을 위한 실행 가능한 대안임을 보여줍니다.
3. 이 접근법은 병렬 생성의 효율성과 고품질 언어 모델링을 결합한 학습된 이진 코드, 적응형 블록 크기, 하이브리드 아키텍처에 대한 향후 연구의 길을 엽니다.

이 작업은 차세대 아키텍처를 향한 한 걸음으로 제시되며, "대규모 어휘 소프트맥스"는 효율성과 표현력을 해방시키기 위해 대체될 수 있는 역사적 유물임을 시사합니다.