더 강한 텍스트 임베딩을 증류해 확산 언어모델의 생성 품질을 높인다

무엇인가

연속 확산 언어모델(DLM)은 텍스트 임베딩에 잠재 확산을 적용하고, 생성된 임베딩을 다시 토큰으로 디코딩하는 방식으로 자기회귀 생성을 대체하려는 시도다. 이 논문이 던지는 핵심 질문은 어떤 임베딩이 가장 확산하기 좋은(diffusible) 잠재 공간을 만드는가다. 기존 연구는 임베딩과 확산 모델을 함께 수정하거나 상대적으로 오래된 임베딩 모델에 의존해 왔기 때문에, 임베딩 공간 자체의 기여를 분리해내기 어려웠다. 저자들은 확산 쪽을 최근의 ELF(Embedded Language Flow) 프레임워크로 고정하고 임베딩만 바꾸는 방식으로 이 질문을 분리해서 다룬다.

어떻게 동작하나

첫 번째 발견은 같은 계열 안에서 임베딩 모델을 키우면 확산성이 좋아진다는 것이다. T5-small에서 T5-base, T5Gemma-1, T5Gemma-2-270M으로 옮겨가면서 OpenWebText의 25% 서브셋(512토큰 시퀀스, OWT-512)에 동일한 ELF-B를 학습시켰고, PPL-엔트로피 곡선상 성능이 계속 개선됐다. ELF의 기본값인 T5-small 임베딩을 T5Gemma-2-270M으로 교체하면 같은 엔트로피에서 생성 PPL이 약 40% 줄어든다. 반면 인코더 전용 모델인 ModernBERT는 판별 과제에서는 강하지만 확산에는 덜 적합했고 반복적인 문장을 생성했다. 즉 판별 성능이 좋은 임베딩이 반드시 확산성이 좋은 것은 아니다.

무엇과 다른가

두 번째 발견은 스케일을 키운 임베딩조차 완벽하게 생성되지는 않는다는 것이다. 생성된 임베딩이 해당 위치의 후보 단어 임베딩 어디에도 가깝지 않은 실패 모드가 관찰된다. 저자들은 디코딩 헤드의 top-k 후보 임베딩까지의 거리가 실제 임베딩 간 중앙 최근접 거리(nn)의 0.91배를 넘으면 invalid, top-1 확률이 0.6 미만이면 uncertain으로 정의했다. 실제 텍스트 임베딩의 top-1 확률은 0.98 이상이다. 1024토큰 생성 시퀀스 하나에 평균 142개의 uncertain 위치가 있고 그중 109개가 invalid였다. 원인은 T5Gemma-2가 원-핫 레이블로 입력을 재구성하도록 학습되어 임베딩이 정보 밀도 높고 서로 멀찍이 떨어져 있다는 점이다. 최적 디노이저는 후보 임베딩들의 가중 평균인데, 후보들이 분리되어 있으면 그 평균이 후보들 사이 허공에 놓여 방향이 모호해지고, 불완전한 샘플링 궤적은 어느 후보에도 닿지 못한 채 끝난다.

어떻게 쓰나

해법은 증류다. T5Gemma-2 인코더를 학생 인코더로 증류하되, 교사 디코더가 출력하는 확률 분포를 소프트 레이블로 학습시킨다. 목적함수는 위치별로 교사 분포와 학생 분포 사이의 KL 발산이며, 262k 어휘 중 상위 256개 후보만 남겨 확률을 재정규화한다. 학생은 교사의 18개 층에서 균등 간격으로 뽑은 9개 층으로 초기화하고 전체 OWT에서 증류한다. 소프트 레이블은 한 위치를 채울 수 있는 여러 후보에 확률을 퍼뜨리므로, 학생은 그 후보 임베딩들을 서로 가깝게 배치한다. 그 결과 더 연결되고 확산하기 쉬운 잠재 공간이 만들어진다.

전제와 한계

절제 실험도 구체적이다. 소프트 레이블 KL 대신 원-핫 레이블에 대한 교차엔트로피나 교사 임베딩에 대한 MSE를 쓰면 학생이 실제 텍스트 엔트로피 5.43에 도달하지 못했다(각각 4.18 이하, 5.26 이하). 층 선택에서는 18층 전체에서 9층으로 초기화해도 교사와 동등하거나 능가했으므로 개선이 용량 감소 때문이 아니다. 무작위 초기화 학생은 수렴하지 못했다. 증류는 OWT에만 했지만 WikiText 같은 다른 텍스트도 여전히 인코딩·디코딩할 수 있었다. 다만 판별 능력은 일부 희생되어 SST-2 분류 성능이 89.4에서 78.1로 떨어졌다.

주요 실험은 OWT-1024와 LM1B에서 동일한 ELF-B, ELF-M을 학습시켜 수행했고, 서로 다른 토크나이저 문제를 피하려고 GPT-2 토크나이저로 재토큰화했다. 중간 크기 DLM이 OpenWebText에서 실제 텍스트 엔트로피 기준 Gen. PPL 17.8(실제 텍스트 PPL 15.4)을 기록해 GPT-2-M을 Gen. PPL에서 앞섰다. T5Gemma-2 임베딩은 NFE 16~32 같은 극단적으로 낮은 스텝에서 붕괴했지만 증류 임베딩은 더 안정적이었다. LM1B에서도 증류 학생이 원본 T5Gemma-2를 능가했다. 효율 측면에서 ELF-M은 NFE 64에서 1024토큰 시퀀스를 1.5초에, NFE 512에서 11.8초에 생성하며, GPT-2-M은 1024 직렬 스텝에 10.7초가 걸린다. 증류 학생은 T5Gemma-2의 인코딩 비용을 절반으로 줄인다.

개발자 관점에서 이 논문의 실용적 메시지는 임베딩 공간 설계가 확산 모델 설계만큼 중요하다는 것이다. 확산 언어모델 파이프라인에서 잠재 공간을 고를 때 판별 성능이나 다운스트림 벤치마크만 보고 결정하면 생성 단계에서 무효 임베딩이 쏟아질 수 있다. 생성된 임베딩이 후보 단어와 얼마나 떨어져 있는지, 디코더 top-1 확률이 충분히 높은지를 진단 지표로 삼을 만하다. 또 증류가 인코딩 비용을 줄여주므로 추론 지연이 중요한 배포 환경에서도 고려할 수 있다.

저자들이 밝힌 한계는 분명하다. 증류는 OWT에서만 수행했고 판별 성능을 희생하는 맞바꿈이 있다. 더 큰 인코더(T5Gemma-2-1B/4B)로의 추가 스케일링은 미래 과제로 남겼다. few-step/one-step 생성, QA 같은 실용 규모 과제, 자기회귀 모델을 임베딩으로 직접 전환하는 것도 열린 문제다. 또한 서로 다른 인코더를 쓴 실험 행들은 사전학습 인코더가 더 많은 텍스트를 봤기 때문에 완전히 통제된 비교가 아니며, 통제된 비교는 T5Gemma-2 대 학생이다. LM1B는 문제가 있는 코퍼스라 비교보다 점검용으로 쓴다고 밝혔다.