플래시블록 긴 컨텍스트 블록 디퓨전 효율을 위한 어텐션 캐싱

초록

플래시블록은 블록 디퓨전에서 현재 블록 외부 토큰에 대한 어텐션이 확산 단계마다 거의 변하지 않는다는 사실을 이용해, 해당 어텐션을 캐시하고 재사용함으로써 KV‑cache 접근과 연산량을 크게 줄인다. 기존 블록 디퓨전 흐름을 그대로 유지하면서도 희소 어텐션과 결합 가능하며, 언어·비디오 디퓨전 모델에서 토큰 처리량 1.44배, 어텐션 시간 1.6배 감소를 달성한다.

상세 분석

플래시블록은 블록 디퓨전 모델이 긴 시퀀스를 처리할 때 겪는 핵심 병목을 정확히 짚어낸다. 블록 디퓨전은 매 확산 단계마다 현재 블록을 업데이트하면서 이전 블록의 KV‑cache를 재활용한다. 그러나 매 단계마다 전체 컨텍스트(이전 블록 전체)에 대해 어텐션을 다시 계산하므로, 컨텍스트가 길어질수록 어텐션 연산과 메모리 이동 비용이 기하급수적으로 증가한다. 논문은 이 현상을 “블록 외부 어텐션의 교차‑스텝 안정성”이라는 관점에서 분석한다. 실험적으로 동일 블록 내에서 인접한 확산 단계 사이에 외부 토큰이 생성하는 어텐션 가중치와 정규화 값이 0.80.9 이상의 높은 코사인 유사도를 보이며, 내부 토큰은 0.40.6 수준으로 크게 변한다는 것을 시각화하였다. 이는 외부 토큰이 이미 충분히 안정된 표현을 가지고 있어, 매 단계마다 재계산할 필요가 없다는 강력한 증거다.

이를 기반으로 플래시블록은 두 가지 핵심 메커니즘을 제안한다. 첫째, 외부 어텐션(A_out)과 그 정규화 로그(L_out)를 현재 단계에서 캐시하고, 이후 단계에서는 KV‑cache를 읽지 않고 캐시된 값을 재사용한다. 둘째, 현재 블록 내부 어텐션(A_in, L_in)은 매 단계 새로 계산한다. 두 부분을 로그‑스페이스에서 합산하는 방식은 FlashAttention의 수치 안정화 기법을 차용해, 반정밀도에서도 오버플로우 없이 정확한 결과를 얻는다.

플래시블록은 또한 선택적 재사용 정책을 도입한다. 블록 내에서 업데이트된 토큰 수(M) 가 사전 정의된 임계값 τ 이하이면 외부 어텐션을 그대로 재사용하고, 그 이상이면 재계산한다. 비디오 디퓨전에서는 헤드별 유사도 γ 를 기준으로 재사용 여부를 결정해, 시각적 변동성이 큰 경우에도 정확성을 유지한다.

희소 어텐션과의 호환성도 설계에 포함된다. 기존 희소 어텐션이 선택한 키 집합을 J_in 으로, 선택되지 않은 나머지를 J_out 으로 두고, 플래시블록은 J_out 에 대한 잔여 어텐션을 캐시한다. 따라서 희소 어텐션이 가져오는 연산 감소와 플래시블록이 제공하는 교차‑스텝 중복 제거가 상호 보완적으로 작동한다.

실험 결과는 두드러진 효율성을 보여준다. 언어 디퓨전 모델에서는 토큰 처리량이 최대 1.44배 증가하고, 어텐션 연산 시간은 1.6배 감소한다. 비디오 디퓨전에서도 유사한 비율의 가속이 관찰되었으며, BLEU·ROUGE·FID 등 품질 지표는 거의 변동이 없었다. 특히, 희소 어텐션과 결합했을 때는 희소 비율이 70% 이상일 때도 품질 저하를 최소화하면서 가속을 유지한다.

전반적으로 플래시블록은 블록 디퓨전의 구조적 특성을 활용해, 기존 모델 아키텍처를 크게 변경하지 않으면서도 장기 컨텍스트에서의 효율성을 크게 향상시키는 실용적인 솔루션이다.