ELF: 임베딩 공간에 머무는 연속 확산 언어 모델(Continuous Diffusion Language Model)에 대한 연구

ELF 소개

이미지와 영상 같은 연속(continuous) 데이터를 생성할 때, 오늘날 사실상의 표준은 확산 모델(Diffusion Model)과 플로 기반(flow-based) 모델입니다. 가우시안 노이즈에서 출발해 점진적으로 데이터를 복원하는 이 패러다임은 Stable Diffusion, FLUX, Sora 같은 모델의 토대가 되었습니다. 그렇다면 같은 방식을 언어 생성에도 그대로 적용할 수 있을까요? 이 질문에서 출발한 것이 확산 언어 모델(Diffusion Language Model, DLM) 연구입니다.

그런데 언어는 이미지와 결정적으로 다른 점이 하나 있습니다. 픽셀 값은 연속적이지만, 텍스트는 정해진 어휘(vocabulary)에서 고른 이산(discrete) 토큰의 나열이라는 점입니다. 이 때문에 DLM 연구는 크게 두 갈래로 나뉘어 발전해 왔습니다. 하나는 토큰을 연속 표현으로 매핑한 뒤 그 공간에서 노이즈를 제거하는 연속 DLM(Diffusion-LM, CDCD, DiffuSeq 등)이고, 다른 하나는 토큰 공간 위에서 직접 이산 확률 과정을 정의하는 이산 DLM(D3PM, MDLM, Duo 등)입니다. 최근의 흐름은 경험적 성능이 더 좋다는 이유로 이산 DLM에 집중되어 왔습니다.

여기서 ELF 연구진이 던진 질문은 다음과 같습니다. "연속 DLM이 이산 DLM에 뒤처지는 것은 언어가 본질적으로 이산적이기 때문인가, 아니면 단지 알고리즘 설계 선택이 충분히 탐구되지 않았기 때문인가?" ELF(Embedded Language Flows)는 후자라는 답을 내놓습니다. 즉, 이산 영역에 대한 최소한의 처리만으로도 연속 DLM을 충분히 강력하게 만들 수 있다는 것을 보이는 연구입니다. 이 논문은 MIT의 Keya Hu, Linlu Qiu(공동 제1저자), Yiyang Lu, Hanhong Zhao, Tianhong Li, Yoon Kim, Jacob Andreas, 그리고 Kaiming He가 함께 작성했으며, 공식 JAX 구현체가 GitHub에 공개되어 있습니다.

ELF, 임베딩 공간에서 동작하는 연속 확산 언어 모델(Continuous Diffusion Language Model) 개요

ELF의 핵심 아이디어는 이름 그대로입니다. 토큰을 임베딩(embedding)으로 바꾼 뒤, 모델이 마지막 단계(t=1 )에 도달하기 전까지 거의 전적으로 연속 임베딩 공간 안에서만 디노이징(denoising)을 수행하고, 오직 맨 마지막 한 번만 공유 가중치(shared-weight) 네트워크로 이산 토큰으로 변환합니다. 위 도식은 주황색 점으로 표현된 데이터(연속 임베딩)와 보라색 선으로 표현된 디노이징 궤적(trajectory)이 가우시안 노이즈로부터 깨끗한 임베딩으로 흘러가는 과정을 보여줍니다. 이산화(discretization)를 마지막 한 스텝으로 미루는 이 단순한 설계 덕분에, 이미지 도메인에서 검증된 분류기 없는 가이던스(Classifier-Free Guidance, CFG)나 정류 흐름(rectified flow) 같은 기법을 거의 그대로 가져다 쓸 수 있습니다.

결과부터 말하면, ELF의 가장 작은 모델인 ELF-B(105M)는 단 32 스텝의 샘플링만으로 OpenWebText에서 생성 혼란도(Generative Perplexity, Gen. PPL) 약 24를 달성합니다. 이는 약 170M 규모의 기존 이산/연속 DLM들을 더 적은 샘플링 스텝으로 능가하는 수치이며, 무엇보다 증류(distillation) 없이, 그리고 기존 방법 대비 약 10배 적은 학습 토큰으로 얻은 결과라는 점이 핵심입니다.

기존 접근법의 한계와 ELF의 발상 전환

ELF의 설계 의도를 이해하려면, 기존 DLM 계열들이 연속과 이산의 경계를 어떻게 다뤄 왔는지를 먼저 봐야 합니다.

연속 임베딩 기반 DLM (Diffusion-LM, CDCD, DiffuSeq 등)은 토큰 임베딩에 가우시안 노이즈를 직접 더해 디노이징을 학습합니다. 하지만 이들 다수는 매 step마다 토큰 단위 교차 엔트로피(cross-entropy) 같은 이산화 손실을 함께 부과하거나, 심플렉스(simplex) 제약, 토큰 정렬된 비트 인코딩 등으로 인해 궤적이 끊임없이 이산 토큰 공간에 묶여 있습니다. 그만큼 플로 다이내믹스(flow dynamics)의 자유도가 제한됩니다.

잠재 확산(Latent Diffusion) 기반 DLM (LD4LG 및 후속 연구)은 이미지의 LDM처럼 동결된 인코더 표현을 압축한 잠재 공간에서 확산을 수행합니다. 그러나 이 방식은 토큰을 복원하기 위한 별도의 디코더 네트워크를 따로 학습해야 하고, 대개 DDPM 스타일의 노이즈 스케줄을 따릅니다.

이산 DLM (MDLM, Duo 등)은 토큰 공간에서 직접 동작합니다. MDLM은 [MASK] 흡수 상태를 두고 반복적으로 마스크를 해제하며, Duo는 토큰을 균일 분포로 확산시켜 추론 중 토큰을 반복 수정할 수 있게 합니다. 현재 DLM의 주류이지만, 연속량을 전제로 만들어진 CFG 같은 강력한 가이던스 기법을 적용하기가 까다롭고 효과도 제한적이라고 보고되어 있습니다.

ELF의 발상 전환은 이 세 갈래의 약점을 정면으로 겨냥합니다. 즉, (1) 매 step의 토큰 단위 감독 없이 제약 없는 연속 임베딩 공간에서만 디노이징하고, (2) 별도 디코더 없이 마지막 step을 자연스러운 연속에서 이산으로의 디코딩 단계로 재해석하며, (3) DDPM이 아닌 연속 시간 플로 매칭(continuous-time Flow Matching)을 선형(정류 흐름) 보간으로 사용합니다. 이 덕분에 ELF는 이미지/영상 생성에서 빠르게 발전 중인 플로 매칭의 성과들을 거의 그대로 흡수할 수 있습니다.

ELF의 핵심 방법론

ELF는 단일 신경망 \text{net}_\theta 하나로 디노이징과 디코딩을 모두 수행하는 미니멀한 구조를 지향합니다. 아래 그림은 학습과 추론 과정을 한눈에 보여줍니다.

토큰에서 연속 임베딩으로

먼저 토큰 시퀀스 \boldsymbol{s} = [s_1, \ldots, s_L] 를 연속 임베딩으로 매핑합니다. 임베딩 방법은 유연하게 선택할 수 있지만, ELF는 기본값으로 동결된 사전 학습 T5-small 인코더(35M, 임베딩 차원 512)에서 얻은 양방향 문맥 임베딩(contextual embedding)을 사용합니다. 이때 차원이 큰 표현을 그대로 쓰는 대신, 512차원을 128차원으로 선형 사영(projection)했다가 다시 모델의 은닉 차원으로 되돌리는 병목(bottleneck) 설계를 적용합니다. 중요한 점은 이 인코더가 학습 시에만 쓰이고 추론 시점에는 추가 모듈을 전혀 더하지 않는다는 것입니다.

임베딩 공간 위의 Flow Matching과 x-예측

연속 표현을 얻은 뒤, ELF는 플로 매칭으로 노이즈에서 데이터로 향하는 흐름을 정의합니다. 데이터 임베딩을 \boldsymbol{x}, 노이즈를 \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I}) 라 할 때, 노이즈가 섞인 잠재 변수는 선형 보간(정류 흐름)으로 정의됩니다.

\boldsymbol{z}_t = t\boldsymbol{x} + (1-t)\boldsymbol{\epsilon}, \quad t \in [0, 1]

연속 시간에서 흐름의 속도장(velocity field)은 \boldsymbol{z} 의 시간 미분, 즉 \boldsymbol{v} = d\boldsymbol{z}/dt = \boldsymbol{x} - \boldsymbol{\epsilon} 으로 주어집니다. 표준 플로 매칭은 신경망으로 \boldsymbol{v} 를 직접 예측하지만, ELF는 최근 이미지 생성 연구를 따라 깨끗한 임베딩 \boldsymbol{x} 자체를 예측합니다(이를 x-예측, \boldsymbol{x}-prediction 이라 합니다). 학습 목적함수는 예측 속도와 실제 속도 사이의 평균 제곱 오차(MSE)입니다.

\mathcal{L}_{\text{MSE}} = \mathbb{E}_{t, \boldsymbol{x}, \boldsymbol{\epsilon}} \frac{1}{(1-t)^2} \|\boldsymbol{x}_\theta(\boldsymbol{z}_t, t) - \boldsymbol{x}\|^2

x-예측은 ELF에서 두 가지 이유로 중요합니다. 첫째, 768차원처럼 고차원 표현 위에서도 플로 매칭이 잘 동작하게 해 줍니다. 둘째, 깨끗한 임베딩을 예측하는 일은 마지막 step에서 깨끗한 이산 토큰을 예측하는 목표와 자연스럽게 정렬됩니다. 실제로 연구진은 \boldsymbol{v}-예측을 마지막 이산화 step과 가중치를 공유하도록 만들면 성능이 크게 나빠진다는 것을 실험으로 확인했습니다.

마지막 step에서 토큰으로 되돌리기

ELF의 핵심 장치는 마지막 시점(t=1 )을 연속에서 이산으로의 디코딩 단계로 재해석하는 것입니다. 별도의 디코더를 두는 대신, 디노이징을 수행하던 바로 그 네트워크가 가중치를 공유한 채 디코더 역할까지 맡습니다.

t \to 1 에서는 \boldsymbol{z}_t 가 이미 깨끗한 \boldsymbol{x} 에 수렴하므로, 그대로 두면 학습에 의미 있는 입력이 되지 못합니다. 그래서 이 마지막 step에만 토큰 단위 손상(corruption) 과정을 적용해 비자명한 입력 \tilde{\boldsymbol{z}} 를 만들고, 동일한 네트워크가 이를 깨끗한 임베딩으로 복원한 뒤 학습 가능한 언임베딩(unembedding) 행렬 W 로 로짓(logit)을 얻습니다. 이때의 손실은 정답 토큰 \boldsymbol{s} 에 대한 교차 엔트로피입니다.

\mathcal{L}_{\text{CE}} = \mathbb{E}_{\tilde{\boldsymbol{z}}} \left[ \text{CrossEnt}(W\boldsymbol{x}_\theta(\tilde{\boldsymbol{z}}), \boldsymbol{s}) \right]

이 네트워크는 시간 조건 t 외에 디노이즈(denoise)인지 디코드(decode)인지를 알려 주는 이진 모드 토큰(mode token) 을 추가로 입력받습니다. 학습 시에는 한 배치 안에서 80%는 MSE 디노이징 분기로, 20%는 교차 엔트로피 디코딩 분기로 처리되며, 두 분기는 마스킹으로 함께 묶여 추가 학습 비용 없이 계산됩니다. 추론 시에는 마지막 step에서만 디코드 모드로 전환해 \operatorname{argmax} 로 토큰을 뽑습니다.

자기 조건화와 분류기 없는 가이던스(CFG)

생성 품질을 제어하는 강력한 기법인 CFG는 본래 점수 함수나 속도장 같은 연속량을 위해 고안되었기에, 연속으로 정식화된 ELF에 자연스럽게 들어맞습니다. 다만 ELF는 클래스 레이블이 없는 설정이므로, 자기 조건화(Self-Conditioning) 로 CFG에 필요한 조건 신호를 만듭니다. 자기 조건화는 한 번의 예측 \hat{\boldsymbol{x}}' 을 다음 예측의 추가 입력으로 넣어 주는 기법으로, 학습 시 50% 확률로 적용하고 나머지는 널(null) 조건을 씁니다. 추론 시에는 직전 step의 예측을 그대로 조건으로 사용하므로 추가 순전파(forward pass)가 들지 않습니다.

원래의 CFG는 추론 시 조건/무조건 예측을 각각 구해 선형 외삽하므로 step당 두 번의 순전파가 필요합니다. ELF는 이미지 생성에서 발전한 학습 시점 CFG(training-time CFG) 기법을 채택해, 단 한 번의 순전파로 \boldsymbol{x}_{\text{cfg}} 를 직접 모델링합니다. CFG 스케일을 키우면 생성 혼란도는 낮아지지만 엔트로피(다양성)도 함께 줄어드는, 전형적인 품질-다양성 트레이드오프가 나타납니다.

조건부 생성으로의 확장

지금까지는 무조건(unconditional) 생성을 기준으로 설명했지만, ELF는 조건부 생성으로 자연스럽게 확장됩니다. 번역이나 요약처럼 입력 시퀀스가 주어지는 경우, 조건 시퀀스의 깨끗한 임베딩을 손상 없이 모델 입력 앞에 덧붙이고 학습/추론 내내 그대로 보존하면, 모델이 셀프 어텐션(self-attention)을 통해 이를 참조합니다. 이 설정에서 조건 신호 \boldsymbol{c} 는 자기 조건화와 접두 임베딩을 모두 포함하며, 텍스트에서 텍스트로의(text-to-text) CFG로 생성 품질을 제어할 수 있습니다.

샘플링: ODE와 SDE

ELF는 연속 시간/공간에서 정의되므로, 추론 시 결정론적 ODE 솔버와 확률적 SDE 류 샘플러를 모두 지원합니다. 가우시안 노이즈 \boldsymbol{z}_0 에서 출발해 d\boldsymbol{z}_t/dt = \boldsymbol{v}_\theta(\boldsymbol{z}_t, t) 를 오일러(Euler) 등으로 수치 적분하고, 마지막 step에서만 디코딩합니다. SDE 변형은 매 step 미세한 노이즈를 주입하면서 시간 변수를 노이즈 쪽으로 살짝 되돌리는 근사로 구현되며(엄밀한 SDE가 아닌 SDE 영감 샘플러), 뒤에서 보듯 적은 스텝 영역에서 오차 누적을 줄여 ODE보다 유리합니다.

ELF-B의 실제 디노이징 궤적

실험 결과 및 성능 분석

평가 설정

무조건 생성은 약 90억(9B) 토큰 규모의 OpenWebText(OWT)에서 시퀀스 길이 L=1024 로 학습하고, 생성 샘플 1,000개에 대해 사전 학습된 GPT-2 Large로 측정한 생성 혼란도(Gen. PPL, 낮을수록 좋음) 와 다양성 지표인 평균 유니그램 엔트로피(unigram entropy, 높을수록 좋음)를 보고합니다. 조건부 생성으로는 WMT14 독일어에서 영어로(De-En) 번역(BLEU)과 XSum 요약(ROUGE-1/2/L)을 평가합니다. 모델은 Muon 옵티마이저, 학습률 0.002, 배치 크기 512로 학습됩니다. 모델 크기는 ELF-B(105M), ELF-M(342M), ELF-L(652M) 세 가지입니다(여기에 학습 시에만 쓰이는 T5 인코더 35M이 더해집니다).

무조건 생성: 시스템 수준 비교

ELF-B(105M)를 약 170M 규모의 이산 DLM(MDLM, Duo)과 연속 DLM(FLM, LangFlow)과 같은 조건에서 비교한 결과, ELF는 단 32 스텝의 SDE 샘플링(자기 조건화 CFG=3)만으로 생성 혼란도 약 24를 달성하며 기존 방법보다 훨씬 적은 추론 연산으로 더 나은 품질을 보였습니다(위 그림 (a)). 더 인상적인 점은, ELF가 적은 스텝 영역에서 증류를 거친 모델들(MDLM+SDTT, Duo+DCD, FMLM)을 어떠한 증류도 없이 능가한다는 것입니다(그림 (b)). 부록에서 추가로 탐구한 점진적 증류(Progressive Distillation) 버전인 ELF+PD는 1~32 스텝 전 구간에서 증류된 베이스라인들을 더욱 앞섭니다.

데이터 효율성의 차이도 큽니다(그림 (c)). 기존 DLM들이 대개 500B 토큰 이상을 학습에 쓰는 반면(예: MDLM/Duo 약 524B), ELF는 약 45B 토큰만으로 동일 수준 이상의 성능을 냅니다. 이는 약 10배 적은 학습 예산이며, 연구진은 토큰을 더 늘려도 추가 향상이 관찰되지 않았다고 보고합니다.

모델 스케일링

세 가지 크기의 ELF를 비교하면, 모델을 키울수록 생성 혼란도와 엔트로피가 그리는 품질-다양성 프론티어(frontier)가 일관되게 개선됩니다. 핵심은 단일 수치가 아니라 이 프론티어의 위치입니다. 비슷한 생성 혼란도에서 더 큰 모델이 더 높은 엔트로피(더 다양한 생성)를 유지합니다. 기본 운영점(64 스텝 SDE, 자기 조건화 CFG=3)에서의 수치는 다음과 같습니다(ELF-B는 32 스텝 기준).

모델 샘플링 스텝(SDE) Gen. PPL ↓ Entropy ↑
ELF-B (105M) 32 24.1 5.15
ELF-M (342M) 64 21.7 5.18
ELF-L (652M) 64 23.3 5.28

샘플러의 효과는 모든 모델 크기에서 일관됩니다. SDE 샘플링이 ODE보다 프론티어를 더 유리한 방향으로 밀어 줍니다.

조건부 생성: 번역과 요약

조건부 생성에서는 검증 셋으로 고른 최적 설정(64 스텝 ODE, 자기 조건화 CFG=1, 입력 조건 CFG=2)을 사용합니다. ELF-B는 비슷한 규모의 자기회귀(AR) 및 확산 베이스라인을 모두 제치고 번역과 요약 양쪽에서 최고 성능을 기록했습니다.

모델 크기 De-En BLEU ↑ XSum R-1 ↑ R-2 ↑ R-L ↑
AR 99M 25.2 30.5 10.2 24.4
MDLM 99M 18.4 33.4 11.6 25.8
Duo 170M (+35M) 21.3 31.4 10.1 25.0
E2D2 99M 24.8 28.4 8.3 22.0
CDCD - 24.9 - - -
ELF-B 105M (+35M) 26.4 36.0 12.2 27.8

핵심 Ablation

연구진은 ELF-B를 기준으로 주요 설계 선택을 분석했습니다.

  • 임베딩 선택: 문맥 임베딩이 비문맥 임베딩보다, 그리고 사전 학습된 임베딩이 무작위/학습형보다 나은 트레이드오프를 보였습니다. 특히 학습형 임베딩은 디노이저와 임베딩을 동시에 최적화하기 어려워 가장 나빴습니다. 결론은 사전 학습된 문맥 임베딩이 ELF에 가장 적합하다는 것입니다.
  • 디코딩 전략: 디노이저와 디코더가 가중치를 공유하는 방식과, 둘을 분리해 2단계로 따로 학습하는 방식을 비교했습니다. 두 방식의 트레이드오프는 비슷했지만, 공유 가중치 방식이 더 낮은 생성 혼란도 영역까지 확장되면서 학습 단계도 단순화했습니다.
  • 샘플러: 적은 스텝 영역에서 SDE 샘플링이 ODE보다 생성 혼란도를 크게 낮췄습니다. 매 step의 확률성이 오차 누적을 줄여 주기 때문입니다.

결론 및 의의

ELF는 연속 시간 플로 매칭을 임베딩 공간에 적용해, 디노이징 궤적을 거의 끝까지 연속으로 유지하고 이산화는 마지막 한 step으로 미루는 연속 확산 언어 모델입니다. 이 미니멀한 설계 덕분에 이미지 도메인에서 검증된 x-예측, 자기 조건화 기반 CFG, SDE 샘플링 같은 기법을 큰 수정 없이 가져올 수 있었고, 그 결과 더 적은 샘플링 스텝과 약 10배 적은 학습 토큰으로, 증류 없이도 기존 이산/연속 DLM을 능가하는 품질-효율 트레이드오프를 달성했습니다.

이 연구가 던지는 메시지는 분명합니다. 연속 DLM이 그동안 뒤처졌던 것은 언어가 이산적이어서가 아니라, 연속과 이산의 경계를 다루는 방식이 충분히 탐구되지 않았기 때문이라는 것입니다. ELF는 그 경계를 마지막 한 step으로 최소화함으로써, 확산 기반 언어 모델링에서 연속 접근법이 여전히 유망한 방향임을 실증적으로 보여 줍니다.

ELF의 설치 및 빠른 검증

ELF는 JAX와 TPU 환경을 기준으로 작성되어 있으며, README는 이미 학습된 체크포인트로 빠르게 검증할 수 있는 평가 스크립트를 제공합니다. HuggingFace Hub의 embedded-language-flows 조직에 모든 사전 학습 체크포인트가 공개되어 있어, --checkpoint_path 인자에 HF 리포지터리 ID만 넘기면 자동으로 다운로드됩니다.

# 의존성 설치 및 WandB 로그인 (선택)
pip install -r requirements.txt
wandb login YOUR_WANDB_API_KEY

# 무조건 생성 평가 (ELF-B)
cd src/
python eval.py \
    --config configs/training_configs/train_owt_ELF-B.yml \
    --checkpoint_path embedded-language-flows/ELF-B-owt

# 조건 생성 평가: XSum 요약
python eval.py \
    --config configs/training_configs/train_xsum_ELF-B.yml \
    --checkpoint_path embedded-language-flows/ELF-B-xsum

# 조건 생성 평가: WMT14 De-En 번역
python eval.py \
    --config configs/training_configs/train_de-en_ELF-B.yml \
    --checkpoint_path embedded-language-flows/ELF-B-de-en

학습도 동일한 YAML 설정으로 한 줄로 실행할 수 있습니다. 기본 ELF-B 설정은 Muon 옵티마이저, 글로벌 배치 사이즈 512, 5 epoch 기준으로 TPU v5p-64에서 약 7.5시간(에폭당 약 1.5시간)이 소요됩니다. configs/config.pyConfig 데이터클래스가 기본 하이퍼파라미터를 담고, configs/training_configs/*.yml이 태스크별 오버라이드를 정의하므로, 사용자는 바꾸고 싶은 항목만 YAML로 작성하면 됩니다.

# 학습 실행
python train.py --config configs/training_configs/train_owt_ELF-B.yml

체크포인트는 매 epoch마다 자동 저장(최대 10개 유지)되며, --resume 없이 재시작하면 최신 체크포인트로 자동 재개됩니다. 설정에 hf_repo_id를 지정하면 결과가 HuggingFace에 자동 업로드됩니다.

메인 브랜치는 JAX/TPU 기준 구현이지만, README에 따르면 PyTorch 구현이 pytorch_elf 브랜치에 공개되어 있어 PyTorch 사용자도 바로 실험을 시작할 수 있습니다.

라이선스

ELF 저장소는 MIT 라이선스로 공개되어 연구/상업 환경 모두에서 자유롭게 사용/수정/재배포할 수 있습니다. 사전 학습 체크포인트는 HuggingFace embedded-language-flows 조직의 각 모델 카드 라이선스를 별도로 확인하는 것이 안전합니다.

:scroll: ELF 논문

:github: ELF 프로젝트 GitHub 저장소

:hugs: ELF 모델 다운로드(HuggingFace)

더 읽어보기




이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:south_korea:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다!
텔레그램(Telegram)이나 Slack/Discord/Teams/Dooray/GoogleChat 등으로도 새 글 알림을 받으실 수 있습니다. :smiley:

:wrapped_gift: 아래:down_right_arrow:쪽에 좋아요:+1:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck: