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

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

왜 중요한가

핵심 기여

MDLM을 위한 모델 스케줄링 전략 제안

전체 생성 과정 중 일부 단계에서 대형 모델 대신 미리 학습된 소형 모델을 사용하는 아키텍처 불가지론적(Architecture-agnostic) 가속 방법을 연구했다. 별도의 추가 학습이나 증류(Distillation) 과정 없이 추론 시점에 즉시 적용 가능하다.

비균일한 단계별 중요도 발견

확산 궤적의 중간 단계가 모델 교체에 가장 민감하며, 초기와 후기 단계는 상대적으로 견고하다는 사실을 입증했다. 이는 기존 이미지 확산 모델에서 보고된 단조로운 중요도 변화와는 대조적인 결과이다.

샌드위치 스케줄링을 통한 효율성 달성

궤적의 양 끝단에 소형 모델을 배치하는 '샌드위치' 패턴을 통해 OpenWebText 데이터셋에서 생성 품질 저하를 최소화하면서 FLOPs를 최대 17%까지 절감했다.

핵심 아이디어 이해하기

기존의 자기회귀(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 블록 연산에 국한되어 있으며 전체 시스템의 병목인 출력 레이어 최적화는 별도의 과제로 남아있다.

실무 활용

추가 학습 비용 없이 기존에 학습된 다양한 크기의 MDLM 체크포인트가 있다면 즉시 추론 속도를 개선할 수 있는 실용적인 기법이다.

실시간 텍스트 생성이 필요한 서비스에서 MDLM의 높은 추론 지연 시간을 단축
제한된 컴퓨팅 자원(Edge Device 등)에서 대형 확산 모델의 품질을 최대한 유지하며 실행
다양한 크기의 모델을 보유한 환경에서 사용자 요구 성능에 맞춰 동적으로 추론 비용을 조절

코드 공개 여부: 비공개

키워드

MDLM(마스크 확산 언어 모델)Model Scheduling(모델 스케줄링)Denoising(노이즈 제거)Inference Optimization(추론 최적화)Transformer(트랜스포머)

용어 해설

마스크 확산 언어 모델(MDLM): — 텍스트의 일부를 마스킹하고 이를 점진적으로 복원(Denoising)하며 생성하는 언어 모델이다. 기존 자기회귀 모델과 달리 여러 토큰을 동시에 업데이트할 수 있지만, 반복적인 추론 단계가 필요하여 연산 비용이 높다.
노이즈 제거(Denoising): — 데이터에 추가된 노이즈를 단계적으로 제거하여 원래의 깨끗한 데이터를 복원하는 과정이다. 확산 모델의 핵심 생성 메커니즘으로, 텍스트 생성에서는 마스킹된 토큰을 예측하는 과정에 해당한다.
키-값 캐싱(KV Caching): — 자기회귀 모델에서 이전 단계의 연산 결과(Key, Value)를 저장해 재사용함으로써 추론 속도를 높이는 기법이다. MDLM은 매 단계 전체 시퀀스를 다시 처리해야 하므로 이 기법을 직접 적용하기 어렵다.
퍼플렉서티(Perplexity): — 언어 모델이 다음 토큰을 얼마나 잘 예측하는지 나타내는 지표로, 값이 낮을수록 모델의 예측 성능과 생성 품질이 우수함을 의미한다.
부동 소수점 연산량(FLOPs): — 모델이 실행될 때 필요한 총 연산 횟수를 나타내는 척도이다. 모델의 크기나 추론 단계 수에 비례하며, 하드웨어 자원 소모량을 측정하는 주요 지표로 쓰인다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Diffusion 언어 모델의 속도와 품질 사이 균형

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

Diffusion 언어 모델의 속도와 품질 사이 균형