기술적 요약: 범주형 흐름 지도의 확장

문제 정의

자기회귀 (AR) 모델이 언어 모델링을 지배하는 반면, 연속 확산 및 흐름 매칭 모델은 가속화된 샘플링 및 분포 기울기 조정과 같은 연속 모달리티에서 상당한 이점을 입증했습니다. 이러한 이점을 이산 데이터로 확장하는 것이 최근 연구의 주요 목표였습니다. 이산 확산과 같은 이전 접근 방식은 이산 격자에서 작동하여 확립된 연속 기법의 적용을 제한합니다. 반면, 연속 상태 프로세스 (예: 범주형 흐름 지도 또는 CFM) 는 연속 심플렉스 또는 $\mathbb{R}^d$ 위의 확률 분포를 진화시켜 이산 데이터를 생성합니다.

최근 몇 단계 영역에서 경쟁력 있는 샘플 품질을 달성한 CFM 의 시연에도 불구하고, 그 확장성은 입증되지 않았습니다. CFM 이 현대의 대규모 언어 모델 (LLM) 에 필요한 수십억 파라미터 및 수조 토큰 영역으로 확장될 수 있는지 여부는 불분명했습니다. 주요 기술적 장벽은 계산 및 메모리 오버헤드입니다. 어휘 크기 $|V|$와 컨텍스트 길이 $L$의 경우, CFM 은 매 단계마다 $L \times |V|$ 크기의 로그 확률 행렬을 구체화해야 하며, $|V| \approx 200,000$이고 $L$이 클 때 이는 금지적일 수 있습니다.

방법론

저자들은 2.1 조 토큰으로 학습된 17 억 파라미터 모델로 CFM 을 확장하기 위한 2 단계 학습 파이프라인을 제안합니다.

1. 흐름 모델 사전 학습

첫 번째 단계는 확률 흐름의 드리프트를 근사하도록 표준 흐름 매칭 모델 (디노이저) 을 학습하는 것입니다.

• 아키텍처: Lipschitz 제약을 준수하여 안정성을 확보한 AdaLN-Zero 변조, 회전 위치 임베딩, SwiGLU MLP 를 갖춘 17 억 파라미터 트랜스포머.
• 목적: 모델은 예측된 종점과 실제 데이터 간의 교차 엔트로피 손실 (식 2.10) 을 최소화하도록 학습되어, 벡터장 $v_\theta$를 효과적으로 학습합니다.
• 디자인 공간 탐색: 확장 전에 저자들은 4 억 파라미터 프록시 모델에서 광범위한 하이퍼파라미터 스윕 (350 개 이상의 구성) 을 수행했습니다. 주요 발견 사항은 다음과 같습니다.
  • 시간 스케줄링: $\lambda \approx 0.75$인 선형 및 "오류 디코딩" 스케줄 ($\alpha_t = (1-\lambda)t + \lambda\gamma_t$) 의 혼합이 최적의 성능을 발휘했습니다. 순수 오류 디코딩 ($\lambda=1$) 은 성능 붕괴를 초래했습니다.
  • 적응형 손실 가중치: 예측 오류에 기반하여 손실을 재가중치 ($w(x) = \text{sg}((\Delta(x) + c)^{-r})$) 하는 것은 생성 퍼플렉시티 (Gen PPL) 를 개선했지만, 분포 충실도 (MAUVE*) 를 유지하기 위해 신중한 튜닝이 필요했습니다.
  • 랜덤 언마스킹: 일관된 조건부 텍스트 생성을 위해서는 임의의 접두사를 언마스크할 확률 $p > 0$로 학습하는 것이 필수적이었으며, 그렇지 않으면 모델은 분포 밖 출력을 생성했습니다.

2. 범주형 흐름 지도 (CFM) 로의 자기 증류

두 번째 단계는 사전 학습된 흐름 모델을 고품질, 소수 단계 생성이 가능한 CFM 으로 변환합니다.

• 메커니즘: 모델은 ODE 를 여러 단계에 걸쳐 적분하는 대신, 시간 $s$에서 $t$로 확률 질량을 한 단계로 운반하는 흐름 지도 $X_{s,t}$를 학습합니다.
• 손실 함수: 저자들은 점진적 자기 증류 (PSD) 와 안정화된 오일러 로그 자기 증류 (ESD) 를 평가했습니다.
  • PSD: 연속 시간 단축 모델 접근 방식을 사용합니다.
  • 발산 선택: 저자들은 Forward KL, Reverse KL, 그리고 Jensen-Shannon 발산 (JSD) 을 비교했습니다. Forward KL 과 JSD 는 소수 단계 영역에서 토큰 엔트로피와 MAUVE*를 개선한 반면, Reverse KL 은 불안정성과 모드 붕괴를 초래했습니다.
  • 안정성: 17 억 규모에서의 ESD 학습은 고분산 야코비안 추정 (유한 차분으로 계산됨) 으로 인해 불안정함이 입증되어, 저자들은 PSD 결과에 집중했습니다.
• 샘플링: 최종 CFM 은 사전 분포에서 샘플링하고 학습된 흐름 지도 $X_{0,1}$을 단일 패스 (또는 소수 패스) 에 적용하여 텍스트를 생성하므로 추론 시간이 크게 단축됩니다.

3. 가능도 추정

표준 가능도 기반 벤치마크 (예: MCQA) 에서 CFM 을 평가하기 위해 저자들은 준이산 ELBO를 유도했습니다. 이 경계는 고차원 이산 설정에서 특히 비용이 많이 드는 ODE 의 역방향 적분을 필요로 하지 않고도 답변 선택지를 점수화할 수 있게 합니다.

주요 기여

1. 확장성 입증: 저자들은 2.1 조 토큰으로 17 억 파라미터 흐름 모델을 성공적으로 학습하고 이를 CFM 으로 자기 증류하여, CFM 이 수십억 파라미터 영역으로 확장될 수 있음을 입증했습니다.
2. 하이퍼파라미터 발견: 4 억 프록시 모델에 대한 체계적인 스윕을 통해 시간 스케줄링, 손실 가중치, 입력 마스킹에 대한 견고한 설정을 식별했으며, 이는 더 큰 규모로 이전됩니다.
3. 소수 단계 생성: 자기 증류된 CFM 은 데이터 수준의 토큰 엔트로피를 유지하면서 최소 4 개의 추론 단계 (NFE) 에서 다양하고 고품질의 텍스트를 생성하며, 저 NFE 영역에서 비증류 흐름 모델을 능가합니다.
4. 가능도 경계: 준이산 ELBO 의 도입은 표준 언어 모델링 벤치마크에서 CFM 을 평가할 수 있게 하여, 다른 비자기회귀 방법과 유사한 성능 범위에 위치시킵니다.
5. 한계 분석: 이 논문은 $L \times |V|$ 행렬을 구체화하는 메모리 병목 현상과 시간 조건부 아키텍처에 대한 제로샷 하이퍼파라미터 이전의 어려움을 명시적으로 다룹니다.

결과

• 품질 vs 단계: 자기 증류된 CFM 은 1024 단계에서 ~58.66 의 Gen PPL 을 달성하여 (AR 기준의 45.33 과 경쟁력 있음) ~0.94 의 MAUVE* 점수를 기록합니다. 중요하게도, 단 4 단계에서도 CFM 은 ~0.88 의 MAUVE*를 달성하여 소수 단계 영역에서 AR 기준 (0.7) 을 크게 능가합니다.
• 토큰 엔트로피: 자기 증류는 낮은 NFE 에서도 높은 토큰 엔트로피 (다양성) 를 유지하는 반면, 비증류 모델은 1~2 단계 영역에서 붕괴 (낮은 엔트로피) 하는 경향이 있습니다.
• 벤치마크 성능: MCQA 벤치마크 (ARC-E, BoolQ, OBQA 등) 에서 CFM 은 다른 비자기회귀 방법 (예: Eso-LM) 과 유사한 정확도를 달성하지만, ELBO 최대화에 직접 학습된 이산 확산 방법보다는 낮습니다.
• 조건부 생성: 자기 증류된 모델은 비증류 모델에 비해 소수 단계 영역에서 $\Delta$PPL(조건부 퍼플렉시티) 에서 상당한 개선을 보입니다.

의미 및 주장

이 논문은 범주형 흐름 지도가 현대 LLM 규모에서 이산 데이터 생성을 위한 실현 가능한 경로라고 주장합니다. 17 억 파라미터로 성공적으로 확장함으로써 저자들은 "흐름 지도 도구상자"—가이드, 분포 기울기 조정, 보상 미세 조정 포함—가 이제 언어 모델링에 접근 가능함을 입증했습니다.

이 연구는 연속 상태 생성이 샘플 품질을 위해 연산을 교환할 수 있음을 시사하며, 전통적인 확산 또는 AR 모델보다 훨씬 적은 함수 평가를 통해 고품질 텍스트 생성을 가능하게 합니다. 저자들은 가능도 추정이 유도된 ELBO 가 경계이므로 개선이 필요한 열린 영역임을 지적하지만, 결과는 CFM 이 이산 확산 기준과 경쟁력 있는 성능을 달성하면서도 연속 흐름 매칭의 이점을 unlocked 함을 나타냅니다. 논문은 CFM 의 확장성이 단순한 그럴듯한 출력 생성을 넘어 언어 모델을 조종하고 미세 조정하는 새로운 길을 열어준다고 결론지었습니다.