💬 NLP

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

이 논문은 불확실한 예측을 가진 이미 선택된 위치에 강하게 어텐션(attend)하는 후보 토큰들을 탐욕적으로 할인함으로써, 정확도를 유지하면서도 추론 단계를 줄여 마스크 확산 언어 모델의 고도로 병렬화된 디코딩 품질을 향상시키는 훈련이 필요 없는 재순위화 규칙인 ADAS를 소개한다.

원저자: Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher, Paolo Favaro

게시일 2026-06-10

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher, Paolo Favaro

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

복잡한 퍼즐을 풀려고 한다고 상상해 보세요. 하지만 조각을 한 번에 하나씩 놓는 대신, 작업을 더 빨리 끝내기 위해 여러 조각을 동시에 놓으려고 합니다. 이것이 바로 **마스크 확산 언어 모델(Masked Diffusion Language Model)**이라는 새로운 유형의 AI가 직면한 과제입니다.

이 모델들은 "마스크(가려진 빈칸)"로 가득 찬 문장에서 시작하여 점진적으로 그 빈칸들을 채워 나가는 방식으로 작동합니다. 목표는 단 한 번의 단계에서 가능한 한 많은 빈칸을 채워 시간을 절약하는 것입니다. 하지만 문제가 하나 있습니다. 특정 빈칸을 채우는 데 자신감이 있다고 해서, 그 이웃에 있는 빈칸도 동시에 채워야 한다는 뜻은 아닙니다. 때때로 두 빈칸은 매우 깊게 연결되어 있어서, 하나를 잘못 추측하면 다른 하나의 추측까지 망쳐버릴 수 있습니다.

이 논문은 AI가 더 현명한 그룹 결정을 내릴 수 있도록 돕는 무료 도구인 **ADAS(Attention-Discounted Adaptive Sampler, 어텐션 할인 적응형 샘플러)**를 소개합니다.

문제점: "취약한 그룹"

AI의 현재 추측 과정을 저녁 메뉴를 결정하려는 친구들의 모임이라고 생각해 보세요.

기존 방식 (표준 샘플러): AI는 각 친구를 개별적으로 봅니다. "앨리스는 피자를 먹고 싶어 할 확률이 90%야. 밥은 파스타를 먹고 싶어 할 확률이 90%지." 그래서 AI는 말합니다. "좋아! 그럼 지금 바로 둘 다 주문하자."
결함: 만약 앨리스와 밥이 항상 다투는 커플이라면 어떻게 될까요? 만약 AI가 이들을 함께 결정하도록 강요한다면, 그들은 서로를 상쇄시키거나, AI는 너무 늦게서야 그들의 결합된 선택이 말이 안 된다는 것을 깨닫게 될 수도 있습니다. 논문에 따르면, AI가 긴밀하게 연결된 두 단어를 동시에 추측하도록 강요받았을 때 정확도가 71%에서 30%로 떨어졌습니다. 이는 마치 한 손으로 두 개의 깨지기 쉬운 달걀을 저글링하려는 것과 같았습니다. 각각은 괜찮지만, 함께 있으면 위험한 상태입니다.

해결책: ADAS ("스마트 그룹화" 도구)

ADAS는 조각을 배치하기 전에 그룹을 살펴보는 현명한 중재자 역할을 합니다. ADAS는 단어들이 서로 어떻게 연관되어 있는지 이해하기 위해 AI가 이미 계산하고 있는 특수한 "어텐션(attention)" 신호를 사용합니다.

ADAS가 작동하는 방식은 다음과 같은 간단한 비유를 통해 설명할 수 있습니다.

점수표: 모든 빈칸에는 "신뢰도 점수"(AI가 정답에 대해 얼마나 확신하는지)가 있습니다.
할인(Discount): ADAS는 이번 라운드에서 AI가 채우기로 이미 결정한 빈칸들을 살펴봅니다. 만약 새로운 후보 빈칸이 아직 불확실하거나 불안정한 상태인 이미 선택된 빈칸을 "바라보고(attention)" 있다면, ADAS는 **할인(discount)**을 적용합니다.
- 비유: 당신이 계주 팀을 뽑고 있다고 상상해 보세요. 아주 빠른 주자가 한 명 있습니다(높은 신뢰도). 하지만 당신은 이 주자가 계속해서 자기 신발 끈에 걸려 넘어지고 있는 팀원을 불안하게 뒤돌아보는 것을 목격했습니다(불확실한 예측). 이때 ADAS는 이렇게 말합니다. "이 주자는 빠르지만, 넘어지는 팀원 때문에 주의가 분산되었습니다. 이 특정 그룹을 위해 이 주자의 우선순위 점수를 낮춥시다(우선순위를 낮춤)."
결과: AI는 여전히 최선의 후보들을 뽑지만, "위험하게 결합된" 조각들을 함께 묶는 것은 피합니다. 그렇다고 이들을 영원히 금지하는 것은 아닙니다. 단지 그룹이 더 안정될 때까지 기다리는 것뿐입니다.

왜 특별한가?

추가 학습 불필요: AI에게 생각하는 법을 다시 가르칠 필요가 없습니다. ADAS는 기존 방식 위에 얹혀지는 "플러그 앤 플레이(plug-and-play)" 방식의 규칙입니다. 이는 카메라 렌즈에 새로운 필터를 추가하는 것과 같습니다. 카메라(AI 모델)는 그대로 유지되지만, 사진(출력물)은 더 선명하게 나옵니다.
모든 규칙과 호환 가능: AI가 5개의 단어를 채우기로 결정하든, 혹은 "너무 불확실해지면" 멈추도록 설정되어 있든, ADAS는 이 모든 규칙과 함께 작동합니다. ADAS는 단지 어떤 단어가 먼저 뽑힐지를 바꿀 뿐입니다.
속도 vs 품질: 이 논문은 수학 문제(방정식 풀이 등)와 코딩 작업에서 이를 테스트했습니다. 연구 결과, AI가 매우 빠르게 움직이려 할 때(많은 단어를 한꺼번에 채울 때) 기존 방식은 많은 실수를 범했습니다. ADAS는 이를 해결하여 평균적으로 약 9~10%포인트의 정확도를 개선하면서도, 시간 비용은 거의 들지 않았습니다(약 3% 정도만 느려짐).

핵심 요약

이 논문은 불확실한 이웃과 너무 밀접하게 연결된 단어의 가치를 "할인"하는 것만으로도, AI가 한 번에 더 많은 빈칸을 안전하게 채울 수 있다고 주장합니다. 이는 취약하고 속도 중심적인 추측 게임을 더 견고한 과정으로 바꾸어 놓으며, 이를 통해 모델이 추가적인 학습이나 복잡한 새로운 하드웨어 없이도 더 빠르고 정확하게 작동할 수 있게 해줍니다.

기술 요약: 마스크 확산 언어 모델을 위한 어텐션 할인 적응형 샘플러 (Attention-Discounted Adaptive Sampler)

문제 정의

마스크 확산 언어 모델(Masked Diffusion Language Models, MDLMs)은 여러 토큰을 한 번의 디노이징 반복(denoising iteration) 과정에서 동시에 드러낼 수 있게 함으로써 인과적 자기회귀(causal autoregressive) 모델과는 차별화된 속도-품질 트레이드오프를 제공한다. 그러나 이러한 병렬성은 취약성을 동반한다. 개별 토큰 예측은 확신도가 높더라도, 여러 토큰을 동시에 결정(commit)하는 것은 해당 위치들이 강하게 결합(coupled)되어 있을 경우 위험할 수 있기 때문이다. 기존의 훈련 불필요 샘플러(예: Top-k, Fast-dLLM, EB-Sampler)는 일반적으로 토큰별 지표(신뢰도, 마진, 엔트로피)를 기준으로 후보를 순위 매기고, 부분 집합 크기를 결정하기 위해 정지 규칙(stopping rule)을 적용한다. 이러한 방법들은 고득점 후보들이 서로 호환 가능하다는 것을 암묵적으로 가정하며, 의존적인 위치들이 함께 언마스킹(unmasking)될 때 발생하는 공동 의존성 오류(joint dependence error)를 간과한다. 이론적 연구에 따르면, 여러 위치를 동시에 결정할 때는 양호한 주변 예측(marginal predictions)만으로는 불충분하며, 오차에는 모델 오차와 구별되는 "공동 의존성" 성분이 포함된다.

방법론: ADAS

저자들은 기본 샘플러의 정지 규칙을 변경하거나 추가 학습을 요구하지 않고, 병렬 마스크 확산 디코딩을 개선하기 위해 설계된 훈련 불필요 재순위화(reranking) 규칙인 ADAS(Attention-Discounted Adaptive Sampler)를 제안한다.

핵심 메커니즘

ADAS는 의존성을 고려한 효용 함수를 도입하여 부분 집합 구축 과정을 수정한다. 단순히 토큰별 점수를 합산하는 대신, ADAS는 이미 선택된 위치 중 예측이 불확실한 위치에 강하게 어텐션(attend)하는 후보들을 할인(discount)함으로써 언마스크할 위치의 집합을 탐욕적(greedily)으로 구축한다.

후보 위치 $i$ 에 대한 주변 효용(marginal utility) $u(i | S, A)$ 는 이미 선택된 집합 $S$ 와 어텐션 행렬 $A$ 가 주어졌을 때 다음과 같이 정의된다:
$u(i | S, A) = c_i - \alpha \sum_{s \in S} A_{is}(1 - c_s)$
여기서:

$c_i$ 는 표준 신뢰도 점수(위치 $i$ 의 가장 가능성 높은 토큰 확률)이다.
$A_{is}$ 는 위치 $i$ 에서 위치 $s$ 로의 어텐션 점수를 나타낸다(디노이저의 마지막 셀프 어텐션 레이어에서 유도됨).
$(1 - c_s)$ 는 이미 선택된 위치 $s$ 의 불확실성을 나타낸다.
$\alpha$ 는 할인 강도를 조절하는 하이퍼파라미터이다.

운영 로직

탐욕적 구축 (Greedy Construction): 빈 집합에서 시작하여, ADAS는 주변 효용이 가장 높은 후보를 반복적으로 선택한다.
의존성 페널티 (Dependency Penalty): 후보는 단순히 다른 선택된 위치와 결합되어 있다는 이유만으로 페널티를 받는 것이 아니라, 특히 선택되었으면서도 불확실한 위치와 결합되어 있을 때 페널티를 받는다. 이는 강한 제약(hard constraint)이라기보다 부드러운 주변 페널티(soft marginal penalty) 역할을 한다.
정지 규칙 보존 (Stopping Rule Preservation): 선택 과정은 기본 샘플러의 허용 조건(예: Top-k의 고정 카디널리티 $k$ , Fast-dLLM의 신뢰도 임계값, 또는 EB-Sampler의 엔트로피 예산)을 만족할 때까지 계속된다.
통합 (Integration): ADAS는 정지 기준을 재설계하거나 새로운 플래너를 학습시킬 필요 없이 기존 샘플러에 플러그인할 수 있는 모듈형 재순위화 단계이다.

주요 기여

공동 의존성 진단: 저자들은 결합된 위치들(예: 동일한 산술 방정식 내의 리터럴들)을 함께 예측하도록 강제하면 정확도가 크게 떨어진다는 것(제어된 진단 실험에서 71%에서 30%로 하락)을 입증했다. 또한, 셀프 어텐션 점수가 의존적인 쌍과 의존적이지 않은 쌍을 효과적으로 구분하며, 의존적인 쌍이 일관되게 더 높은 어텐션 값을 보임을 보여주었다.
ADAS 도입: 정지 규칙과 순위 매기기를 분리한, 훈련이 필요 없는 어텐션 할인 방식의 탐욕적 선택기를 제안한다. 이는 최소한의 런타임 오버헤드(포워드 패스당 3.1%)를 가지며 Top-k, Fast-dLLM, EB-Sampler와 호환된다.
실증적 검증: 두 가지 MDLMs(LLaDA-8B-Base, Dream-7B-Base)와 네 가지 벤치마크(GSM8K, MATH500, HumanEval, MBPP)에 대한 광범 l한 평가를 통해 저-NFE(고병렬) 환경에서의 일관된 성능 향상을 보여준다.

실험 결과

평가는 병렬 언마스킹이 가장 공격적이고 오류가 발생하기 쉬운 "저-NFE(low-NFE)" 환경에 집중된다.

성능 이득: 모든 데이터셋과 정지 기준에 걸쳐 ADAS는 매치된-NFE(matched-NFE) 성능을 개선한다.
- LLaDA-8B-Base: 평균 +9.11 퍼센트 포인트 개선.
- Dream-7B-Base: 평균 +10.46 퍼센트 포인트 개선.
강건성: 이득은 엔트로피 기반, 신뢰도 기반, 고정 예산 정지 규칙 전반에서 일관되게 나타난다. 특히 결합된 위치들이 일관되지 않은 출력을 유발하기 쉬운 코드 생성 벤치마크(HumanEval, MBPP)에서 개선 효과가 두드러진다.
베이스라인과의 비교: ADAS가 적용된 샘플러는 병렬 토큰 수가 증가함에 따라 KLASS 및 DAPD와 같은 다른 훈련 불필요 방식들과 대등하거나 이를 능가하는 성능을 유지한다. DAPD는 어텐션을 사용하여 하드 그래프 제약(독립 집합)을 구성하는 반면, ADAS의 소프트 할인 방식은 하드 제약이 취약할 수 있는 고병렬 환경에서 더 견고하다는 것이 증명되었다.
오버헤드: 이 방법은 주로 마지막 어텐션 레이어에 접근하고 탐욕적 업데이트를 수행하기 때문에 모델 포워드 패스당 3.1%의 런타임 오버헤드를 발생시킨다.

의의 및 주장

본 논문은 **소프트 어텐션 할인 재순위화(soft attention-discounted reranking)**가 마스크 확산 언어 모델의 고병렬 디코딩 품질을 개선하는 단순하고 모듈화된 효과적인 전략이라고 주장한다.

저자들은 토큰별 휴리스틱이 공동 의존성을 설명하지 못하는 특정 실패 모드를 ADAS가 해결한다는 점을 강조한다. 모델 자체의 셀프 어텐션을 결합 강도의 대리 지표(proxy)로 사용함으로써, ADAS는 학습된 정책, 룩어헤드 탐색(lookahead search) 또는 복잡한 그래프 제약을 요구하지 않고도 의존적인 토큰을 동시에 결정하는 위험을 완화한다. 결과는 병렬 디코딩을 개선하기 위해서는 더 나은 정지 규칙뿐만 아니라, 각 단계에서 드러낼 위치의 부분 집합을 구성하기 위한 더 나은 목적 함수(objective)가 필요함을 시사한다.

저자들은 셀프 어텐션이 확률적 의존성의 대리 지표일 뿐이라는 점과, 이 방법의 효용성을 더 넓은 모델 규모와 작업 도메인에서 조사할 필요가 있다는 점을 언급하며 한계점에 대해 겸허한 태도를 유지한다. 그들은 ADAS를 기존의 학습된 언마스킹 정책 및 룩어헤드 방법과 상호 보완적인 관계로 위치시킨다.