Skip to the Good Part: Representation Structure & Inference-Time Layer Skipping in Diffusion vs. Autoregressive LLMs

이 논문은 확산 언어 모델이 자기회귀 모델과 달리 계층적 표현 구조와 초기 층의 중복성을 가지며, 이를 활용한 정적 계층 건너뛰기 기법으로 추가적인 구조 변경 없이 추론 시 연산량을 18.75%까지 줄이면서도 성능을 유지할 수 있음을 보여줍니다.

핵심

🎨 핵심 비유: "점화식 (AR)" vs "조각상 다듬기 (Diffusion)"

우리가 글을 쓸 때, 인공지능도 크게 두 가지 방식으로 글을 만듭니다.

1. 기존 방식 (Autoregressive, AR):

   • 비유: 한 글자씩 써가는 '점화식'.
   • 설명: "안녕"이라고 쓸 때, '안'을 먼저 쓰고, 그다음 '녕'을 생각합니다. 이전 단어가 다음 단어에 절대적인 영향을 미칩니다.
   • 특징: 한 번 실수하면 바로잡기 어렵고, 매 단어를 쓸 때마다 뇌 (레이어) 전체가 다시 긴장하며 새로운 정보를 추가합니다. 그래서 중간 과정을 생략하면 글이 완전히 망가집니다.

2. 새로운 방식 (Diffusion, dLLM):

   • 비유: 거친 조각상을 다듬는 '조각가'.
   • 설명: 처음엔 전체가 흐릿한 구름 (노이즈) 으로 되어 있습니다. 조각가는 전체를 한눈에 보며, "여기는 너무 거칠다", "저기는 모양이 맞지 않다"라고 전체를 보며 점진적으로 다듬어 나갑니다.
   • 특징: 처음엔 전체적인 윤곽 (대략적인 의미) 을 잡고, 나중엔 디테일을 다듬습니다. 처음 몇 단계는 전체적인 흐름만 잡기 때문에, 그 과정을 생략해도 최종 결과물이 크게 달라지지 않습니다.

---

🔍 이 논문이 발견한 놀라운 사실 3 가지

1. "처음에 시작하면 끝까지 따라가는 습관" (초기화 편향)

연구진은 **"기존 방식 (AR) 으로 훈련된 모델을 새로운 방식 (Diffusion) 으로 다시 훈련하면, 정말 새로운 사람이 될까?"**를 궁금해했습니다.

• 결과: 아니요! 습관은 쉽게 바뀌지 않습니다.
• 비유: 이미 '점화식'으로 글을 쓰는 습관이 들은 사람 (Dream-7B) 이 '조각가' 훈련을 받아도, 여전히 한 글자씩 꼼꼼히 확인하는 습관 (AR 성향) 을 버리지 못합니다.
• 의미: 모델의 '성격'은 처음에 어떻게 가르쳤는지 (초기화) 에 따라 결정되며, 나중에 훈련 방식을 바꿔도 그 성격이 쉽게 사라지지 않습니다.

2. "조각가에게는 '여분'이 있다!" (중복성 발견)

반면, 처음부터 '조각가' 방식으로 훈련된 모델 (LLaDA) 은 완전히 달랐습니다.

• 발견: 이 모델은 처음 몇 단계 (레이어) 에서 이미 전체적인 그림을 그립니다. 그다음 단계들은 그 그림을 조금 더 선명하게 할 뿐, 전혀 새로운 정보를 추가하지 않습니다.
• 비유: 100 단계의 작업을 하는데, 16 단계는 "대략적인 윤곽 잡기"만 하고, 7100 단계는 "디테일 다듬기"만 합니다. 만약 1~6 단계를 생략하고 7 단계부터 시작해도, 최종 결과물은 거의 비슷합니다.
• 결론: 새로운 방식 (Diffusion) 은 기존 방식보다 '여분의 작업 (중복성)'이 훨씬 많습니다.

3. "일하는 시간을 20% 줄여도 결과는 그대로!" (레이어 스킵)

이러한 '여분의 작업'을 이용해, 인공지능이 일하는 시간을 줄여보았습니다.

• 방법: "처음에 대략적인 그림만 그리는 단계 (중복된 레이어) 는 아예 건너뛰자!"라고 정했습니다.
• 결과:
  • 새로운 방식 (Diffusion): 일하는 양을 약 19% 줄였는데, 성능은 90% 이상 유지되었습니다. (완벽한 효율!)
  • 기존 방식 (AR): 일하는 양을 조금만 줄여도 (7% 만 줄여도) 성능이 반토막이 났습니다. (너무 취약함)

---

💡 이 연구가 우리에게 주는 메시지

1. 더 빠르고 저렴한 AI: 새로운 방식 (Diffusion) 으로 만든 AI 는 불필요한 계산을 많이 하므로, 이를 잘라내면 전기세와 시간을 아낄 수 있습니다.
2. 모델을 바꿀 때는 조심하자: 기존에 훈련된 모델을 새로운 방식으로 바꿀 때, 겉모습만 바뀌고 속성은 그대로일 수 있다는 점을 깨달았습니다.
3. 효율의 새로운 길: 메모리를 줄이는 기술 (KV Cache) 과는 별개로, 계산 과정 자체를 줄이는 새로운 방법을 제시했습니다.

🏁 한 줄 요약

"기존 AI 는 한 줄 한 줄 꼼꼼히 써야 하지만, 새로운 AI 는 처음에 대략적인 그림을 먼저 그리기 때문에, 그 '초기 작업'을 생략해도 결과가 거의 똑같습니다. 이 사실을 이용해 AI 의 일하는 시간을 20% 줄여도 성능은 그대로 유지할 수 있습니다!"

이 연구는 인공지능이 더 가볍고 빠르게 작동할 수 있는 새로운 길을 열어주었습니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 최근 확산 기반 언어 모델 (Diffusion Language Models, dLLMs) 은 자기회귀 (Autoregressive, AR) 모델과 유사한 성능을 달성하고 있습니다. 그러나 AR 모델은 왼쪽에서 오른쪽으로 토큰을 점진적으로 예측하여 표현을 구성하는 반면, dLLMs 는 전체 시퀀스를 반복적으로 탈노이즈 (denoising) 하는 방식으로 학습됩니다.
• 문제: 두 모델의 학습 목표 (Objective) 가 근본적으로 다르기 때문에, 내부 표현 (Internal Representations) 이 층 (Layer) 을 따라 어떻게 조직화되는지에 대한 이해가 부족합니다. 특히, dLLMs 의 학습 목표가 표현의 계층적 추상화 (Hierarchical Abstraction) 에 어떤 영향을 미치고, 이것이 추론 시간 (Inference-time) 효율성 향상 (예: 레이어 스킵) 으로 이어질 수 있는지 명확하지 않았습니다.
• 가설: 확산 (Diffusion) 학습 목표는 AR 모델과 다른 내부 기하학적 구조를 형성하며, 특히 초기 층에서의 표현 중복성 (Redundancy) 을 유발하여 추론 시 레이어를 생략 (Skipping) 하더라도 성능 저하가 적을 것이라고 가정했습니다.

2. 방법론 (Methodology)

저자들은 학습 목표와 초기화 (Initialization) 의 영향을 분리하기 위해 세 가지 모델 군을 비교 분석했습니다:

1. Native dLLM: LLaDA (확산 모델로 처음부터 학습됨).
2. Native AR Model: Qwen2.5-7B (자기회귀 모델).
3. AR-initialized dLLM: Dream-7B (Qwen2.5 로 초기화 후 확산 학습을 거친 모델).

A. 표현 분석 (Representational Analysis)

• 층별 및 토큰별 유사도 측정: 시퀀스 내 모든 토큰에 대해 인접한 층 ($h_\ell$과 $h_{\ell+1}$) 간의 코사인 유사도 (Cosine Similarity) 를 측정했습니다.
• 관찰 지표:
  • 계층적 추상화: 초기 층에서 표현이 얼마나 "거칠게 (Coarse)" 형성되는지.
  • 최근성 편향 (Recency Bias): 새로운 토큰이 들어올 때 표현이 얼마나 급격히 변하는지.
  • 초기화 편향: AR 로 초기화된 확산 모델이 확산 학습 후에도 AR 의 표현 패턴을 유지하는지 확인.

B. 추론 시간 레이어 스킵 (Inference-Time Layer Skipping)

• 고정된 스킵 정책 (Static, Task-agnostic): 학습 단계에서 분석된 유사도 패턴을 기반으로, 유사도가 높은 (중복된) 레이어들을 추론 시 정적으로 생략하는 정책을 도입했습니다.
• 구현 방식:
  • KV 캐시 공유나 아키텍처 변경 없이, 기존 사전 학습된 모델에 적용 가능합니다.
  • 연속된 레이어를 생략하지 않도록 하여 표현의 연속성을 유지하는 알고리즘을 사용했습니다.
  • Algorithm 1: 유사도가 임계값 ($\theta=0.95$) 이상인 레이어를 선택하되, 인접한 레이어가 이미 생략되지 않았는지 확인하여 생략 목록을 구성합니다.

3. 주요 기여 및 발견 (Key Contributions & Findings)

1) 학습 목표에 따른 표현 구조의 근본적 차이

• Native dLLM (LLaDA): 초기 층에서 매우 높은 유사도 (>0.95) 를 보이는 "중복된 표현" 영역이 존재하며, 이는 거친 표현을 형성한 후 후기 층에서 정교한 세밀화 (Refinement) 를 수행하는 계층적 추상화 (Coarse-to-fine) 구조를 가집니다. 또한 최근성 편향 (Recency Bias) 이 최소화되어 전역적 (Global) 인 표현을 보입니다.
• Native AR Model (Qwen2.5): 모든 층에서 토큰별로 표현이 점진적으로 세밀하게 업데이트되며, 강한 최근성 편향을 보입니다. 층 간 표현이 밀접하게 결합되어 있어 중복성이 낮습니다.
• AR-초기화 dLLM (Dream-7B): 확산 학습을 받았음에도 불구하고, AR 초기화 (Qwen2.5) 의 표현 패턴을 강력하게 유지합니다. 이는 초기화가 학습 목표보다 표현 구조에 더 큰 영향을 미친다는 것을 보여줍니다.

2) 레이어 스킵을 통한 효율성 증대

• Native dLLM: 초기 층의 높은 중복성을 활용하여 6 개의 레이어 (FLOPs 18.75% 감소) 를 생략하더라도, 추론 성능 (Reasoning, Code Generation) 을 88% 이상 유지했습니다.
• AR 모델 및 AR-초기화 모델: 동일한 조건 (2 개 레이어 생략, FLOPs 7.14% 감소) 에서 성능이 급격히 저하되었습니다 (Qwen2.5 는 3575% 수준, Dream-7B 는 6080% 수준). 이는 AR 모델의 표현이 층 간에 밀접하게 의존하고 있음을 시사합니다.

4. 실험 결과 (Results)

• 벤치마크: GSM8K (수학), MATH-500, HumanEval, MBPP (코드 생성) 등 다양한 과제를 평가했습니다.
• 성능 유지율:
  • LLaDA: 6 레이어 스킵 시 모든 태스크에서 88%~102% 의 성능 유지.
  • Qwen2.5: 2 레이어 스킵 시 35%~75% 로 급감.
  • Dream-7B: 2 레이어 스킵 시 60%~81% 로 저하 (LLaDA 보다는 낮고 Qwen2.5 보다는 높으나, 여전히 AR 특성을 보임).
• 계산 효율성: Native dLLM 은 AR 모델 대비 2.6 배 더 많은 FLOPs 절감을 달성하면서도 1.4 배 더 높은 품질 유지를 보였습니다.
• KV 캐시와의 관계: 레이어 스킵은 깊이 (Depth) 기반 계산을 줄이는 반면, KV 캐시는 토큰 기반 중복을 줄이므로 두 기법은 상호 보완적 (Orthogonal) 입니다.

5. 의의 및 결론 (Significance)

1. 학습 목표와 표현 구조의 연결: 확산 학습 목표가 모델 내부의 표현을 계층적이고 중복성 있게 재구성한다는 것을 최초로 체계적으로 증명했습니다.
2. 초기화의 지속성: AR 모델로 초기화 후 확산 학습을 거친 모델 (Dream-7B) 이 여전히 AR 의 표현 특성을 유지한다는 것은, 모델 적응 (Adaptation) 시 초기화 편향이 학습 목표보다 더 강력할 수 있음을 경고합니다.
3. 실용적 효율성: 아키텍처 변경 없이, 오직 학습된 표현의 중복성만 활용하여 캐시 (KV-cache) 와 독립적인 추론 가속화를 가능하게 했습니다. 이는 확산 언어 모델의 실용성을 높이는 중요한 단계입니다.
4. 향후 과제: 동적 (Dynamic) 또는 입력 적응형 스킵 정책, KV 캐시와의 결합, 멀티모달 확산 아키텍처로의 확장 등을 제안합니다.

요약하자면, 이 논문은 확산 기반 언어 모델이 AR 모델과 다른 내부 표현 구조 (초기 층의 높은 중복성, 계층적 추상화) 를 가지며, 이를 활용하면 별도의 구조 변경 없이도 추론 비용을 크게 줄일 수 있음을 증명했습니다. 반면, AR 모델로 초기화된 확산 모델은 이러한 이점을 충분히 누리지 못함을 보여주었습니다.