💬 NLP

Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation

본 논문은 교사 궤적의 맹목적인 확률적 점프를 훈련 중 에너지 유도 탐색으로 대체하여, 소수 단계 학생 모델이 다단계 교사 및 더 큰 기준 모델보다 현저히 낮은 퍼플렉시티와 더 빠른 추론을 달성할 수 있게 하는 증류 방법인 궤적 모양 이산 흐름 매칭 (TS-DFM) 을 소개합니다.

원저자: Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade, Manuel R. Ciosici, Yizhe Zhang, Irina Belousova

게시일 2026-05-11

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade, Manuel R. Ciosici, Yizhe Zhang, Irina Belousova

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

한 명의 학생에게 완벽한 이야기를 쓰는 법을 가르치려 한다고 상상해 보세요.

옛날 방식: "맹목적인 점프" 교사
과거 연구자들은 **이산 흐름 매칭 (Discrete Flow Matching, DFM)**이라는 방법을 사용했습니다. 이 교사를 한 명의 학생을 빈 페이지에서 완성된 이야기까지 안내하려 하지만, 매우 기이한 방식으로 시도하는 사람으로 생각해 보세요.

교사는 "고양이가 매트 위에 앉았다"와 "보라색 바나나가 날아간다"가 섞인 무작위 뜬구름 잡는 말로 가득 찬 페이지로 시작합니다. 최종 이야기에 도달하기 위해 교사는 수백 번의 미세한 맹목적인 추측을 해야 합니다. 매 단계마다 그들은 결정해야 합니다: "이 단어를 바꿔야 할까? 그렇다면 무엇으로?"

문제는 교사가 결과를 보지 않은 채 이러한 결정을 내린다는 점입니다. 마치 어두운 숲을 걷다가 한 걸음을 내디디고 지뢰를 밟지 않았기를 바라는 것과 같습니다. 교사가 초기에 한 번의 나쁜 추측을 하면 (예를 들어 "고양이"를 "박쥐"로 바꾸는 대신 "고양이"로 유지해야 했을 때), 그 실수는 고정됩니다. 그 이후의 모든 단계는 그 실수 위에 쌓여야 합니다. 교사가 1,000 번째 단계를 마칠 때쯤이면 이야기는 오류로 가득 차 있지만, 학생은 그 messy하고 오류로 가득 찬 버전을 외우도록 강요받습니다. 그것이 그들에게 제공된 유일한 예시이기 때문입니다.

새로운 아이디어: "궤적이 교사가 되는 것"
이 논문의 저자들은 문제가 학생이 충분히 똑똑하지 않기 때문이 아니라, 교사의 경로가 고장 났기 때문이라고 주장합니다. 교사가 너무 많은 맹목적인 점프를 하고 있는 것입니다.

그들은 **TS-DFM (Trajectory-Shaped Discrete Flow Matching, 궤적 형태 이산 흐름 매칭)**이라는 새로운 방법을 제안합니다. 간단한 비유로 작동 방식을 설명해 보겠습니다:

나침반 비유

교사가 다시 어두운 숲을 걷고 있지만, 이번에는 마법 나침반(에너지 나침반이라고 함)을 가지고 있다고 상상해 보세요.

맹목적인 점프 (옛날 방식): 교사는 한 걸음을 내디디려다가 무작위 방향을 선택합니다.
안내된 점프 (새로운 방식): 한 걸음을 내디디기 전에 교사는 멈춥니다. 그들이 취할 수 있는 다섯 가지 가능한 경로를 상상합니다. 마법 나침반을 들어 올리면, 다섯 가지 경로 중 어떤 것이 가장 일관성 있고 고품질의 이야기로 이어지는지 즉시 알려줍니다.
- 경로 A: "고양이가 매트 위에 앉았다." (좋음)
- 경로 B: "박쥐가 매트 위에 앉았다." (괜찮지만 이상함)
- 경로 C: "고양이가 모자 위에 앉았다." (문맥이 틀림)
- ...이와 같습니다.

나침반은 "경로 A 로 가라"고 말합니다. 교사는 그런 다음 그 특정 고품질의 한 걸음을 내딛습니다.

두 단계 전략

이 논지는 이 나침반을 위한 교묘한 두 단계 과정을 설명합니다:

큰 그림 점검 (시퀀스 단계): 나침반은 전체 문장을 봅니다. 지금까지의 이야기 중 가장 좋은 전체 버전을 선택합니다. 이는 이야기가 뜬구름 잡는 말로 흘러가지 않도록 보장합니다.
미세 조정 점검 (토큰 단계): 전체 문장이 좋아 보일지라도, 특정 단어가 약간 어긋날 수 있습니다. 시스템은 그런 다음 개별 단어에 대한 교사의 "직감"(수학적 확신)을 점검합니다. 교사가 단어가 "앉았다 (sat)"여야 한다고 매우 확신하는데 경로가 "앉다 (sit)"를 선택했다면, 시스템은 조용히 그것을 다시 바꿔줍니다.

중요하게도, 이 나침반은 교사가 훈련 중일 때만 사용됩니다. 학생이 교훈을 배우면 나침반은 버려집니다. 학생은 나중에 이야기를 쓸 때 나침반이 필요하지 않습니다. 그들이 필요로 하는 것은 교사가 보여준 올바른 경로를 기억하는 것뿐입니다.

결과: 더 빠르고 더 똑똑함

이 논문은 1 억 7 천만 파라미터 모델 (중간 크기 AI) 로 이를 테스트했습니다.

속도: 옛날 방식은 문장을 쓰기 위해 1,024 단계가 필요했습니다. 새로운 방식은 단 8 단계로 이를 수행합니다. 이는 128 배 더 빠릅니다.
품질: 새로운 방식이 128 배 더 빠르지만, 그것이 쓰는 이야기들은 실제로 더 좋습니다(낮은 "퍼플렉시티", 즉 덜 혼란스럽고 더 일관성 있음). 이는 느리고 1,024 단계가 필요한 교사보다 낫습니다.
"병목 현상": 이 논문은 이러한 AI 모델이 얼마나 좋아질 수 있는지의 한계가 학생의 뇌 크기가 아니라, 교사가 보여주는 경로의 품질에 있음을 증명합니다. 경로를 고치면 결과도 고쳐집니다.

요약

운전 배우기를 생각해 보세요.

옛날 방식: 도로를 확인하지 않고 왼쪽, 오른쪽, 직진 중 하나를 선택하라고 말하는 운전 강사. 천 번의 작은 무작위 회전을 합니다. 당신은 운전을 배우지만, 결국 도랑에 빠집니다.
새로운 방식 (TS-DFM): 강사에게 GPS(나침반) 가 있습니다. 모든 회전 전에 다섯 가지 가능한 경로를 확인하고, 가장 안전한 것을 선택하여 그곳으로 안내합니다. 당신은 단 몇 번의 회전으로 완벽한 경로를 배웁니다.

이 논문은 훈련 중에 최상의 경로를 선택할 수 있도록 교사에게 "나침반"을 제공함으로써, 품질을 잃지 않고 AI 가 텍스트를 놀랍도록 빠르게 쓰도록 가르칠 수 있음을 보여줍니다. 이는 무작위 텍스트와 "마스크"(숨겨진 단어) 로 시작하는 텍스트 모두에 작동하며, 이전 방법들이 잘 처리하지 못했던 문제를 해결합니다.

기술 요약: 궤적이 교사이다: 에너지 항법 증류에 의한 소수 단계 이산 흐름 매칭

문제 제기

이산 흐름 매칭 (Discrete Flow Matching, DFM) 모델은 소스 시퀀스 (예: 균일 무작위 토큰 또는 마스킹 토큰) 를 일관된 언어로 반복적으로 변환하여 텍스트를 생성합니다. 품질은 자기회귀형 대규모 언어 모델 (LLM) 과 일치하지만, 이 과정은 계산적으로 매우 부담스러워 생성당 수백 또는 수천 번의 순전파 (정제 단계) 가 필요한 경우가 많습니다.

소수 단계 이산 흐름 매칭 (Few-Step Discrete Flow Matching, FS-DFM) 과 같은 증류 기술은 학생 모델이 반복적인 교사의 출력을 훨씬 적은 단계로 재현하도록 훈련함으로써 이 과정을 압축하려고 시도합니다. 그러나 증류된 학생 모델은 충분한 용량을 갖추었음에도 불구하고 종종 교사 모델보다 성능이 떨어집니다. 표준적인 설명은 이를 모델 용량의 한계로 돌립니다.

본 논문은 병목 현상이 학생 모델의 용량이 아니라 훈련 궤적의 품질에 있다고 주장합니다. 표준 DFM 증류 (특히 RK-4 추정기를 사용하는 경우) 에서 중간 상태는 "맹목적인 확률적 점프"를 통해 구성됩니다. 궤적의 각 중간 지점에서 속도 로지트 (velocity logits) 는 시퀀스 품질에 대한 평가 없이 확률적 과정을 통해 구체적인 토큰으로 변환됩니다. 초기 중간 지점에서 발생한 단일한 나쁜 확률적 결정은 후속 단계를 거쳐 전파되어 속도 평가와 최종 목표를 훼손합니다. 학생 모델은 이러한 "이동된 (drifted)" 목표를 충실히 모방하여 최적의 경로가 아닌 누적된 오류를 학습하게 됩니다.

방법론: 궤적 형태 이산 흐름 매칭 (TS-DFM)

저자들은 훈련 단계 동안 맹목적인 확률적 점프를 **안내된 항법 (guided navigation)**으로 대체하는 TS-DFM 프레임워크를 제안합니다. 이러한 "궤적 형성 (navigation shaping)"은 훈련 중에만 발생하며, 추론은 변경되지 않습니다.

1. 에너지 나침반

부분적으로 드러난 흐름 상태의 품질을 평가하기 위한 경량 스칼라 에너지 모델 ( $E_\phi$ ) 이 도입되었습니다.

아키텍처: FineWeb-Edu 로 훈련된 6 블록 트랜스포머 (90M 파라미터).
입력: 시간 조건부 없이 원시 토큰 시퀀스에 점수를 매겨 내재적인 시퀀스 품질을 평가합니다.
훈련: 나침반은 "생성 인식 부정 (generation-aware negatives)"을 사용하는 노이즈 대비 목적함수로 훈련됩니다. 이러한 부정 샘플은 흐름을 따라 뒤로 한 걸음씩 이동하여 덜 드러난 상태를 만든 다음, 고정된 교사 모델을 앞으로 실행하여 이를 재구성함으로써 생성됩니다. 이는 나침반이 실제 흐름 상태와 구별해야 하는 그럴듯하지만 손상된 시퀀스를 생성합니다.
성능: 나침반은 생성 관련 손상과 실제 흐름 상태를 구분하는 데 98.5%~99.8% 의 정확도를 달성합니다.

2. 시퀀스 - 토큰 항법

모든 토큰 위치를 개별적으로 평가하는 비실현 가능한 비용을 피하기 위해 TS-DFM 은 각 RK-4 중간 지점에서 2 단계 정책을 적용합니다.

1 단계 (시퀀스 수준): 현재 상태로부터 온도 다양성을 가진 CTMC 전이를 사용하여 $K$ 개의 다양한 후보 연속체를 생성합니다. 에너지 나침반은 모든 후보를 단일 배치 순전파로 점수화하고, 가장 낮은 에너지를 가진 후보를 선택합니다. 이를 통해 궤적이 전역적으로 낮은 에너지 영역으로 진입하도록 보장합니다.
2 단계 (토큰 수준): 선택된 후보는 속도 네트워크 자체의 신뢰도를 품질의 대리 지표로 사용하여 토큰 수준에서 정제됩니다. 후보가 속도의 고신뢰도 예측과 일치하지 않는 위치는 교체 대상이 됩니다. 이 단계는 추가적인 에너지 모델 추론 비용을 발생시키지 않습니다.
에너지 안전 장치: 최종 에너지 점검은 정제된 시퀀스가 작은 허용 오차 ( $\epsilon_{safe}$ ) 를 초과하여 에너지를 증가시키지 않도록 보장하여, 국소적 정제가 전역 품질을 저하시키지 않도록 합니다.

3. 시간 임계값 활성화

항법은 시간 임계값 $\tau$ 에 의해 게이트됩니다. 초기 흐름 시간 ( $t < \tau$ ) 에는 시퀀스가 소스 노이즈에 지배되어 에너지 기반 구분이 비효율적입니다. 항법은 충분한 언어 구조가 나타났을 때 ( $t \ge \tau$ ) 에만 활성화되어 나침반이 신호가 풍부한 상태에서 작동하도록 보장합니다.

주요 기여

병목 현상의 재정의: 본 논문은 소수 단계 DFM 증류에서 학생 모델의 용량이 아닌 궤적 품질이 주요 제한 요소임을 규명합니다. 맹목적인 확률적 점프에서 발생하는 누적 오류가 감독 신호를 저하시킨다는 것을 입증합니다.
궤적 형성: 에너지 나침반을 통한 안내 선택으로 맹목적인 점프를 대체하여 훈련 궤적을 형성하는 새로운 프레임워크입니다.
시퀀스 - 토큰 정책: 전역 시퀀스 선택과 국소 토큰 정제를 결합하여 토큰별 에너지 평가의 필요성을 피하는 계산 효율적인 2 단계 접근법입니다.
훈련 전용 오버헤드: 모든 형성 메커니즘은 훈련 중에만 적용됩니다. 학생 모델의 추론 절차, 아키텍처, 속도는 표준 증류 모델과 동일하게 유지됩니다.

실험 결과

170M 파라미터 및 1.3B 파라미터 언어 모델을 사용하여 균일 및 마스킹 소스 분포로 실험을 수행했습니다.

교사 대비 성능: 균일 소스를 사용한 170M 모델에서 8 단계 TS-DFM 학생 모델은 GPT-2 퍼플렉시티 56.1을 달성했습니다. 이는 1,024 단계 교사 (82.8 PPL) 대비 32% 감소이며, FS-DFM 베이스라인 (87.6 PPL) 대비 36% 감소입니다.
속도: 8 단계 TS-DFM 모델은 더 나은 퍼플렉시티를 달성하면서 1,024 단계 교사보다 128 배 빠릅니다.
마스킹 소스: 개선 효과는 마스킹 소스에서 더욱 두드러져, 8 단계에서 TS-DFM 이 FS-DFM 베이스라인 대비 5.6 배 개선을 달성했습니다.
최첨단 성능: TS-DFM 은 일부 경쟁사보다 작은 모델과 더 적은 훈련 데이터를 사용함에도 불구하고 모든 평가자 (GPT-2, LLaMA-2, LLaMA-3) 에서 기존 소수 단계 이산 생성 방법 (Duo 및 SDTT 포함) 을 능가합니다.
확장성: 1.3B 파라미터에서 TS-DFM 과 FS-DFM 간의 격차는 확대되어 (8 단계에서 41% PPL 감소), 형성된 학생 모델은 단 8 단계에서 1,024 단계 교사를 능가합니다.
추론 작업: 이 방법은 1.3B 규모에서 수학 추론 (GSM8K) 으로 전이되어, 4 단계 및 8 단계 TS-DFM 학생 모델이 전체 단계 교사의 정확도를 초과합니다.

중요성 및 주장

본 논문은 궤적 품질이 소수 단계 이산 흐름 매칭의 근본적인 병목 현상이라고 주장합니다. 중간 상태를 확률적 노이즈에 의해 고정된 것으로 받아들이는 대신 품질 최적화의 대상으로 간주함으로써 TS-DFM 은 현재 증류 방법의 성능 한계를 돌파합니다.

저자들은 이 접근 방식이 **증류 메커니즘 자체와 직교 (orthogonal)**하다고 강조합니다. 이는 궤적 기반 증류 프레임워크 위에 적용될 수 있습니다. 이 방법은 추론 비용을 증가시키지 않고 오직 일회성 훈련 오버헤드 (표준 FS-DFM 비용의 약 2.2 배) 에만 의존하여 최첨단 퍼플렉시티를 달성합니다. 결과는 이산 생성의 경우 "교사" (궤적 구성 과정) 가 종종 제한 요소이며, 학생 용량을 단순히 증가시키는 것보다 교사의 경로를 개선하는 것이 더 큰 효과를 가져온다는 것을 시사합니다.