모든 노이즈 제거 단계가 동일하지는 않다: 더 빠른 마스크 확산 언어 모델을 위한 모델 스케줄링

마스크 확산 언어 모델(MDLM)은 생성 품질은 뛰어나지만 반복적인 추론 단계로 인해 연산 비용이 매우 높다는 단점이 있다. 이 논문은 생성 과정의 중간 단계가 가장 중요하다는 사실을 발견하고, 중요도가 낮은 초기와 후기 단계에 작은 모델을 배치하여 품질 저하 없이 속도를 높이는 실용적인 방법을 제시한다.

기존의 자기회귀(Autoregressive) 모델은 이전 토큰들을 바탕으로 다음 토큰 하나를 예측하며, KV Caching을 통해 연산 효율을 높인다. 반면 마스크 확산 언어 모델(MDLM)은 전체 문장의 마스킹된 부분을 여러 단계에 걸쳐 동시에 복원(Denoising)하는데, 매 단계마다 전체 시퀀스에 대해 Transformer 연산을 수행해야 하므로 연산량이 매우 많다.

이 논문은 모든 복원 단계가 동일한 연산 능력을 필요로 하는지 의문을 제기한다. 확산 과정은 완전한 노이즈 상태에서 시작해 점진적으로 깨끗한 텍스트로 변해가는 과정이다. 연구팀은 대형 모델과 소형 모델의 예측 차이를 분석한 결과, 생성의 시작(완전 노이즈)과 끝(거의 복원됨) 단계에서는 소형 모델도 대형 모델과 유사한 성능을 내지만, 문장의 구조가 결정되는 중간 단계에서는 대형 모델의 복잡한 추론 능력이 필수적임을 발견했다.

결과적으로 생성 과정의 양 끝단에는 가벼운 모델을 쓰고, 가장 민감한 중간 단계에만 무거운 모델을 집중 배치하는 '샌드위치' 구조를 통해 전체적인 연산 비용을 줄이면서도 최종 생성물의 품질은 대형 모델만 사용했을 때와 유사하게 유지할 수 있다.

대형 모델(Heavy, 12개 블록)과 소형 모델(Light, 4~10개 블록)을 준비하고, 전체 T 단계의 복원 과정 중 특정 구간에서 어떤 모델을 사용할지 결정하는 스케줄링 함수 s(i)를 정의한다. 별도의 재학습 없이 동일한 노이즈 스케줄로 학습된 서로 다른 크기의 모델들을 추론 시점에 교체하며 사용한다.

단계별 중요도를 분석하기 위해 두 가지 정량적 지표를 사용한다. 첫째는 마스크 토큰에 대한 Cross-Entropy Loss 차이인 Δloss(t)이다. 동일한 노이즈 입력 zt에 대해 대형 모델과 소형 모델의 손실값 차이를 계산하여 값이 클수록 해당 단계에서 대형 모델의 필요성이 높음을 의미한다. 둘째는 토큰 분포 간의 KL Divergence인 ΔKL(t)이다. 두 모델이 예측한 확률 분포의 거리인 Σ p_H log(p_H / p_L)를 계산하여 모델 간의 불일치 정도를 측정한다.

최적의 스케줄을 찾기 위해 전체 1000단계를 10개의 구간으로 나누고, 소형 모델을 배치할 4개의 구간을 선택하는 210가지 조합에 대해 전수 조사를 수행했다. 이를 통해 각 구간이 교체되었을 때 전체 생성 품질(Perplexity)에 미치는 영향력을 수치화하여 '샌드위치' 패턴의 우월성을 검증했다.

OpenWebText 데이터셋 실험에서 12블록 모델의 단계 중 25%를 4블록 모델로 교체했을 때, 샌드위치 스케줄(양 끝단 교체)은 Perplexity 44.31을 기록하여 전체 대형 모델(42.85) 대비 품질 저하를 최소화하면서 FLOPs를 16.7% 절감했다. 반면 중간 단계를 교체했을 때는 Perplexity가 47.5 이상으로 급격히 악화되었다.

LM1B 데이터셋에서도 동일한 경향성이 확인되어 이 현상이 특정 데이터셋에 국한되지 않음을 입증했다. 소형 모델의 크기를 4블록에서 10블록까지 변화시켰을 때, 모델이 커질수록 품질 저하는 줄어들지만 연산 절감 효과도 감소하는 연속적인 Speed-Quality Trade-off 관계를 보여주었다.

실제 실행 시간(Wall-clock time) 측면에서는 출력 레이어(Vocabulary Projection)의 고정 비용으로 인해 FLOPs 절감분만큼의 속도 향상은 나타나지 않았으나, 100% 소형 모델 사용 시 최대 28.3%의 속도 향상을 기록했다. 이는 추후 Liger-Kernel 등 최적화된 커널을 적용할 경우 더 개선될 여지가 있다.

MDLM의 복원 궤적에서 단계별 민감도가 '피크(Peaked)' 형태를 띠는 것은 연속적인 이미지 확산 모델의 '단조로운(Monotonic)' 경향성과 구별되는 특징이다. 이는 텍스트의 이산적 특성과 마스킹 메커니즘으로 인해 특정 노이즈 수준(t ≈ 0.4~0.6)에서 문맥적 결정이 집중되기 때문으로 분석된다.

실험에 사용된 모델은 768 히든 사이즈와 4 MLP 비율을 가진 Transformer-Encoder 구조이며, 4/6/8/10/12 블록 깊이만 다르게 구성되었다. 모든 모델은 동일한 GPT-2 토크나이저와 OpenWebText(1M 스텝) 또는 LM1B 데이터셋으로 학습되었다.

성능 평가 시 사용된 Generative Perplexity는 사전 학습된 GPT-2 모델을 통해 측정되었으며, 샘플의 다양성을 확인하기 위해 토큰 레벨 엔트로피를 병행 측정하여 모델 교체 시에도 생성 결과의 다양성이 훼손되지 않음을 확인했다.

현재 연구는 비교적 작은 규모의 모델과 데이터셋에서 수행되었으므로, Llama나 Qwen과 같은 초대형 모델 스케일에서도 동일한 중간 단계 민감도 패턴이 유지되는지 추가 검증이 필요하다. 또한, 현재의 속도 향상은 Transformer 블록 연산에 국한되어 있으며 전체 시스템의 병목인 출력 레이어 최적화는 별도의 과제로 남아있다.