이산 확산 모델 증류에서 어휘 전체 비교를 양방향 활성집합으로 줄인다

무엇인가

이 논문은 병렬 디코딩을 쓰는 이산 확산 언어·비전언어 모델을 소수 스텝으로 증류하는 문제를 다룬다. 기존 방식은 각 위치에서 교사와 학생의 토큰별 KL을 최소화하는데, 블록 분포를 위치별 범주분포의 곱으로 표현하는 구조에서는 이 KL이 그 곱 형태 분포를 정확히 일치시키긴 한다. 저자들이 지적하는 문제는 두 가지다. 첫째, 확률 공간에서의 최적화 기하가 공격적인 few-step 증류에서 충분한 감독을 주지 못하고, 둘째, 교사 쪽 후보만 잘라 쓰는 관행 때문에 학생이 선호하는 상태가 최적화 지원 집합 바깥으로 밀려난다. 대안인 Concrete Score Matching(CSM)은 정규화 상수가 해석적으로 소거되는 상대 로그 확률 비를 맞추지만, 어휘 전체 쌍대 비교는 O(|V|²)라 큰 어휘에서는 비용이 감당되지 않는다. 참고로 위치별 목적함수는 모두 위치 간 조건부 독립 가정을 공유하므로, 이 논문도 결합 분포 자체를 복원하지는 않는다고 명시한다.

어떻게 동작하나

첫 번째 구성 요소는 동적 양방향 활성 부분공간이다. 교사 로짓과 학생 로짓을 각각 정렬한 뒤 Kneedle의 최대거리 기준으로 경계를 찾는다. 정렬된 로짓을 x_i=(i-1)/(K_max-1), y_i=(f̄_i-f̄_Kmax)/max(f̄_1-f̄_Kmax, ε)로 정규화하고, 대각선 기준선 y=1-x까지의 수직 거리 (1-x_i)-y_i를 최대화하는 인덱스를 경계 k로 삼는다. 이 지점이 고확률 모드에서 평평한 꼬리로 넘어가는 전이다. 교사에서 k^t, 학생에서 k^s를 각각 구한 뒤 Ω = Unique(TopK(f_T,k^t) ∪ TopK(f_S,k^s))로 합집합을 만든다. 교사 부분집합은 목표 모드를 보존하고 학생 부분집합은 학생이 현재 선호하는 상태를 덮는다. 이로써 쌍대 비교가 O(|V|²)에서 O(K²)로 줄고, K≪|V|다.

무엇과 다른가

두 번째는 평균중심 이산 스코어 매칭이다. 지수 확률비를 직접 계산하는 대신 로짓 공간에서 쌍대 로그비 목적함수를 만든다. 두 상태 i, j의 불일치는 D_i=f_S,i-f_T,i에 대해 Δ_i,j = D_i - D_j로 단순해지고, 가중 제곱합 (1/2W_Ω)ΣΣ p_i p_j (D_i-D_j)²를 최소화한다. 여기서 p는 학생의 어휘 전체 소프트맥스 확률, α는 활성집합 내로 재정규화한 가중치다. 핵심은 분산 항등식으로, 이 O(K²) 쌍대 목적함수가 Σ p_i (D_i - D̄)²와 정확히 같아진다. D̄는 α로 가중한 활성집합 평균 차이 f̄_S - f̄_T이므로, 결국 f̃_S,i - f̃_T,i라는 중심화 로짓 차이를 맞추는 O(K) 목적함수가 된다. 중심화는 상태 간 차이를 보존하면서 공통 가산 이동에 불변이라, 절대 로짓을 맞추는 Direct Logit Distillation보다 학생의 해 공간이 넓다. 구현은 학생 로짓과 정지기울기(detach)한 교정 방향 2p_i(f̃_S,i-f̃_T,i)의 내적으로 처리한다.

어떻게 쓰나

세 번째는 노이즈 레벨 간 전이다. 같은 순방향 확산 궤적에서 뽑은 두 상태 x_t1, x_t2(t2>t1)에 대해, 하위 노이즈 상태 t1을 기준으로 양방향 활성집합 Ω_trans와 가중치를 고정한다. 두 레벨의 중심화 로짓 차이 Δ_S = f̃_S,t1 - f̃_S,t2, Δ_T = f̃_T,t1 - f̃_T,t2를 구해 그 불일치 e_trans를 Σ p_i e²로 정렬하고, 이 전이 항은 t1과 t2에 반대 방향의 교정을 건다. 최종 목적함수는 두 시점의 정적 BCSM 항에 λ_trans를 곱한 전이 항을 더한 형태다.

전제와 한계

실험은 이산 언어모델과 비전언어 모델 양쪽에서 이뤄졌다. LLaDA-8B-Instruct와 Dream-7B-Base를 Qwen2.5-7B 응답 기반 Bespoke-Stratos-17k(최대 600 토큰)로 증류했고, 비전언어는 블록 크기 4인 8B 교사에서 블록 크기 32인 Bard-4B-VL 학생을 17만 장의 이미지-텍스트 쌍으로 학습했다. H800 4장에서 각각 약 23시간, 16시간, 9시간이 걸렸다. LLaDA-8B에서 HumanEval 42.1(NFE=42.9)로 기존 모델 36.0과 CDLM 40.2를 앞섰고, MATH에서는 30.1(NFE=88.8)로 기존 23.7, D2F 29.1(NFE=142.5)을 넘었다. Dream-7B의 동적 변형은 HumanEval 57.9(NFE=110.8)로 full-step 베이스라인을 웃돈다. 비전언어에서는 NFE=1에서 MMMU 35.67로 미증류 학생의 32스텝 점수 32.14를 상회했고, NFE=8에서는 RWQA와 MMVU 모두 32스텝 학생을 넘었으며, Math는 동적 변형이 NFE=16에서 61.73으로 32스텝 학생 61.93에 거의 붙는다. 어휘 전체 CSM 베이스라인은 모든 설정에서 메모리 부족으로 실행되지 않았다고 보고한다.

소거 실험은 OCR 과제, NFE=16 대 미증류 32스텝 기준 70.20에서 이뤄졌다. 중심화를 빼고 절대 로짓을 맞추면 59.60, 어휘 전체 소프트맥스 KD는 71.10, 제안 방식은 71.30이었다. 가중치를 균등 1/|Ω|로 두면 51.90, 활성집합 내부 소프트맥스로만 가중하면 55.10으로, 어휘 전체 확률 질량을 반영하는 가중이 중요함을 보인다. 교사 후보만 쓰는 동적 선택은 68.60, 양방향은 71.30이었고, 고정 Top-K는 64.00~66.60에 그쳤다. 탐색 창 K_max를 64에서 768로 늘리면 65.40에서 70.90으로 개선되고 기본값 1024에서 71.30으로 포화된다. 전이 정렬은 NFE=1에서는 40.90으로 양쪽이 동일해 이득이 없고, NFE=8에서 65.90→66.80, NFE=16에서 71.30→71.90으로 다단계에서만 효과가 나타났다. λ_trans는 0.1에서 71.90으로 최적이고 0.15, 0.2에서는 70.50, 69.00으로 떨어졌다.

실무 관점에서 이 논문이 주는 것은 어휘 크기가 큰 모델에서도 감당 가능한 증류 손실의 구체적인 형태다. 어휘 전체 쌍대 비교를 하지 않으면서 상대 로짓 구조를 보존하려면, 교사 상위 후보만 쓰지 말고 학생 상위 후보를 합집합으로 포함시키고, 경계는 고정 Top-K 대신 로짓 곡선의 꺾임점으로 정하며, 손실은 확률비가 아니라 중심화 로짓 차이로 계산하라는 것이 요지다. 다만 전이 항은 단일 스텝 추론에서는 이득이 없고, λ_trans를 키우면 오히려 정적 정렬이 약해지므로 스텝 예산에 맞춰 조정해야 한다. 또한 K_max가 충분히 크지 않으면 후보 커버리지가 부족해진다.

저자들이 밝힌 한계는 분명하다. 이 방법을 포함한 모든 위치별 목적함수는 교차 위치 조건부 독립 가정을 벗어나지 못해 결합 분포의 의존성을 표현하지 못하고, 위치별로는 그럴듯하지만 함께는 일관되지 않은 조합에 확률 질량을 배정할 수 있다. MS COCO Caption 2014 검증셋 20,252개로 잰 비참조 n-gram 비율에서는 BCSM이 1-gram 30.2%(베이스라인 38.6~45.0%), 4-gram 72.7%(베이스라인 74.5% 초과)로 가장 낮았지만, 이는 국소적 참조 불일치를 줄였다는 의미이지 명시적인 교차 위치 의존성 모델링은 아니다. 긴 시퀀스의 정확 복원은 여전히 어려워 출력이 잘리거나 일부 문자가 어긋나는 실패 사례가 남는다고 적는다.