Large Flow Language Models

논문 링크

토큰을 하나씩 뽑는 순차 병목

대형 언어 모델은 학습 때 문장 전체를 함께 계산할 수 있지만, 글을 만들 때는 앞 토큰을 알아야만 다음 토큰을 고를 수 있다. GPU가 아무리 빨라도 이 순서를 한 번에 건너뛸 수 없어 생성 지연이 커진다. 이 글은 이 병목을 블록 단위 flow로 푸는 Lafa(Large flow language model)를 다룬다.

Lafa는 이미 학습된 자기회귀 모델을 되살려 쓴다. 한 번의 계산으로 블록 전체의 후보 토큰을 만들고, 자기회귀 예측으로 앞에서부터 검증해 확정한다. 결과적으로 같은 품질 수준에서 기존 자기회귀 대비 최대 2.54배의 처리량을 얻는다. 정확한 조건과 숫자는 뒤에서 다시 짚는다.

왜 flow인가, 그리고 블록 단위로 쪼개는 이유

기존 병렬 생성 방법인 이산 확산(discrete diffusion)은 토큰을 [MASK]로 지웠다가 되살리는 방식이다. 한 스텝에서 여러 위치를 독립적으로 샘플링하면 위치 사이 관계가 끊겨 품질이 떨어질 수 있다.

flow는 이와 다르게 토큰을 지우지 않는다. 토큰 임베딩 벡터에 가우시안 노이즈를 섞은 뒤, 노이즈에서 깨끗한 임베딩으로 이어지는 연속 경로를 거꾸로 적분해 토큰을 만든다. 연속 공간에서 움직이기 때문에 이미지·비디오 생성에서 나온 가속·제어 기법을 그대로 쓸 수 있다. 다만 기존 flow 언어 모델은 소규모·단순 과제에서만 보여졌고, 한 차례 스케일링 시도는 학습 비효율을 지적했으며, 수십억 파라미터로 키우는 방법은 아직 정립되지 않았다(§1, §2).

Lafa는 flow를 블록 단위로 쪼갠다. 전체 문장을 통째로 만드는 대신, 앞서 확정한 글(prefix)을 조건으로 크기 $B$의 토큰 블록을 하나씩 생성하고, 블록 사이는 인과적으로 이어 붙인다. 문장 확률을 블록의 곱으로 인수분해한 셈이다.

블록 단위로 쪼개면 두 가지를 얻는다. 첫째, prefix의 키·값(KV) 캐시—이전 토큰 계산에서 나온 어텐션 중간 결과—를 블록마다 다시 쓸 수 있다. 둘째, 병렬 생성을 블록 안쪽에 한정할 수 있다(§2.1). 토큰을 임베딩 위에서 다루면 토큰 한 개의 차원이 어휘 크기 $|V|$에서 임베딩 차원 $d$로 줄어 학습도 확장하기 쉬워진다.

블록 안의 비동기 flow

블록의 모든 위치가 같은 속도로 노이즈에서 벗어나면, 왼쪽에서 오른쪽으로 읽는 텍스트의 방향성이 사라진다. Lafa는 각 위치 $i$에 자기만의 진행 시계 $\tau^i(t)$를 준다. 앞쪽 위치의 시계가 더 빨리 돌아 앞 토큰이 먼저 드러난다. 노이즈 $x_0$에서 깨끗한 임베딩 $x_1$으로 가는 경로는 다음과 같다.

$\tau$는 고정하지 않는다. 학습 중 손실이 큰 구간을 천천히 지나가도록 데이터에서 배운다. 이렇게 하면 학습 샘플과 추론 스텝이 어려운 노이즈 수준에 몰린다. 학습된 스케줄은 일반적으로 앞 위치가 먼저 진행하는데(Fig. learned_time의 가로축은 기준 시간 $t\in[0,1]$, 세로축은 위치별 $\tau^i(t)$, 곡선은 블록 위치 1, 4, 7, …, 16에 대응한다), 이 좌→우 편향은 수동으로 강제한 것이 아니라 학습에서 나온 것이다(§2.1).

설명용 예시. prefix “오늘 날씨가"의 임베딩 3개가 이미 확정되어 있고, 블록 크기 $B=3$으로 “맑다/흐리다/.“를 만들어야 한다고 하자.

확정된 토큰의 KV를 이어 붙이고 다음 블록으로 넘어간다. 이 예시는 이해를 돕기 위한 것이며, 뒤의 측정값과는 무관하다.

자기회귀 모델에서 flow를 배우는 설계

핵심은 flow 모델을 처음부터 키우는 것이 아니라, 이미 학습된 자기회귀 모델을 블록 디노이저로 개조하는 것이다. 디노이저는 노이즈 낀 블록을 보고 각 위치의 깨끗한 토큰 확률을 예측하는 함수다. 문제는 이 개조가 자기회귀 예측과 prefix 캐시를 해치면 안 된다는 점이다. 논문은 세 가지 설계로 해결한다(§2.2).

첫째, tokenwise 시간 조건부. prefix에는 시간 1(노이즈 없음)을, 블록의 각 위치에는 자기 시간 $\tau^i(t)$를 넣는다. 이 값은 적응형 레이어 정규화(adaLN)—레이어 출력을 입력에 따라 다시 정규화하는 장치—로 주입된다. adaLN을 항등 함수로 초기화하면, 시간 조건이 처음에는 기존 계산을 그대로 두고 점차 학습된다.

둘째, 블록 양방향 어텐션. 블록 내부 토큰끼리는 서로 볼 수 있게 해, 아직 덜 정리된 뒷 토큰이 앞 토큰을 돕게 한다. prefix는 기존처럼 인과적으로만 본다. 덕분에 prefix의 키·값이 블록 상태에 의존하지 않아 디노이징 스텝마다 재사용할 수 있다.

셋째, 임베딩 norm 고정. 임베딩을 자유롭게 갱신하면 노이즈 스케일에 비해 L2 norm이 커져 디노이징이 쉬워지는 편법이 생긴다. 사전학습된 norm을 고정하고 방향만 학습한다.

학습 손실은 세 항을 합친다. prefix에 대한 자기회귀 교차 엔트로피, 블록 토큰에 대한 디노이징 교차 엔트로피, 예측 임베딩과 실제 임베딩의 제곱 오차다(식 6). 자기회귀 손실을 같이 두면 flow 정확도까지 올라간다. 이 공동 학습 덕분에 flow를 배우면서도 자기회귀 능력이 보존되거나 오히려 좋아진다. 멀티 토큰 예측이 다음 토큰 예측을 개선한다는 기존 관찰과도 일치한다(§2.2, §3.2).

위 그림(원문 Figure training)은 병렬 학습의 상태 변화를 보여준다. 위쪽 깨끗한 시퀀스는 인과 어텐션으로 자기회귀 예측을 만들고, 아래쪽 복사본은 블록 크기 3으로 나뉘어 각 블록이 앞의 깨끗한 토큰만 보고 블록 내부는 양방향으로 디노이징 예측을 만든다. 한 번의 forward pass로 자기회귀 손실과 모든 블록의 디노이징 손실을 동시에 얻어, 겹치는 prefix를 다시 계산하지 않는다.

ablation으로 확인한 설계 효과

0.5B·블록 4·32 오일러 스텝에서 각 설계를 한 번씩 되돌려 놓고 평균 정확도를 비교했다(Table design-ablations).

어텐션·시간·임베딩·손실 네 축에서 모두 원래 설계가 낫다. 특히 자기회귀 손실을 더하면 flow 정확도가 오히려 올라가며, AR 보존과 flow 생성이 서로를 해치지 않는다.

flow map 증류로 스텝 수 줄이기

flow 디코딩은 블록마다 디노이저를 여러 번 불러야 해 비용이 크다. 이를 줄이기 위해 flow map 증류를 쓴다. 상태를 시간 $s$에서 $t$로 직접 옮기는 함수 $\hat X_{s,t}$를 학생 모델로 배워, 오일러 스텝 여러 번을 한 번의 점프로 대체한다. 사선($s=t$) 예측은 교사와 맞추고, 비사선($s\lt t$) 예측은 짧은 점프 두 번을 합친 결과에 맞추는 반군(semigroup) 방식이다. 학생은 prefix 캐시를 유지하며 최소 1스텝($\hat X_{0,1}$) 생성까지 가능하다(§2.2).

초안과 검증을 한 번에: self-speculative decoding

flow 디코딩으로 AR 수준 품질을 내려면 스텝을 여러 번 밟아야 한다. Lafa는 비싼 다중 스텝을 버리고, 1스텝 flow로 블록 초안을 뽑은 뒤 자기회귀 예측기로 검증하는 방식을 기본으로 쓴다(§2.3).

초안은 노이즈에서 한 번의 디노이저 평가로 얻는다. 초안의 각 위치를 자기회귀 예측과 앞에서부터 비교하고, 일치하는 가장 긴 prefix만 확정하며 나머지는 버린다. 첫 토큰부터 다르면 자기회귀 쪽 토큰으로 교체한다. 최종 출력은 모델의 탐욕적(greedy) 자기회귀 출력과 정확히 같다. flow는 속도를 올리는 데만 쓰이고, 품질 기준은 탐욕적 자기회귀와 같아진다.

단순 구현에서는 검증 forward와 다음 초안 forward가 따로 필요해 손해가 난다. 논문은 이를 융합(fused) 디코딩으로 없앤다. 이전 초안의 검증과 다음 블록의 초안 생성을 한 번의 forward pass에 섞는 것이다.

위 그림(원문 Figure fused_decoding)은 블록 크기 3에서 한 단계가 어떻게 흐르는지 보여준다. 가로 방향은 시퀀스 위치, 세로 블록은 각 forward에 들어간 입력 조각이다.

깨끗한 초안은 prefix처럼 붙여 인과 어텐션으로 검증 로짓을 만들고, 동시에 “앞 초안이 $j$개 수용되는 경우"에 해당하는 후보 노이즈 블록을 붙여 한 번에 처리한다. 검증 결과 실제로 수용된 길이의 후보 블록을 골라 다음 초안으로 삼아, 별도의 재초안 forward를 없앤다.

이 방식은 학습된 모델 하나로 초안과 검증을 모두 수행하므로 self-speculative(자기 추측)라 부른다. 별도의 작은 초안 모델이나 추가 파라미터가 필요 없다. 대안으로 자기회귀 검증 대신 토큰 확률이 임계값을 넘는 prefix만 수용하는 confidence 디코딩도 가능하지만, 임계값을 낮추면 속도 대신 정확도가 떨어지므로 논문은 self-speculative를 기본으로 둔다(§3.3, Appendix confidence_results).

결과: 품질은 유지하고 속도를 올린다

모델은 Qwen2.5 0.5B·7B Instruct와 Qwen3 8B Base를 초기화로 쓰고, Dolci-Instruct-SFT 데이터로 5–15 에포크(9–27B 토큰)만 학습한다. 이는 비슷한 규모의 확산 모델(Dream, LLaDA, Efficient-DLM)이 보고한 최대 500B 토큰 대비 최소 18배 적다(§3, Table training-hyperparameters). 벤치마크는 지시 따르기·추론·코딩을 아우르는 IFEval, BBH, MMLU-Pro, GPQA, GSM8K, HumanEval 여섯 가지다.

flow 학습이 자기회귀를 해치지 않는다

여섯 벤치마크 평균 단일 토큰 디코딩 정확도를 비교하면(Table ntp_main), Lafa는 모든 규모에서 같은 초기화·데이터·학습 예산으로 만든 AR SFT를 넘는다. 8B에서는 Efficient-DLM 63.5%에 근접하면서 학습 토큰은 최소 18배 적다. flow를 배우는 과정이 오히려 자기회귀 예측을 개선한 셈이다(Table 1, §3.2).

병렬성: TPF가 오르면 정확도는 어디까지 유지되는가

TPF(tokens per forward)는 한 번의 forward pass로 확정되는 토큰 수다. AR은 정의상 1이다. self-speculative의 출력은 탐욕적 자기회귀와 같으므로, TPF 비교는 같은 AR 품질 위에서의 병렬성 비교가 된다.

위 그림(원문 Figure selfspec_frontier)의 가로축은 TPF, 세로축은 여섯 벤치마크 평균 정확도다. 7B에서 블록 4는 TPF 2.6·정확도 60.9%로 DeepSeek-MTP를 정확도와 TPF 모두 넘고, 블록 16은 TPF 4.9에서 정확도 58.0%로 AR SFT 58.4%에 근접하면서 병렬성을 거의 두 배로 올린다. 전 규모에서 약 2.5–5 TPF 구간에서 AR 수준 정확도를 유지한다(Table selfspec_lafa_sixbench).

실제 시간: TPS로 측정한 처리량

TPF는 병렬성만 보여줄 뿐 실제 시간을 말해주지 않는다. forward pass 비용이 방법마다 다르기 때문에, H100 GPU 1장·배치 1에서 초당 토큰 수(TPS)를 측정했다.

위 그림(원문 Figure throughput)의 가로축은 TPS, 세로축은 IFEval·GSM8K·HumanEval 평균 정확도다. Lafa 7B는 블록 16에서 100.2 TPS로 AR SFT의 39.5 TPS 대비 2.54배의 처리량을 비슷한 평균 정확도(74.5% vs 73.9%)에서 낸다. 블록 4에서는 79.9 TPS로 약 2배 빠르면서 평균 정확도(76.5%)가 더 높다. DeepSeek-MTP(53.6 TPS)와 MARS(48–49 TPS)는 AR보다 정확도가 같거나 높아도 최대 1.4배에 그친다. 확산 기반 방법(Fast-dLLM v2, Dream, Efficient-DLM)은 설정에 따라 더 높은 TPS를 내기도 하지만 평균 정확도가 AR SFT 아래이며, 빠른 설정일수록 정확도 손실이 크다(Table throughput_full, §3.3).

‘품질 유지’의 의미

self-speculative의 출력은 탐욕적 자기회귀와 같으므로, 2.54배는 같은 AR 품질을 유지한 채 얻는 속도 향상이다. flow 디코딩 자체는 다중 스텝이 필요하다. 7B·블록 4에서 32스텝 평균은 75.2%로 이미 AR SFT(73.9%)를 넘고, 64스텝에서는 75.9%에 이른다(Table flow_full, §3.1). 여러 스텝이 필요한 flow를 1스텝 초안 + 검증으로 바꾸는 것이 실질적인 가속의 원천이다. flow map 증류는 1–2스텝 구간의 정확도를 개선하며, 7B에서 2스텝 학생이 4스텝 교사에 근접한다.

한계와 적용 판단

이 논문이 스스로 남긴 한계는 주로 flow map 증류와 관련 있다. 증류된 1·2스텝 모델은 스텝 수를 줄이면서 정확도를 올리지만, 더 많은 스텝을 밟는 flow 디코딩과의 품질 격차가 남는다. 자기 조건화(self-conditioning)와 prefix guidance를 증류에 충분히 녹이는 것은 후속 과제로 남겨두었다(§3.1, §4). confidence 디코딩은 임계값을 낮추면 정확도가 떨어져 기본 방식으로 쓰지 못했다(Appendix confidence_results).

리뷰어 관점에서 주의할 점도 있다. 첫째, 2.54배는 7B·블록 16·H100 1장·배치 1 조건에서 측정한 것으로, 더 큰 모델이나 배치·긴 문맥에서 같은 배수가 나온다는 보장은 없다. 둘째, 평가가 지시·수학·코딩 벤치마크로 한정되어 있어 다른 영역 일반화는 원문에서 확인되지 않는다. 셋째, 8B는 Instruct 대신 Base에서 시작했다(논문은 Instruct 대응물의 추가 학습이 성능을 떨어뜨렸다고 보고한다). 이 때문에 8B 스케일의 IFEval 수치가 전반적으로 낮게 나온다(AR SFT도 38.1%).

이 조건을 감안하면, Lafa는 자기회귀 모델의 탐욕적 출력을 그대로 유지하면서 처리량을 높이고 싶은 서빙 환경에 적합하다. 이미 학습된 모델이 있고 토큰을 조금만 추가 학습할 수 있다면, 초안 모델을 따로 두지 않고도 speculative decoding의 이점을 얻을 수 있다. 반면 출력 분포 자체를 바꾸고 싶거나, 수 스텝의 flow 디코딩을 실시간으로 돌릴 수 있을 만큼 지연 여유가 큰 경우에는 이 방법의 강점이 덜 부각된다.

이 글의 그림은 arXiv:2610.28559 원본에서 가져왔습니다 (CC BY 4.0). 크기와 형식만 바꿨습니다.