미니배치 최적수송과 퍼플렉시티 상한을 이용한 이산 흐름 매칭

초록

이 논문은 이산 흐름 매칭(Discrete Flow Matching)에서 상태 전이 횟수를 최소화하기 위해 동적 최적수송 목표를 도입하고, 이를 미니배치 방식으로 최적화한다. 또한 확률 추정이 어려운 이산 흐름에 대해 두 가지 퍼플렉시티 상한을 제시해 훈련 및 평가를 이론적으로 정당화한다. 마지막으로 다중 마스크 흐름(Multimask Flow)을 제안해 마스크 기반 모델보다 퍼플렉시티와 다양성 모두에서 우수한 성능을 보이며, 1024 단계에서 32 단계로 추론 속도를 32배 가속한다.

상세 분석

본 연구는 연속 흐름 매칭에서 사용되는 정류(rectification) 전략이 이산 흐름에서는 확률 경로가 확률적(stochastic)이라 적용되지 못한다는 근본적인 한계를 지적한다. 이를 해결하기 위해 저자들은 “동적‑최적수송(dynamic‑optimal‑transport)” 형태의 목적함수를 정의한다. 이 목적은 각 시점 t에서 상태 i가 현재 토큰 x_i^t와 다른 토큰으로 전이될 확률 u_{i t}(x_i, x_t)를 토큰 간 유사도 s(x_i, x_i^t)와 곱해 누적함으로써, 전이 비용이 높은(즉, 유사도가 낮은) 이동을 억제한다. 수식 (11)에서 제시된 이 함수는 전통적인 베나무–브레니어(Benamou‑Brenier) 정리와 유사하게 동적 형태와 정적(Kantorovich) 형태가 동일한 최적값을 갖는다는 정리를 통해 정당화된다(정리 3.1). 특히 s를 해밍 거리(1‑δ)로 두면 비용함수는 두 시퀀스 간 해밍 거리와 동일해, 전이 횟수를 직접 최소화하는 효과를 얻는다. 토큰 임베딩의 L2 거리로 정의하면 연속 흐름과 동일한 2‑norm 비용을 재현한다.

이론적 기반 위에 저자들은 미니배치 최적수송을 적용한다. 전체 데이터셋이 아닌 작은 배치에 대해 비용을 추정하고, 배치 내 샘플 간 커플링 π를 최적화함으로써 대규모 어휘와 긴 시퀀스에서도 효율적인 학습이 가능하도록 설계했다. 이는 기존 마스크 흐름이 단일 마스크 토큰만을 사용해 커플링이 고정된 반면, 다중 마스크 흐름(Multimask Flow, DFM‑MMF)을 도입해 V_s개의 서로 다른 마스크 토큰을 활용한다는 점에서 차별화된다. 초기 시점 t=0에서는 마스크 토큰만이 존재하고, 각 위치에 ε=1/(V_s·L) 정도의 작은 질량을 할당해 “가상의 격자”를 만든 뒤, 최적수송을 통해 이 질량을 실제 데이터 토큰으로 옮긴다. 이렇게 하면 마스크와 데이터 사이의 커플링이 자유롭게 형성되어, 미니배치 OT가 적용 가능해진다.

퍼플렉시티는 이산 모델에서 직접적인 확률밀도 추정이 어려워 평가 지표로 활용하기 힘들다. 저자들은 두 가지 상한을 제시한다. 첫 번째는 동적 목표식에서 기대 전이 비용을 이용해 로그‑우도에 대한 상한을 만든 것이고, 두 번째는 커플링 π의 엔트로피와 비용을 결합한 변형된 포괄적 경계이다. 이 경계들은 모델이 실제 퍼플렉시티보다 낮은 값을 달성할 수 없음을 보장하며, 훈련 손실로 직접 사용될 수 있다. 실험 결과, 이러한 상한을 이용한 훈련은 기존 교차 엔트로피 기반 훈련과 비교해 퍼플렉시티가 0.2~0.4 정도 개선되고, 특히 GPT‑2 규모의 모델을 OWT(OpenWebText) 데이터에 적용했을 때 1024 단계에서 32 단계로 추론 스텝을 감소시켜도 동일한 퍼플렉시티를 유지한다는 점을 확인했다.

요약하면, 이 논문은 (1) 이산 흐름 매칭에 동적 최적수송 프레임워크를 도입해 전이 비용을 직접 최소화, (2) 미니배치 OT를 활용해 대규모 어휘와 긴 시퀀스에서도 효율적인 학습을 가능하게 함, (3) 퍼플렉시티 상한을 이론적으로 도출해 훈련·평가의 정당성을 확보, (4) 다중 마스크 흐름을 설계해 기존 마스크 흐름보다 퍼플렉시티와 다양성 모두에서 우수한 성능을 달성했다는 점에서 의미가 크다.