💬 NLP

Learnability-Informed Fine-Tuning of Diffusion Language Models

본 논문은 다양한 확산 시간 단계에서 토큰 학습 난이도를 사용 가능한 문맥과 동적으로 정렬하여 추론 벤치마크에서 기존 지도 미세조정 기준을 크게 능가하는 확산 언어 모델을 위한 학습 가능성 기반 미세조정 알고리즘인 LIFT 를 소개합니다.

원저자: Shubham Parashar, Atharv Chagi, Jacob Helwig, Lakshmi Jotsna, Sushil Vemuri, James Caverlee, Dileep Kalathil, Shuiwang Ji

게시일 2026-05-25

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Shubham Parashar, Atharv Chagi, Jacob Helwig, Lakshmi Jotsna, Sushil Vemuri, James Caverlee, Dileep Kalathil, Shuiwang Ji

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

이 논문은 간단한 언어와 창의적인 비유를 사용하여 설명합니다.

핵심 아이디어: AI 에게 더 잘 "학습"하도록 가르치기

복잡한 수학 문제를 푸는 법을 학생에게 가르치려 한다고 상상해 보세요. 이를 위한 두 가지 주요 방법이 있습니다:

오래된 방법 (자기회귀): 학생이 문제를 읽고 한 단어씩 답을 작성합니다. 마치 문장을 타이핑하는 것과 같습니다.
새로운 방법 (확산): 학생은 빈 칸으로 가득 찬 페이지 (마스킹된 토큰) 로 시작하여, 답이 명확해질 때까지 모든 빈칸을 한 번에 채우려고 시도하며 추측을 정제합니다. 이를 **확산 언어 모델 (DLM)**이라고 합니다.

이 논문의 연구자들은 "새로운 방법"이 빠르다는 점을 발견했지만, 표준적인 방법 ( 지도 미세 조정 또는 SFT) 으로 가르치려 할 때 막힌다는 점도 발견했습니다. 이는 90% 의 단어가 누락된 페이지를 학생에게 건네주되, "the"나 "and"와 같은 가장 흔한 단어만 추측하도록 요청하는 것과 같습니다. 학생은 지루해하고 새로운 것을 배우지 못합니다. 반대로, 페이지가 거의 비어있을 때 희귀하고 어려운 단어를 추측하도록 요청하면, 학생은 좌절하고 해내지 못합니다.

문제: "무엇"과 "언제"가 맞지 않음

저자 수백만 개의 학습 예제를 분석하여 이러한 모델이 학습하는 방식에 특정 불일치가 있음을 발견했습니다:

"무엇": 희귀하고 어려운 단어 (예: 특정 수학 용어) 는 배우기 어렵습니다. 흔한 단어는 쉽습니다.
"언제": 확산 과정에서 모델은 매우 혼란스럽고 대부분이 비어있는 페이지 (무엇도 배우기 어려움) 로 시작하여 점차 더 많은 문맥을 드러냅니다 (배우기 쉬움).

불일치:

과정 초기 (많은 문맥): 모델은 흔한 단어를 쉽게 추측할 수 있지만, 쉬운 것에 너무 바빠서 어렵고 희귀한 단어는 무시합니다.
과정 후기 (매우 적은 문맥): 모델은 대부분이 비어있는 페이지를 바라봅니다. 희귀한 단어를 추측하려 하지만, 이를 수행할 충분한 정보가 없어 실패합니다.

표준 학습 방법은 모든 단계에서 모든 것을 가르치려 하므로 비효율적입니다. 이는 학생에게 눈가리개를 하고 사전을 암기하게 한 뒤, 소음 차단 헤드폰을 끼고 미적분 문제를 풀게 하는 것과 같습니다.

해결책: LIFT (학습 가능성 기반 미세 조정)

저자들은 LIFT라는 새로운 방법을 고안했습니다. LIFT 는 학생이 현재 얼마나 도움을 받고 있는지에 따라 무엇을 가르치고 언제 가르칠지 정확히 아는 스마트한 튜터라고 생각하세요.

LIFT 의 작동 방식:

페이지가 대부분 비어있을 때 (후기 단계): 튜터는 "좋아, 아직 어렵고 희귀한 단어를 추측하지는 말자. 너에게는 단서가 충분하지 않아. 대신, 이렇게 적은 정보로도 실제로 알아낼 수 있는 쉽고 흔한 단어를 연습하자"라고 말합니다.
페이지가 대부분 명확할 때 (초기 단계): 튜터는 "좋아, 이제 단서가 많이 생겼구나. 쉬운 단어 연습은 그만두고 이전에 추측하지 못했던 어렵고 희귀한 단어에 집중하자"라고 말합니다.

사용 가능한 정보의 양에 따라 "쉬운" 목표와 "어려운" 목표 사이를 역동적으로 전환함으로써, LIFT 는 모델이 항상 유용한 것을 배우고 불가능한 추측이나 지루한 반복에 시간을 낭비하지 않도록 보장합니다.

결과: 더 똑똑하고 빠름

팀은 AIME 수학 경시대회와 같은 어려운 수학 추론 벤치마크에서 LIFT 를 테스트했습니다.

성능: LIFT 는 이전 방법보다 훨씬 뛰어난 성과를 보였습니다. 일부 어려운 수학 테스트에서는 표준 학습 방식에 비해 모델의 정확도를 3 배 향상시켰습니다.
효율성: 이것이 가장 인상적인 부분입니다. 일반적으로 이렇게 똑똑한 모델을 얻으려면 강화 학습 (RL) 을 사용해야 하는데, 이는 수일 동안 대규모 시뮬레이션을 실행하는 것과 같아 수천 시간의 슈퍼컴퓨터 시간이 소요됩니다. LIFT 는 500 배 적은 컴퓨팅 파워로 유사한 결과를 달성했습니다.

결론

이 논문은 모델이 특정 유형의 정보를 학습할 능력이 있는 시기를 이해함으로써, 확산 언어 모델을 훨씬 더 효과적으로 훈련시킬 수 있다고 주장합니다. 모델에게 한 번에 모든 것을 학습하도록 강요하는 대신, LIFT 는 전략적인 코치처럼 올바른 시기에 올바른 숙제를 부여합니다. 이는 AI 가 추론 작업에서 더 똑똑해지게 하면서 막대한 에너지와 비용을 절약합니다.

기술적 요약: 확산 언어 모델의 학습 가능성 기반 미세 조정

문제 정의

확산 언어 모델 (DLM) 은 토큰을 병렬로 생성하여 autoregressive 모델보다 추론 처리량 측면에서 상당한 이점을 제공하지만, 추론 능력을 향상시키기 위한 사후 학습 (post-training) 은 여전히 어렵습니다. 특히 autoregressive 모델의 표준적인 방법인 지도 미세 조정 (SFT) 은 종종 DLM 의 성능을 향상시키지 못하며, 오히려 사전 학습 (pre-training) 대비 성능을 저하시킬 수도 있습니다. 이러한 실패의 근본 원인은 충분히 연구되지 않았습니다. 저자들은 기존 SFT 가 학습 가능성 (learnability) 개념, 즉 확산 과정에서 어떤 토큰이 학습되는지, 그리고 언제 학습되는지를 간과하고 있음을 규명했습니다.

분석 결과 두 가지 중요한 차원이 드러났습니다:

무엇 (토큰 빈도): 희귀 토큰은 학습 중 노출이 제한적이기 때문에 공통 토큰보다 본질적으로 예측이 더 어렵습니다.
언제 (확산 시간): 토큰 예측의 난이도는 사용 가능한 컨텍스트의 양에 따라 달라집니다. 순방향 확산 과정에서 시간 $t \to 1^+$ 로 갈수록 입력이 심하게 마스킹되어 컨텍스트가 줄어듭니다. 이러한 조건에서 희귀 토큰은 비례적으로 매우 어렵거나 사실상 학습 불가능해집니다. 반면, $t \to 0^-$ 로 갈수록 컨텍스트가 풍부해져 희귀 토큰은 학습 가능해지고, 공통 토큰은 예측이 매우 쉬워집니다.

기존 SFT 는 모든 확산 단계에서 모든 토큰을 균일하게 학습하려 시도하여, 컨텍스트가 부족한 상황에서는 학습 불가능한 희귀 토큰을 학습하려 하거나 컨텍스트가 풍부한 상황에서는 trivial 한 공통 토큰을 학습하려 하는 비효율적인 학습을 초래합니다.

방법론: LIFT

이러한 과제를 해결하기 위해 저자들은 DLM 을 위한 효율적인 SFT 기반 사후 학습 알고리즘인 LIFT(Learnability-Informed Fine-Tuning) 를 제안합니다. LIFT 는 토큰 난이도와 사용 가능한 컨텍스트 간의 상호작용에 기반하여 각 확산 시간 단계에서 어떤 토큰을 감독할지 적응적으로 선택합니다.

핵심 메커니즘

LIFT 는 표준 무작위 마스킹 대신 "학습 가능성 기반 마스크"를 구성하여 작동합니다:

2 차 마스킹: 주어진 확산 시간 $t$ 에 대해, LIFT 는 2 차 마스킹 비율 $\rho \sim U(0, 1-t)$ 을 샘플링하여 더 손상된 입력 $x_{t+\rho}$ 를 생성합니다.
신뢰도 추정: 모델은 $x_{t+\rho}$ 의 마스킹된 위치에서 정답 토큰에 대한 신뢰도 점수 $c_k = p_\theta(x^k_0 | x_{t+\rho})$ 를 계산합니다.
적응적 선택: 현재 확산 시간 $t$ $t$ 에 기반하여 LIFT 는 감독할 토큰의 부분집합 $S_t$ $S_{t}$ 를 선택합니다:
- 높은 손상 ( $t \to 1^+$ ): 모델은 Top-K 토큰 (가장 높은 신뢰도/가장 쉬운 토큰) 을 선택합니다. 컨텍스트가 부족한 상황에서는 모델이 신뢰할 수 있게 학습할 수 있는 토큰에 집중합니다.
- 낮은 손상 ( $t \to 0^-$ ): 모델은 Bottom-K 토큰 (가장 낮은 신뢰도/가장 어려운 토큰) 을 선택합니다. 컨텍스트가 풍부한 상황에서는 이제 학습 가능한 희귀하고 어려운 토큰에 집중합니다.
- 중간 손상: 모델은 표준 기존 SFT(무작위 선택) 로 돌아갑니다.
손실 계산: 학습 손실은 선택된 부분집합 $S_t$ 에만 제한되며, 이는 Negative Evidence Lower Bound (NELBO) 목적 함수를 수정합니다.

변형

LIFT: 신뢰도 추정을 위한 한 번의 순전파와 손실 계산을 위한 한 번의 순전파, 총 두 번의 순전파가 필요합니다.
LIFT-A: 단일 순전파를 $t+\rho$ 에서 수행하고 선택되지 않은 토큰에 대한 손실을 0 으로 만들기 위해 게이팅 마스크를 적용하는 경량 근사 변형입니다. 이는 목적 함수에 약간의 편향을 도입하지만 계산 효율성을 크게 높입니다.

주요 기여

SFT 실패에 대한 체계적 분석: 본 논문은 DLM 에서 기존 SFT 가 실패하는 이유에 대한 최초의 체계적 분석을 제공하며, 이 실패를 토큰 난이도 (빈도) 와 특정 확산 시간 단계에서 사용 가능한 정보 간의 불일치로 규명합니다.
LIFT 알고리즘: "무엇"(토큰 식별/빈도) 과 "언제"(확산 시간) 간의 상호작용을 명시적으로 모델링하여 학습 유틸리티를 최적화하는 최초의 사후 학습 접근법인 LIFT 를 제안합니다.
효율성과 성능: LIFT 는 DLM 사후 학습에 대한 새로운 파레토 프론티어를 수립하여, 검증 가능한 보상과 함께 강화 학습 (RLVR) 방법과Comparable 한 성능을 달성하면서도 계산 자원을 수배 이상 적게 요구합니다.

실험 결과

저자들은 GSM8K, MATH, Countdown, Sudoku, 그리고 까다로운 AIME'24 및 AIME'25 데이터셋을 포함한 여섯 가지 추론 벤치마크에서 두 가지 DLM 기본 모델 (LLaDA-8B-Instruct 및 LLaDA 1.5) 에 대해 LIFT 를 평가했습니다.

성능 향상: LIFT 는 모든 벤치마크에서 기존 SFT 및 기타 베이스라인 (GIFT, CART) 보다 일관되게 우수한 성과를 거두었습니다.
- AIME'24에서 LIFT 는 기존 SFT 대비 3.3 배의 상대적 개선을 달성했습니다 (Pass@16 기준 16.7 대 6.7).
- AIME'25에서 LIFT 는 2.0 배의 상대적 개선을 달성했습니다 (Pass@16 기준 6.7 대 3.3).
- GSM8K, MATH, Countdown 에서도 개선이 관찰되었습니다.
계산 효율성: LIFT 는 약 1.8 H100 GPU 시간의 학습으로 이러한 성과를 달성했습니다. 반면, 유사한 성능을 달성하기 위해 RL 기반 베이스라인 (d1) 은 2,000 H100 GPU 시간 이상이 필요했습니다. 이는 계산량이 약 1000 배 감소한 것입니다.
강건성: LIFT 는 서로 다른 학습 데이터셋 (s1K 대비 더 큰 12K 데이터셋) 과 생성 길이 전반에 걸쳐 강건성을 입증했습니다.
Ablation 연구: 실험은 "무엇"과 "언제"를 함께 모델링하는 것이 중요함을 확인했습니다. 토큰 난이도 (Top-K 또는 Bottom-K) 만 고려하거나 확산 시간을 무시한 (무작위 선택) 변형들은 특히 AIME 와 같은 복잡한 추론 작업에서 LIFT 의 성능을 따라잡지 못했습니다.

중요성과 주장

본 논문은 LIFT 가 복잡한 추론 작업을 위해 DLM 을 실현 가능하게 만드는 중요한 진전이라고 주장합니다. 다양한 확산 단계에서 사용 가능한 정보와 학습 목적을 정렬함으로써 LIFT 는 각 학습 단계의 유틸리티를 극대화합니다.

저자들은 LIFT 가 새로운 계산 효율적 파레토 프론티어를 수립했다고 강조합니다. 이는 DLM 에서 정교한 추론 능력 향상이 반드시 강화 학습 (RL) 의 막대한 계산 오버헤드를 필요로 하지 않음을 보여줍니다. 대신, 신중하게 설계된 학습 가능성 기반 SFT 전략은 RL 성능을 비용의 일부로 달성할 수 있습니다. 이 연구는 향후 DLM 개발이 단순히 계산량을 늘리거나 RL 근사를 채택하는 것보다 토큰 특성과 확산 역학 간의 상호작용을 이해하는 데 우선순위를 두어야 함을 시사합니다.