MDM-Prime-v2: 이진 인코딩과 인덱스 셔플링을 통한 확산 언어 모델의 연산 최적 스케일링

기존의 확산 기반 언어 모델은 자동 회귀 모델(ARM)에 비해 연산 효율성이 떨어지는 한계가 있었습니다. 이 논문은 토큰을 비트 단위로 쪼개고 무작위로 섞는 기법을 통해 확산 모델의 학습 효율을 극대화하여, 동일한 연산 자원으로 훨씬 더 높은 성능을 내는 방법을 제시합니다.

기존 언어 모델은 다음 토큰을 하나씩 예측하는 자동 회귀(Autoregressive) 방식이 주류였으나, 최근에는 마스킹된 부분을 한꺼번에 복원하는 확산(Diffusion) 모델이 주목받고 있습니다. 하지만 확산 모델은 토큰 단위의 이산적 구조 때문에 학습 효율이 떨어지는 문제가 있었습니다. MDM-Prime은 토큰을 여러 개의 서브토큰으로 쪼개어 세밀하게 복원하는 방식을 제안했으나, 얼마나 작게 쪼개야 할지(ℓ)와 어떤 방식으로 쪼개야 효율적인지에 대한 이론적 근거가 부족했습니다. 본 논문은 수학적 분석을 통해 토큰을 비트 단위(이진수)로 쪼개는 것이 변분 하한을 가장 타이트하게 만든다는 것을 증명했습니다. 또한, 자주 쓰이는 BPE 토큰의 인덱스 구조가 서브토큰의 엔트로피를 낮춰 학습을 방해한다는 점을 발견하고, 인덱스를 무작위로 섞는 셔플링을 도입해 이 문제를 해결했습니다.

전체 접근 방식은 MDM-Prime 프레임워크를 확장하여 이진 인코딩과 인덱스 셔플링을 결합한 MDM-Prime-v2를 구축하는 것이다. 토큰 x₀를 서브토큰 시퀀스 y₀로 변환하는 서브토크나이저 fℓ를 정의한다. 이진 인코딩 메커니즘은 ℓ = ⌈log₂ V⌉로 설정하여 각 토큰을 비트 시퀀스로 변환한다. [어휘 사전 크기 V를 입력으로] → [로그 연산을 통해 최소 비트 수 ℓ을 결정하고] → [각 토큰 인덱스를 ℓ비트 이진수로 변환하여] → [변분 하한 Lvb가 최소화되는 최적의 입도(granularity)를 확보한다]. 인덱스 셔플링은 BPE 토큰의 편향된 분포를 해결한다. [기존 BPE 인덱스를 입력으로] → [무작위 순열 함수 fshuffle을 적용해 인덱스를 재배치한 뒤] → [이진 인코딩을 수행하여] → [서브토큰의 엔트로피를 이론적 최대치인 log b에 근접하게 높여 예측 정확도를 향상시킨다].

OpenWebText 벤치마크에서 연산 최적화 비교 시 MDM-Prime-v2는 7.77 Perplexity를 기록하여 ARM(12.99), MDM(18.94), MDM-Prime(13.41)을 크게 앞질렀다. 스케일링 분석 결과, MDM-Prime-v2는 동일한 손실값(Loss)을 달성하는 데 있어 ARM보다 약 21.8배 더 적은 연산량(FLOPs)을 필요로 하는 것으로 나타났다. 1.1B 파라미터 모델 확장 실험에서 SciQ, McTaco 등 상식 추론 태스크에서 기존 ARM 기반 모델(TinyLLaMA 등)보다 높은 제로샷 정확도를 달성했다.

MDM-Prime-v2는 Diffusion Transformer(DiT) 아키텍처를 채택하며, 서브토큰 임베딩을 결합(Concatenation)하여 트랜스포머 레이어에 입력한다. 이는 모델의 파라미터 수나 연산량(FLOPs)을 늘리지 않으면서도 토큰 내부의 세밀한 정보를 모델링할 수 있게 한다. 이론적으로 변분 하한(Variational Bound) Lvb가 토큰 입도 ℓ에 대해 단조 비증가(monotonically non-increasing)함을 증명했다. 특히 ℓ = ⌈log₂ V⌉일 때 하한이 가장 타이트해지며, 이는 각 토큰을 비트 단위로 표현하는 이진 인코딩이 최적임을 수학적으로 뒷받침한다. BPE 토큰화의 인덱스 할당 방식이 서브토큰의 엔트로피를 낮추는 문제를 해결하기 위해 인덱스 셔플링(Index Shuffling)을 도입했다. 셔플링을 통해 서브토큰이 균등 분포(Uniform Distribution)에 가까워지도록 유도하며, 이는 예측의 확실성을 높여 손실 함수 값을 유의미하게 낮춘다. Chinchilla 스케일링 법칙을 적용한 분석에서 MDM-Prime-v2는 ARM보다 더 큰 데이터셋(D)을 효율적으로 학습할 수 있는 특성을 보였다.

토큰 간의 조건부 독립성 가정으로 인해 마스킹 비율이 매우 높은 초기 단계에서의 결합 확률 모델링 능력이 다소 저하되는 한계가 존재한다.