💬 NLP

Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

본 논문은 양방향 마마(Mamba) 스캐닝을 활성 디노이징 블록으로 제한하여 정확한 캐싱을 가능하게 함으로써, 기존의 전체 시퀀스 및 어텐션 기반 확산 베이스라인 모델들보다 우수한 퍼플렉시티(perplexity)와 현저히 높은 롱 컨텍스트 추론 처리량을 달나성하는 하이브리드 확산 언어 모델인 BDLM Mamba-H를 소개한다.

원저자: Pranshu Chaturvedi, Parth Shroff, Tarun Suresh, Hangoo Kang, Kaiyue Wen

게시일 2026-07-07

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Pranshu Chaturvedi, Parth Shroff, Tarun Suresh, Hangoo Kang, Kaiyue Wen

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

당신이 아주 긴 이야기를 쓰려고 한다고 상상해 보세요. 하지만 당신에게는 엄격한 규칙이 하나 있습니다. 오직 방금 쓴 마지막 몇 페이지만을 기억할 수 있다는 것입니다. 문장을 하나 쓸 때마다, 당신은 공책의 맨 처음으로 다시 걸어가서 지금까지 쓴 모든 단어를 읽은 다음, 다음에 무엇을 쓸지 결정해야 합니다.

이것이 바로 현재의 "대규모 언어 모델(LLM, AI 챗봇)"이 긴 문맥을 다룰 때 겪는 어려움입니다. 대화가 길어질수록, AI는 새로운 단어를 하나 생성할 때마다 전체 기록을 매번 "다시 읽어야" 합니다. 이는 매우 느리고 많은 에너지를 낭비하게 만듭니다. 마치 매 걸음마다 점점 더 무거워지는 배낭을 메고 마라톤을 하는 것과 같습니다.

이 논문은 이러한 문제를 해결하기 위해 AI가 "기억하고" "생각하는" 방식을 바꾸는 새로운 방법인 BDLM Mamba-H라는 새로운 학습법을 소개합니다.

다음은 쉬운 비유를 사용한 상세 설명입니다.

1. 문제점: "무거운 배낭"

현재의 AI 모델은 질문에 답하기 전에 매번 교과서 전체를 다시 읽어야 하는 학생과 같습니다.

문제점: 이야기가 길어질수록 "배 backpack"(기억)은 너무 무거워집니다. 컴퓨터는 실제로 생각(연산)하는 시간보다 데이터를 옮기는 데(메모리 대역폭) 더 많은 시간을 소비하게 됩니다.
기존의 해결책: 일부 연구자들은 다른 종류의 메모리(Mamba라고 불리는)를 사용하여 배낭을 더 가볍게 만들려고 시도했습니다. 다른 이들은 한 번에 한 단어씩 쓰는 대신 여러 문장을 한꺼번에 쓰는 방식(Block Diffusion)을 시도했습니다.
결함: 연구자들이 이 두 가지 아이디어를 결합하려고 했을 때, 문제가 발생했습니다. "Mamba" 메모리 시스템은 제대로 작동하기 위해 이야기 전체를 역방향으로 훑어봐야 했습니다. 하지만 전체 이야기를 역방향으로 훑어보면, 이미 완료한 부분에 대한 "책갈피"를 저장할 수 없습니다. 매번 모든 것을 다시 읽어야만 합니다.

2. 해결책: "로컬 역방향 스캔 (Local Reverse Scan)"

저자들은 영리한 묘수를 제안합니다. 현재의 단락 안에서만 뒤를 돌아보는 것입니다.

당신이 챕터별로 책을 쓰고 있다고 상상해 보세요.

기존 방식 (전체 역방향): 새로운 문장을 쓰기 위해, 당신은 매번 마지막 페이지부터 첫 번째 페이지까지 책 전체를 다시 읽습니다.
새로운 방식 (BDLM Mamba-H):
1. "깨끗한" 부분: 한 챕터가 끝나면, 그 챕터를 안전한 금고에 넣습니다. 그리고 그 챕터에 대한 간단한 "요약 카드"(캐시)를 만듭니다. 당신은 다시는 금고를 열 필요가 없으며, 단지 요약 카드만 사용하면 됩니다.
2. "활성" 부분: 현재 챕터를 쓰는 동안에는, 문장들이 잘 흐르는지 확인하기 위해 해당 챕터 내에서 앞뒤로 훑어볼 수 있습니다.
3. 마법 같은 점: 당신이 현재 챕터 내부에서만 뒤를 돌아보기 때문에, 완료된 챕터들의 요약 카드는 업데이트될 필요 없이 유효하게 유지됩니다.

3. 결과: 속도와 지능

연구진은 이 새로운 방식이 기존 방식들에 비해 지능을 유지하면서도 속도를 높이는 두 가지 주요 목표를 달 수 있는지 테스트했습니다.

지능적인가? 네. 표준 독해 테스트(C4-en)를 통해 테스트했을 때, 새로운 모델(BDLM Mamba-H)은 작은 규모의 모델들(8,700만 파라미터) 사이에서 가장 높은 점수를 받았습니다. 모델을 더 크게 만들었을 때도(3억 5,000만 파라미터), 다른 최상위 모델들과 대등한 수준의 지능을 유지했습니다.
빠른가? 엄청난 차이가 있습니다.
- 중간 길이(65,000 단어)에서, 새로운 모델은 기존의 "전체 역방향" 방식보다 19.7배 더 빨랐습니다.
- 매우 긴 길이(262,000 단어)에서, Mamba를 사용하지 않는 기존의 가장 뛰어난 "블록" 방식보다 3.7배 더 빨랐습니다.

핵심 요약

이 새로운 모델을, 완성된 챕터를 다시 읽을 필요가 없도록 책갈피를 만들어 두면서도, 동시에 현재 챕터를 자유롭게 편집할 수 있게 된 작가라고 생각하십시오.

텍est 블록 내에서만 "뒤를 돌아보는 것"을 제한함으로써, AI는 메모리 트래픽에 발이 묶이지 않고 매우 긴 이야기를 생성할 수 있습니다. 이 논문은 이것이 긴 문맥을 처리하기 위한 실용적이고 강력한 구조임을 입증하며, 모든 기록을 매번 다시 처리하지 않고도 속도와 고품질의 글쓰기를 모두 가질 수 있음을 보여줍니다.

기술 요약: 부분 양방향성을 갖춘 하이브리드 블록 확산 언어 모델 학습

문제 정의
대규모 언어 모델(LLM)의 고처리량 장문맥 생성은 연산 능력보다는 주로 메모리 대역폭에 의해 제약됩니다. 표준 자기회귀 디코딩에서는 매 단계마다 메모리로부터 누적된 키/값(KV) 캐시를 스트리밍해야 하며, 이로 인해 컨텍스트 길이가 길어질수록 대역폭 요구량이 선형적으로 증가하는 반면 생성되는 토큰은 단 하나뿐입니다. 이를 완화하기 위해 두 가지 병렬 접근 방식이 존재하지만—선형 시간 믹서(예: Mamba)를 통한 메모리 액세스 감소와 블록 생성을 통한 병렬성 증대—이 둘을 결합하는 데에는 특정한 기술적 난관이 존재합니다.

DiffuMamba와 같은 초기 하이브리드 확산 모델은 전체 디노이징 윈도우에 대해 역방향 스캔을 포함하는 완전 양방향 Mamba 스캔을 사용합니다. 이 역방향 스캔은 활성 토큰에 의존하므로, 내부 상태가 프리픽스 전용(prefix-only)이 아닙니다. 결과적으로 이러한 상태들은 정밀하게 캐싱하거나 재사용할 수 없으며, 이는 블록 확산의 효율성 이점을 상쇄합니다. 핵심 과제는 양방향 믹싱의 이점을 유지하면서도 정확한 캐싱이 가능한 블록 단위 생성을 지원하는 하이브리드 아키텍처를 설계하는 것입니다.

방법론
저자들은 Mamba와 어텐션 레이어를 결합한 하이브리드 아키텍처인 BDLM Mamba-H를 제안합니다. 핵심 혁신은 역방향 Mamba 스캔을 전체 시퀀스가 아닌 활성 디노이징 블록으로만 제한하는 부분 양방향성(partial bidirectionality) 메커니즘입니다.

주요 방법론적 구성 요소는 다음과 같습니다:

캐시 정렬 학습 목적 함수(Cache-Aligned Training Objective): 모델은 완료된 프리픽스 블록이 다음을 저장하도록 하는 "올-블록(all-block)" 목적 함수를 채택합니다:
- 어텐션 레이어를 위한 KV 상태.
- Mamba 레이어를 위한 순방향 전용 경계 상태(재귀 상태).
  역방향 스캔은 현재 활성 블록( $b_t$ ) 내로 국한되어, 프리픽스 캐시( $C_\ell(c)$ )가 역방향 확산 타임스텝( $t$ )과 무관하게 깨끗한 프리픽스(clean prefix)만의 함수가 되도록 보장합니다.
타임스텝 인수 분해(Timestep Factorization): 캐시 유효성을 유지하기 위해, 아키텍처는 레이어 계산을 타임스텝이 없는 프리픽스 연산자와 타임스텝이 조건화된 활성 블록 연속체로 인수 분해합니다. 이는 깨끗한 프리픽스 캐시가 특정 역방향 확산 단계에 의존하게 되는 것을 방지하며, 이는 추론 시 유효한 재사용을 위한 필수 요구 사항입니다.
학습 프로토콜: 모델은 블록 단위의 이산 확산(discrete diffusion) 목적 함수를 사용하여 학습됩니다. 프리픽스 캐시는 깨끗한 프리픽스를 스캔하여 구축되며, 다운스트림 블록 손실은 이러한 프리픽스 상태에 저장된 표현을 직접 학습함으로써 학습 계산을 추론 시 샘플러의 캐시 인터페이스와 일치시킵니다.

주요 기여

캐시 정렬 학습: 저자들은 캐시된 객체(KV 상태 및 Mamba 프리픽스 상태)가 추론 시의 블록 캐시와 정확히 일치하는 BDLM 올-블록 목적 함수를 정식화했습니다. 이를 통해 다운스트림 블록 손실이 재귀 상태 재사용에 사용되는 프리픽스 스캔을 학습할 수 있게 합니다.
제어된 검증 및 규모 확장: 저자들은 DCLM 데이터셋을 사용하여 87M 파라미터에서 종합적인 스윕(sweep)을 수행하였으며, Complete(d)P 스타일의 학습률 전이 규칙을 사용하여 선택된 모델들을 350M 파라미터로 확장했습니다.
장문맥 처리량 이득: 연구는 확장된 컨텍스트 길이에서의 추론 처리량을 평가하며, 하이브리드 접근 방식이 컨텍스트 길이가 증가함에 따라 풀 시퀀스 베이스라인 및 순수 어텐션 기반 블록 확산 모델보다 성능이 크게 우수함을 입증합니다.

결과

검증 품질 (87M 파라미터): BDLM Mamba-H는 87M 모델 중 가장 우수한 C4-en 검증 퍼플렉시티(61.6)를 기록하며, BDLM 어텐션(76.5), 풀 시퀀스 DiffuMamba-H(83.7), 그리고 풀 시퀀스 어텐션(87.6)을 능가했습니다.
검증 품질 (350M 파라미터): 350M 규모에서 BDLM Mamba-H는 BDLM 어텐션과 경쟁력 있는 수준을 유지하며, BDLM 어텐션의 36.97 대비 38.2의 C4-en 퍼플렉시티를 달성했습니다.
처리량 성능:
- 65K 토큰에서 BDLM Mamba-H는 풀 시퀀스 DiffuMamba-H 베이스라인 대비 19.7배의 처리량을 달성했습니다.
- 262K 토큰에서 BDLM Mamba-H는 BDLM 어텐션 대비 3.7배의 처리량을 달성했습니다.
- 풀 시퀀스 모델은 올-블록 목적 함수의 오버헤드로 인해 짧은 컨텍스트(8192 토큰)에서는 더 빠르지만, 캐싱을 사용하는 네이티브 BDLM 모델은 장문맥 환경에서 탁월한 확장성을 보여줍니다.

의의 및 주장
본 논문은 BDLM Mamba-H가 **장문맥 확산(long-context diffusion)**을 위한 실행 가능한 아키텍처임을 주장합니다. 역방향 Mamba 스캔을 활성 블록으로 제한함으로써, 모델은 블록 간의 정확한 캐싱을 가능하게 하여 이전의 양방향 확산 하이브리드에서 나타났던 상태 재사용 문제를 해결합니다. 저자들은 이 접근 방식이 하이브리드 Mamba-어텐션 모델이 경쟁력 있는 검증 품질을 유지하면서도 장문맥 추론을 위한 상당한 처리량 개선을 제공할 수 있다고 주장하며, 블록 부분 역방향 하이브리드가 확산 언어 모델(dLLM)을 확장하기 위한 실질적인 방향임을 강조합니다.

본 논문은 최적성에 대해서는 신중한 태도를 유지하며, 규모 확장에 사용된 학습률 전이 규칙이 기존 파라미터화 법칙에 근거한 강력한 후보이지만 이산 확산 언어 모델(dLLM)에 대해 증명된 최적은 아니라고 언급했습니다. 또한, 블록 목적 함수가 장문맥 생성을 개선하지만, 현재의 하드웨어 구성에서 풀 시퀀스 베이스라인에 비해 학습 처리량 페널티(약 25~34% 느림)를 발생시킨다는 점을 인정했습니다.

1. 문제점: "무거운 배낭"

2. 해결책: "로컬 역방향 스캔 (Local Reverse Scan)"

3. 결과: 속도와 지능

핵심 요약

유사한 논문