기술 요약: FlashBlock

1. 문제 정의

블록 디퓨전(block diffusion)은 블록 단위 인과적 추론(block-wise causal inference)과 KV 캐싱을 결부터하여 디퓨전 언어 모델 및 비디오 생성과 같은 긴 문맥(long-context) 생성 모델의 추론 효율성을 개선해 왔으나, 긴 문맥 설정에서는 여전히 상당한 오버헤드 문제를 겪고 있습니다. 표준 블록 디퓨전에서는 모든 디퓨전 단계마다 전체적으로 성장하는 KV 캐시에 대해 어텐션을 재계산해야 합니다. 문맥 길이가 길어짐에 따라, 어텐션 계산과 KV 캐시 접근에 드는 결합 비용이 추론 지연 시간(latency)의 지배적인 요인이 됩니다.

본 논문은 블록 디퓨전의 미개척된 특성인 **단계 간 중복성(cross-step redundancy)**을 식별합니다. 구체적으로, 디노이징 과정 동안 현재 업데이트 중인 블록 외부의 토큰들은 크게 정적인 상태를 유지하며, 이로 인해 "블록 외부(block-external)" 토근들의 어텐션 기여도는 인접한 디퓨전 단계들 사이에서 매우 안정적인 양상을 보입니다. 반면, 현재 블록 내부("block-internal")의 어텐션은 토큰들이 활발하게 정제되거나 언마스킹(unmasking)됨에 따라 크게 변화합니다. 기존 방식들은 이러한 안정성을 활용하지 못하고, 중복되는 블록 외부 어텐션을 반복적으로 재계산합니다.

2. 방법론: FlashBlock

저자들은 밑바닥의 디퓨전 프로세스를 수정하지 않고도 계산량과 메모리 I/O를 줄이기 위해, 블록 외부 어텐션 출력을 캐싱하고 재사용하는 메커니즘인 FlashBlock을 제안합니다.

핵심 메커니즘

FlashBlock은 전체 어텐션 출력을 두 가지 구성 요소로 분해합니다:

1. 블록 내부 어텐션 ($A_{in}$): 활발하게 변화하고 있는 현재 블록($J_{in}$) 내의 토큰들에 대해 계산됩니다.
2. 블록 외부 어텐션 ($A_{out}$): 현재 블록 외부($J_{out}$)의 토큰들에 대해 계산되며, 단계가 지나도 안정적인 상태를 유지합니다.

이 방법은 이전 디퓨전 단계에서의 블록 외부 어텐션 출력($A_{out}$)과 그 로그 정규화 인자($L_{out} = \log Z_{out}$)를 캐싱합니다. 이후 단계에서는 $J_{out}$을 위해 전체 KV 캐시에 다시 접근하는 대신, 캐싱된 값을 재사용합니다. 전체 어텐션 출력은 수치적 안정성을 보장하기 위해 로그 공간 합성(log-space composition)(FlashAttention과 유사함)을 사용하여 캐싱된 $A_{out}$과 새로 계산된 $A_{in}$을 결합하여 재구성됩니다:
$$A_{full} = \frac{e^{L_{out}-m} A_{out} + e^{L_{in}-m} A_{in}}{e^{L_{out}-m} + e^{L_{in}-m}}$$
여기서 $m = \max(L_{out}, L_{in})$입니다.

적응형 재사용 전략

• 임계값 기반 재사용: 이산 디퓨전 모델(discrete diffusion models)에서 토큰 업데이트는 급격하게 일어날 수 있습니다. 본 방법은 업데이트된 토큰의 수($M^{s+1}$)에 대한 임계값 $\tau$를 도입합니다. 만약 $M^{s+1} < \tau$이면 캐싱된 어텐션을 재사용하고, 그렇지 않으면 정확성을 보장하기 위해 다시 계산합니다.
• 헤드별 선택 (비디오): 단계 간 변동성이 더 높은 비디오 디퓨전의 경우, 본 방법은 어텐션 헤드별로 유사도를 추정합니다. 유사도가 임계값 $\gamma$를 초과하는 헤드만 캐시를 재사용합니다.
• 희소 어텐션(Sparse Attention)과의 호환성: FlashBlock은 희소 어텐션과 직교(orthogonal)합니다. 이는 희소 어텐션에서 선택되지 않은 키(key)들을 "블록 외부" 집합으로 취급하여 그 잔여 기여도를 캐싱하는 방식입니다. 이를 통해 FlashBlock은 공격적인 희소화(sparsification)로 인한 정보 손실을 보완할 수 있습니다.

학습 적응 (선택 사항)

급격한 토큰 언마스킹으로 인해 발생할 수 있는 이산 디퓨전 언어 모델(dLLM)의 분포 불일치 문제를 해결하기 위해, 저자들은 재사용 인식 증류(reuse-aware distillation) 전략을 제안합니다. FlashBlock을 사용하는 학생 모델은 LoRA 어댑터를 활용하여 가볍게 적응하면서, 밀집(dense) 어텐션을 사용하는 고정된 교사 모델의 출력 분포를 모방하도록 학습됩니다.

3. 주요 기여

1. 경험적 통찰: 본 논문은 블록 외부 어텐션의 강력한 단계 간 안정성과 블록 내부 어텐션의 높은 변동성을 대조하는 경험적 연구를 제공합니다.
2. FlashBlock 메커니즘: 로그 공간 집계를 통해 이 중복성을 활용하여 계산 및 KV 캐스 접근을 줄이는 새로운 캐싱 메커니즘을 제시합니다.
3. 희소성와의 시너지: FlashBlock이 희소 어텐션의 보완적인 잔차 재사용(residual reuse) 전략으로서, 공격적인 희소화로 인한 품질 저하를 완화함을 입증합니다.
4. 효율성 이득: 생성 품질에 거의 영향을 미치지 않으면서 디퓨전 언어 모델 및 비디오 생성 모델의 추론 처리량(throughput)과 지연 시간을 크게 개선했습니다.

4. 실험 결과

실험은 디퓨전 언어 모델(Trado-8B/4B)과 비디오 디퓨전 모델(LongLive-1.3B)을 대상으로 수행되었습니다.

• 처리량(Throughput): FlashBlock은 디퓨전 언어 모델에서 최대 1.44배 높은 토큰 처리량을 달연했습니다 (예: 블록 크기 4일 때 312에서 451 tokens/s로 증가).
• 지연 시간(Latency): 어텐션 시간이 최대 1.6배 단축되었습니다.
• 품질: 생성 품질은 밀집(dense) 베이스라인과 유사한 수준을 유지했습니다. 수학적 추론(GSM8K, MATH500) 및 코딩 벤치마크(HumanEval, MBPP)에서 정확도 차이는 최소 수준(블록 크기에 따라 0.2% ~ 1.6% 이내)이었습니다.
• 긴 문맥 스케일링: 문맥 길이가 100k에서 800k로 증가함에 따라 FlashBlock과 베이스라인 간의 지연 시간 격차가 벌어졌으며, 특정 임계값 하에서 FlashBlock은 베이스라인에 비해 약 절반의 지연 시간 증가율을 보였습니다.
• 희소 어텐션 결합: SparseD와 결합했을 때(밀도 20%, 30%, 40%), FlashBlock은 SparseD 단독 사용 시보다 일관되게 높은 정확도를 보였습니다 (예: 20% 밀도에서 GSM8K 정확도를 34.72%에서 42.68%로 향상). 이는 희소화로 인해 손실된 정보를 효과적으로 복구함을 보여줍니다.

5. 의의

본 논문은 긴 문맥 블록 디퓨전의 핵심 병목 구간인 '안정적인 어텐션 기여도의 중복된 재계산' 문제를 FlashBlock이 해결한다고 주장합니다. 안정적인 "블록 외부" 문맥을 동적인 "블록 내부" 업데이트와 분리함으로써, 이 방법은 성장하는 KV 캐스와 관련된 계산 및 메모리 I/O 비용을 실질적으로 줄여줍니다.

저자들은 이 접근 방식이 기존의 희소 어텐션 기술과 **직교(orthogonal)**하며, 효율성과 정확성을 모두 향상시키는 보완적인 전략임을 강조합니다. 결과적으로 FlashBlock은 텍스트와 비디오 도메인 모두에서 효율적인 긴 문맥 생성을 위한 확장 가능한 솔루션을 제공하며, 생성된 콘텐츠의 품질을 저해하지 않으면서도 더 빠른 추론을 가능하게 합니다.