기술 요약: 재학습이 아닌 정렬: 표현 정렬을 통한 자기회귀 언어 모델의 확산 언어 모델 적응

문제 제기

확산 언어 모델 (DLM) 은 비순차적 생성, 양방향 편집, 반복적 정제 측면에서 표준 자기회귀 (AR) 모델보다 뚜렷한 장점을 제공합니다. 그러나 DLM 을 확장하는 것은 계산 비용이 매우 많이 듭니다. 이론적으로 DLM 은 AR 모델의 단일 왼쪽에서 오른쪽으로의 경로에 비해 $L!$개의 생성 경로를 학습해야 하므로, 종종 훨씬 더 많은 계산 자원이 필요합니다.

최근 접근법들은 사전 학습된 AR 체크포인트를 DLM 으로 변환함으로써 이러한 비용을 완화하려 시도합니다. 이러한 방법들은 AR 가중치에서 DLM 을 성공적으로 초기화하지만, 주로 목적 함수, 주의 패턴 또는 샘플링 관습의 수정을 동반한 지속적인 탈노이즈 훈련에 의존합니다. 결정적으로, 기존 레시피는 변환된 DLM 이 원래 AR 사전 학습 동안 학습된 내부 표현 기하학을 명시적으로 보존하도록 제약하지 않습니다. 이는 다음과 같은 핵심 질문을 제기합니다: AR 와 DLM 을 언어 표현을 재학습해야 하는 별개의 패러다임으로 취급해야 할까요, 아니면 AR 사전 학습이 학습한 의미 구조를 변환 과정에서 명시적으로 보존할 수 있을까요?

방법론: REPR-ALIGN

저자들은 REPR-ALIGN을 제안합니다. 이는 표현 학습 문제로서가 아니라 정렬 문제로 DLM 훈련을 다루는 최소주의적 적응 방법입니다. 핵심 가설은 언어 생성의 '어려운 부분'인 의미 및 구문 구조 (표현) 학습은 AR 사전 학습이 이미 마스터했으며, 남은 과제는 이러한 특징들을 반복적이고 임의의 순서인 디코딩 메커니즘에 적응시키는 것뿐이라는 것입니다.

이 방법은 동일한 아키텍처를 가진 두 개의 모델을 포함합니다:

1. 교사: 인과적 주의 (causal attention) 를 가진 고정된 사전 학습 자기회귀 (AR) 트랜스포머.
2. 학생: 동일한 AR 가중치에서 초기화된 양방향 주의 (bidirectional attention) 를 가진 마스킹 확산 언어 모델 (DLM).

훈련 절차

훈련 중, 학생 모델은 손상된 시퀀스 (토큰이 마스킹됨) 를 소비하며 표준 마스킹 탈노이즈 목적 함수를 사용하여 마스킹된 토큰을 예측하려 합니다. 동시에 고정된 교사 모델은 인과적 주의 하에서 정제된 (clean) 시퀀스를 처리합니다.

AR 모델의 내부 기하학을 보존하기 위해, 저자들은 **레이어별 표현 정렬 손실 (layer-wise representation alignment loss)**을 도입합니다. 트랜스포머 스택의 모든 레이어 $\ell$에 대해, 학생의 은닉 상태 ($h^{(\ell)}_D$) 는 코사인 유사도 손실을 사용하여 교사의 은닉 상태 ($h^{(\ell)}_{AR}$) 와 정렬됩니다.

전체 훈련 목적 함수는 다음과 같습니다:
$$\mathcal{L}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda \mathcal{L}_{align}(\theta)$$
여기서:

• $\mathcal{L}_{diff}$는 표준 마스킹 탈노이즈 교차 엔트로피 손실입니다.
• $\mathcal{L}_{align}$은 모든 레이어와 마스킹된 위치를跨한 학생과 교사 은닉 상태 간의 평균 코사인 거리입니다.
• $\lambda$는 가중치 하이퍼파라미터 (기본값 10) 입니다.

주요 설계 선택 사항은 다음과 같습니다:

• 어댑터 부재: 이 방법은 인과적 주의 마스크를 양방향으로 전환하는 것 이상으로 추가 파라미터나 아키텍처 변경이 필요하지 않습니다.
• 고정된 교사: AR 교사는 고정되어 안정적인 표현 좌표계를 제공합니다.
• 시프트 관습: 이 방법은 다음 토큰 인덱싱과 위치별 탈노이즈를 조화시키기 위해 AR-to-DLM 적응에 사용되는 표준 시프트 관습을 준수합니다.

주요 기여

1. 표현 보존의 식별: 이 논문은 AR 모델의 내부 표현 기하학에 대한 명시적 보존이 현재 AR→DLM 변환 레시피에서 누락된 요소임을 식별합니다.
2. 최소주의적 변환 레시피: REPR-ALIGN 은 레이어별 코사인 정렬을 통해 DLM 학생을 고정된 AR 교사에 고정하는 간단하고 어댑터가 없는 방법을 도입합니다. 이 접근법은 모델 규모 전반에 걸쳐 변환 품질과 샘플 효율성을 향상시킵니다.
3. 데이터 및 계산 효율성: 저자들은 AR→DLM 변환이 본질적으로 데이터를 많이 요구하지 않는다고 입증합니다. 표현 정렬을 통해 0.8B 토큰의 작은 데이터 하위 집합으로 훈련하는 것이 동일한 스텝 예산 하에서 50B 토큰 전체 데이터셋으로 훈련하는 것보다 더 좋은 성능을 낼 수 있습니다. 또한, 큰 블록 (임베딩 및 MLP) 을 고정하면 품질을 저하시키지 않고 훈련 처리량을 최대 2$\times$까지 향상시킵니다.
4. 확장성 검증: 저자들은 이 레시피를 사용하여 4B 파라미터 "oDLM"을 훈련했습니다. Dream-7B(AR 초기화를 기반으로 구축된 강력한 공개 DLM) 와 비교했을 때, oDLM 은 더 적은 파라미터와 훨씬 가벼운 데이터 및 계산 예산을 사용함에도 불구하고 HumanEval 및 HumanEval+ pass@10 점수가 더 좋았습니다.

실험 결과

이 방법은 Qwen3 체크포인트 (0.6B, 1.7B, 4B 규모) 를 사용하여 코드 생성 벤치마크 (HumanEval, HumanEval+, MBPP, MBPP+) 에서 평가되었습니다.

• 적응 속도: REPR-ALIGN 은 일관되게 적응을 가속화했습니다. HumanEval pass@10 에서 정렬된 모델은 초기 변환 기간 내내 AR 파인튜닝 및 스크래치 훈련 모두를 능가했습니다. 200k 스텝에서 0.6B 모델은 6.1 점, 1.7B 모델은 기준 대비 9.4 점의 향상을 보였습니다.
• 확장성: 정렬의 이점은 모델 크기가 커짐에 따라 증가했습니다. 1.7B 모델은 0.6B 모델보다 더 큰 절대적 개선을 보였으며, 이는 더 큰 AR 사전 학습 기하학이 보존할 가치가 더 있음을 시사합니다.
• 데이터 효율성: 정렬을 사용하여 0.8B 토큰 무작위 하위 집합으로 훈련하는 것이 동일한 스텝 수에서 50B 토큰 전체 스트림으로 훈련하는 것보다 더 높은 pass@10 점수를 얻었습니다.
• 고정: 토큰 임베딩 및 MLP 블록을 고정하면 처리량이 약 2$\times$ 증가하면서도 성능을 유지하거나 약간 향상시켜, 주요 학습 신호가 표현을 재학습하는 것이 아니라 탈노이즈 메커니즘을 적응시키는 데 있음을 뒷받침합니다.
• 절대 실험 (Ablations):
  • 지표: 코사인 정렬이 L2 거리보다 우수하여, 유용한 신호는 거리 (크기) 가 아니라 기하학적 (방향) 임을 시사합니다.
  • 가중치: $\lambda=10$이 고정과 적응 사이의 최적 균형을 제공했습니다.
  • 레이어: 모든 레이어를 정렬하는 것이 pass@1 에서 가장 좋은 결과를 제공했으며, 중간 및 상부 레이어는 pass@10 에 특히 중요했습니다.

중요성 및 주장

이 논문은 생성 순서와 관계없이 언어 표현은 보편적이라고 주장합니다. DLM 훈련을 정렬 문제로 재정의함으로써, 저자들은 확산 사전 학습의 높은 비용이 AR 사전 학습이 이미 제공한 표현의 불필요한 재학습에 크게 기인한다고 주장합니다.

이 작업의 중요성은 다음과 같은 확산 언어 모델 훈련을 위한 '가이드 기술'을 제공한다는 점에 있습니다:

• 계산 감소: 표준 사전 학습 비용의 일부로 경쟁력 있는 성능을 달성합니다.
• 훈련 단순화: 복잡한 어댑터나 무거운 파인튜닝의 필요성을 제거하고 간단한 정렬 목적 함수에 의존합니다.
• 접근성 민주화: 낮은 데이터 환경과 제한된 계산 자원으로 고품질 DLM 훈련을 실현 가능하게 합니다.

저자들은 그들의 작업을 "재학습"에서 "정렬"로의 전환으로 위치시키며, 이산 생성의 미래 발전은 처음부터 훈련 패러다임을 구축하는 것보다 더 나은 메커니즘 정렬에서 나올 수 있음을 시사합니다.