💬 NLP

Triplet-Block Diffusion RWKV

본 논문은 트리플릿 블록 레이아웃을 통해 인과적 RWKV 모델의 $O(L)$ 추론 효율성과 병렬 양방향 이산 확산을 통합한 새로운 아키텍처인 $B^3D-RWKV$ 를 소개하며, 기존 모델과 비교 가능한 정확도를 유지하면서 디코딩 처리량을 1.6 배 향상시킵니다.

원저자: Ke Lin, Yiyang Luo, Zhaolong Su, Yunya Song, Anyi Rao

게시일 2026-05-26

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ke Lin, Yiyang Luo, Zhaolong Su, Yunya Song, Anyi Rao

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

다음은 논문 "Triplet-Block Diffusion RWKV"(B3D-RWKV) 에 대한 설명으로, 창의적인 비유를 통해 간단한 개념으로 분해하여 정리한 것입니다.

큰 문제: "일방통행로" 대 "그룹 프로젝트"

이야기를 쓰는 두 가지 다른 방식을 상상해 보세요:

엄격한 인과적 모델 (일방통행로):
전통적인 AI 를 생각하면, 이미 입력한 단어만 볼 수 있는 작가와 같습니다. 그들은 왼쪽에서 오른쪽으로 단어 하나씩 씁니다. 다음에 무엇이 올지는 볼 수 없습니다.
- 장점: 새로운 단어를 추가할 때마다 처음부터 다시 읽을 필요가 없으므로 긴 문서를 읽는 속도가 매우 빠릅니다.
- 단점: 두 단어를 동시에 쓸 수 없기 때문에 텍스트를 생성하는 속도는 느립니다. 1 번 단어를 끝내야 2 번 단어를 시작할 수 있습니다.
확산 모델 (그룹 프로젝트):
반면, "무의미한 글"이나 "마스크"([MASK]와 같은) 로 가득 찬 빈 페이지에서 시작하는 다른 AI 를 상상해 보세요. 단어 하나씩 쓰는 대신, 전체 페이지를 한 번에 보고 누락된 단어가 무엇인지 추측한 후 몇 가지를 수정하고 이야기가 논리적으로 맞을 때까지 이 과정을 반복합니다.
- 장점: 여러 단어를 동시에 수정할 수 있어 (병렬 처리) 잠재적으로 훨씬 빠를 수 있습니다.
- 단점: 이를 위해서는 전체 문맥 (앞에 온 것과 뒤에 올 것) 을 동시에 볼 수 있어야 합니다. 이는 보통 매우 비싸고 느린 컴퓨터 아키텍처를 필요로 합니다.

갈등: 이 두 가지를 쉽게 섞을 수 없습니다. "일방통행로" 작가는 앞을 내다볼 수 없지만, "그룹 프로젝트"는 모든 것을 한 번에 봐야 하기 때문입니다.

해결책: "트리플릿-블록" 트릭

Ke Lin 과 동료 저자들은 트리플릿-블록 레이아웃이라는 교묘한 학습 트릭을 고안했습니다. 그들은 작가의 두뇌를 바꾸지 않고도 "일방통행로" 작가가 "그룹 프로젝트" 팀처럼 행동하도록 가르치는 방법을 찾아냈습니다.

이 트릭이 어떻게 작동하는지 연습 비유를 통해 설명해 보겠습니다:

왼쪽에서 오른쪽으로 대본만 읽을 수 있는 배우 (AI) 라고 상상해 보세요. 누락된 대사를 추측해야 하는 장면을 배워야 하지만, 정확하게 추측하려면 현재 위치보다 나중에 나오는 대사를 알아야 합니다.

저자들은 모든 장면을 위해 3 단계 연습을 구성했습니다:

1 부 (마스크된 복사본): 장면을 읽지만, 대사의 절반이 검은 테이프로 가려져 있습니다 ([MASK]). 이를 왼쪽에서 오른쪽으로 읽습니다.
2 부 (손실 가능한 마스크된 복사본): 검은 테이프가 붙은 정확히 같은 장면을 다시 읽습니다. 여기서 테스트를 받습니다. 누락된 대사를 추측해야 합니다.
- 마법 같은 점: 방금 1 부를 읽었기 때문에 당신의 두뇌 (AI 의 내부 기억) 는 1 부의 모든 보이는 대사를 이미 흡수했습니다. 2 부를 엄격하게 왼쪽에서 오른쪽으로 읽는 중이지만, 두뇌는 1 부에 저장된 덕분에 장면의 오른쪽 부분에서 온 문맥을 이미 "알고" 있습니다.
- 결과: 여전히 왼쪽에서 오른쪽으로 읽으면서도 과거와 미래를 모두 아는 "의사 양방향" 지식을 가지고 누락된 대사를 추측할 수 있게 됩니다.
3 부 (깨끗한 복사본): 완전하고 완벽한 장면을 한 번 더 읽습니다. 이는 다음 장면을 준비할 수 있도록 두뇌를 초기화합니다.

이 트리플릿 패턴 (마스크됨 -> 마스크됨/테스트 -> 깨끗함) 을 반복함으로써 AI 는 "미래"(실제로는 학습 시퀀스의 이전 블록) 에서 온 정보를 사용하여 누락된 단어를 예측하는 법을 배우며, 원래의 "일방통행로" 속도를 유지합니다.

결과: 빠르고 정확함

이 팀은 이 방법을 사용하여 B3D-RWKV-7.2B라는 모델을 구축했습니다. 그들이 발견한 바는 다음과 같습니다:

속도: "일방통행로" 아키텍처 (RWKV) 를 유지했기 때문에, 모델의 디코딩 속도가 놀라울 정도로 빠릅니다. 그들은 동일한 모델의 표준 버전보다 1.6 배 더 빠르다고 주장합니다.
지능: 질문 답변이나 이야기 쓰기 같은 일반 작업에서 다른 최상위 모델들과同等한 성능을 발휘합니다.
절충점: 논문은 완벽한 단계별 논리가 필요한 매우 복잡한 수학 문제의 경우, 확산의 "추측" 성질로 인해 때때로 작은 오류가 발생할 수 있다고 지적합니다. 그러나 대부분의 작업에서는 그럭저럭 잘 해냅니다.

한 마디로 요약한 내용

이 논문은 역설을 해결합니다: 어떻게 빠르고 왼쪽에서 오른쪽으로 쓰는 작가를 똑똑하고 모든 것을 보는 편집자처럼 만들 수 있을까?

그들은 작가가 장면을 연속해서 세 번 연습하도록 가르침으로써 이를 달성했습니다. 첫 번째 연습은 작가의 기억에 문맥을 채우고, 두 번째는 그 기억을 이용해 누락된 부분을 추측하는 연습을 하게 하며, 세 번째는 다음 장면을 위해 판을 깨끗이 합니다. 이를 통해 그들은 확산 모델의 병렬 처리 능력을 얻으면서도 선형 작성자의 속도를 유지할 수 있게 되었습니다.

그들이 주장하지 않는 것:

이것이 의료 진단이나 임상 용도로 작동한다고 주장하지 않습니다.
이것이 모든 AI 문제의 만능 해결책이라고 주장하지 않습니다. 복잡한 수학 구조에서는 약간 어려움을 겪는다고 인정합니다.
명시적으로 모델의 안전 기능을 변경하지 않았으므로, 원래 모델이 가지고 있던 편향성을 그대로 계승한다고 밝힙니다.

기술 요약: 트리플릿-블록 확산 RWKV (B3D-RWKV)

문제 제기

인과적 트랜스포머 언어 모델은 두 가지 주요 구조적 한계에 직면해 있습니다. 첫째, 엄격한 순차적 디코딩으로 인해 병렬화가 불가능하다는 점이며, 둘째, 단계당 이차적 주의 비용으로 인해 긴 문맥 추론이 계산적으로 비용이 많이 든다는 점입니다. 선형 시간 인과 모델 (예: RWKV) 과 이산 확산 모델은 각각 이러한 약점을 독립적으로 해결하지만, 이들의 통합은 본질적으로 일관성이 없었습니다. 확산 모델은 토큰 블록을 병렬로 탈노이즈하기 위해 양방향 주의가 필요한 반면, 인과 모델은 엄격히 단방향입니다. 기존에 이를 결합하려는 시도들은 종종 선형 시간 백본의 효율성을 훼손하는 구조적 수정을 요구하거나 처음부터 학습을 필요로 했습니다.

방법론

저자들은 근본적인 인과적 백본을 변경하지 않고 $O(L)$ 추론 효율성과 병렬 양방향 이산 확산을 통합하는 확산 스타일 RWKV 변형인 B3D-RWKV를 제안합니다. 핵심 혁신은 엄격히 인과적인 아키텍처 내에서 양방향 접근을 시뮬레이션하도록 설계된 학습 및 추론 전략인 트리플릿-블록 레이아웃입니다.

1. 트리플릿-블록 레이아웃 (학습)

인과 모델에서 확산을 가능하게 하기 위해, 크기 $B$ 의 논리적 생성 블록은 학습 샘플 내에서 세 개의 연속된 물리적 블록으로 확장됩니다:

$b_1$ (마스크된 복사본): 논리적 블록의 마스크된 버전.
$b_2$ (손실 가능한 마스크된 복사본): 탈노이즈 손실이 계산되는 동일한 마스크된 복사본.
$b_3$ (클린 정답 복사본): 후속 블록을 위한 순환 상태를 갱신하는 마스크되지 않은 버전.

의사 양방향 접근:
인과 백본이 토큰을 엄격히 왼쪽에서 오른쪽으로 처리하기 때문에, $b_2$ 의 임의의 마스크된 위치에 도달하는 숨겨진 상태는 이미 $b_1$ 전체를 처리한 상태입니다. $b_1$ 과 $b_2$ 가 동일한 마스크 패턴을 공유하므로, $b_1$ 의 마스크되지 않은 토큰들은 모델이 $b_2$ 의 마스크된 토큰을 예측하기 전에 숨겨진 상태에 효과적으로 "흡수"됩니다. 이는 모델이 논리적 블록 내의 자신의 마스크되지 않은 문맥에 의사 양방향 접근을 할 수 있게 하여, 확산의 양방향 요구 사항을 충족하면서도 RNN/SSM 백본의 단방향성을 유지합니다.

학습 목적:
모델은 마스크 예측 목적을 사용하여 학습됩니다. 손실은 다음으로 구성됩니다:

교차 엔트로피 손실 ( $L_{CE}$ ): 토큰이 마스크된 감독 위치에서 계산됨.
신뢰도 인식 병렬 손실 ( $L_{CAP}$ ): LLaDA-2.0 에 따라, 이 항은 이미 올바르게 예측된 위치의 확률 분포를 날카롭게 만들어 추론 샘플러가 단계당 더 많은 토큰을 확정하도록 장려합니다.

2. 추론: 블록 단위 반복적 탈노이즈

추론 중 모델은 한 번에 하나의 논리적 블록을 생성합니다. 이 과정은 다음을 포함합니다:

모든 토큰이 마스크된 블록을 초기화합니다.
모델이 마스크된 위치에 대한 상위 1 위 확률을 예측하는 반복적 탈노이즈 단계를 실행합니다.
상위 1 위 확률이 임계값 $\tau$ 를 초과하는 토큰을 확정합니다.
임계값을 통과하는 토큰이 $k_{min}$ 보다 적을 경우 진행을 보장하기 위해 상위 $k_{min}$ 개의 가장 확신 있는 토큰을 확정하는 폴백 메커니즘이 작동합니다.
모든 위치가 확정되면, 클린 블록이 접두사에 추가되고 다음 논리적 블록이 시작됩니다.

주요 기여

B3D-RWKV-7.2B: 마스크 예측 목적을 사용하여 7B 규모로 학습된 최초의 확산 스타일 선형 시간 RNN 언어 모델입니다. 이는 원래 매개변수를 수정하지 않고 RWKV-7 백본에 병렬 토큰 선택을 통합합니다.
트리플릿-블록 프레임워크: 엄격한 인과성과 전파 상태 속성을 만족하는 엄격히 인과적인 백본을 블록 확산 모델로 변환하는 보편적인 학습 레시피입니다. 이 접근법은 백본을 양방향 블록으로 수정하는 구조적 복잡성을 피합니다.
효율성과 성능: 이 방법은 경쟁력 있는 추론 능력을 유지하면서 표준 인과적 디코딩 대비 상당한 처리량 개선을 달성합니다.

실험 결과

모델은 일반 지식 (MMLU, ARC, PIQA, RACE) 및 추론/수학 작업 (GSM8K, MATH, GPQA) 을 포함한 8 개 작업 세트를 평가했습니다.

정확도: B3D-RWKV-7.2B 는 RWKV-7 베이스라인과 유사 규모의 다른 확산 모델 (예: LLaDA-8B, Dream-7B) 과 비교 가능한 정확도를 달성합니다. 특히 의사 양방향 지각이 추론을 향상시킨 것으로 여겨지는 ARC-Challenge 및 RACE 와 같은 추론 작업에서 베이스라인을 크게 능가합니다.
처리량: 모델은 1K 에서 512K 토큰에 이르는 문맥 길이에 걸쳐 RWKV-7 베이스라인 대비 평균 1.6 배의 속도 향상을 보여줍니다. 조정된 샘플링 매개변수를 사용하면 약간의 품질 트레이드오프로 2.02 배의 속도 향상을 달성할 수 있습니다.
수학 추론: 논문은 복잡한 수학 작업 (예: MATH) 에서 인과적 베이스라인 대비 성능이 약간 하락했다고 지적합니다. 저자들은 이를 부분 점수가 부여되지 않는 기호 수학 문제에 필요한 고도로 구조화된 단계별 검증을 처리하는 데 어려움을 겪는 병렬 디코딩의 특성 때문이라고 귀인합니다.

중요성 및 주장

이 논문은 구조적 변경 없이 사전 학습된 인과 언어 모델을 확산 언어 모델로 변환하는 효율적인 경로를 제공한다고 주장합니다. 트리플릿-블록 레이아웃을 활용함으로써 B3D-RWKV 는 선형 시간 모델의 $O(L)$ 추론 효율성과 확산의 병렬 디코딩 능력을 성공적으로 통합합니다.

저자들은 이 작업을 인과 및 확산 패러다임 간의 비호환성에 대한 구조적 해결책으로 위치시킵니다. 그들은 현재 단일 백본 (RWKV-7) 에 대한 검증, 학습 중 3 배의 물리적 시퀀스 길이 비용, 그리고 특정 작업에서의 정확도 후퇴에 기여할 수 있는 RL 정렬이 없는 상대적으로 작은 SFT 데이터셋의 사용 등 한계를 인정합니다. 그러나 핵심 기여는 구조적 재공학적 개입이 아닌 데이터 레이아웃을 통해 엄격히 인과적인 선형 시간 백본에서 고통처리 병렬 확산 생성이 가능하다는 것을 입증한 데 있습니다.