💬 NLP

LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

이 논문은 블록 단위 확산 언어 모델에서 KV 인플레이션 문제를 해결하고 메모리 병목 현상을 완화하기 위해, 안정적인 토큰의 캐시된 어텐션 결과를 재사용하고 활성 토큰에만 희소 어텐션을 적용하는 'LOSA(Locality-aware Sparse Attention)'를 제안하여 속도와 정확도를 동시에 향상시킨다고 요약할 수 있습니다.

원저자: Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang, Michael Mahoney, Chenfeng Xu, Kurt Keutzer, Amir Gholami

게시일 2026-04-15

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang, Michael Mahoney, Chenfeng Xu, Kurt Keutzer, Amir Gholami

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

📚 배경: 인공지능이 글을 쓰는 방식의 변화

과거의 AI(자율적 생성 모델)는 한 번에 한 글자씩 순서대로 글을 썼습니다. 마치 연필로 글을 쓸 때, 한 글자를 쓰고 나서 다음 글자를 쓰는 것처럼 말이죠.

하지만 최신 기술인 **'확산 언어 모델 (Diffusion Language Models)'**은 다릅니다. 이 AI 는 한 번에 여러 글자 (블록) 를 동시에 고쳐가며 글을 완성합니다. 마치 흩어진 퍼즐 조각을 한 번에 여러 개씩 맞춰가며 그림을 완성하는 것과 같습니다. 이 방식은 복잡한 추론이나 긴 글을 다룰 때 훨씬 유리합니다.

🚨 문제점: "불필요한 도서관 방문" (KV Inflation)

하지만 이 방식에는 치명적인 단점이 있었습니다.

비유:
AI 가 글을 고칠 때마다, **수천 권의 책이 꽂혀 있는 거대한 도서관 (기억 공간)**에서 필요한 정보를 찾아야 합니다.

기존 방식 (Naive Sparse Attention): AI 가 한 번에 10 개의 퍼즐 조각 (글자) 을 고친다고 칩시다. 각 조각마다 도서관에서 필요한 책 1 권씩을 찾아야 한다면, 10 개의 조각이 서로 다른 책을 원할 수 있습니다.

결과: AI 는 10 개의 조각을 위해 10 권의 서로 다른 책을 모두 도서관에서 꺼내야 합니다. 이렇게 되면, "한 번에 10 권을 찾아야 한다"는 목표가 무너져서, 실제로는 10 권이 아니라 100 권을 찾아야 하는 것처럼 느려집니다. 이를 논문에서는 **'KV 팽창 (KV Inflation)'**이라고 부릅니다.

✨ 해결책: LOSA (로컬리티 인지형 희소 주의)

LOSA 는 이 문제를 해결하기 위해 **인공지능의 '변화 패턴'**을 관찰했습니다.

관찰:
AI 가 퍼즐을 고치는 과정에서, 실제로 크게 변하는 조각은 아주 일부뿐입니다. 나머지 대부분의 조각은 이미 거의 완성되어 있어, 거의 변하지 않습니다.

변하는 조각 (Active Tokens): 새로운 정보가 필요해서 도서관에서 책을 새로 찾아야 합니다.

안정된 조각 (Stable Tokens): 이미 지난 번에 찾아본 책과 똑같은 정보를 필요로 하므로, 다시 도서관에 갈 필요가 없습니다.

🛠️ LOSA 의 작동 원리: "기억을 재사용하는 똑똑한 사서"

LOSA 는 다음과 같이 작동합니다.

변화 감지: AI 가 글을 고칠 때, 어떤 조각이 크게 변했는지 먼저 확인합니다.
선택적 방문:
- 크게 변한 조각들: 도서관 (기억 공간) 에 가서 필요한 책 (정보) 을 찾아옵니다.
- 안정된 조각들: "아, 이거는 지난 번에 봤던 거잖아!"라고 생각하며 지난번에 찾아온 책을 그대로 재사용합니다. 도서관에 다시 갈 필요가 없습니다.
결과:
- 도서관에 가는 횟수가 급격히 줄어듭니다.
- 불필요한 이동 (메모리 접근) 이 사라져서 속도가 매우 빨라집니다.
- 중요한 정보는 모두 기억하고 있기 때문에 정확도도 떨어지지 않습니다.

🏆 성과: 왜 이것이 중요한가요?

이 연구는 여러 실험을 통해 LOSA 의 효과를 입증했습니다.

속도: 기존 방식보다 최대 4 배 이상 빨라졌습니다. (RTX A6000 그래픽 카드 기준)
정확도: 속도를 높였음에도 불구하고, 긴 글을 이해하는 능력은 거의 그대로 유지되거나 오히려 더 좋아졌습니다. (기존 방식보다 정확도가 9 점이나 높았다는 결과도 있음)
효율성: 불필요한 정보 검색을 줄여 메모리 사용량도 1.5 배 이상 줄였습니다.

💡 한 줄 요약

LOSA는 인공지능이 긴 글을 쓸 때, "변하지 않는 부분은 다시 찾아보지 않고 기억을 재활용하고, 변한 부분만 집중적으로 찾아보게" 만들어서, 속도는 4 배 빨라지고 정확도는 그대로 유지되게 해주는 똑똑한 기술입니다.

이 기술이 상용화되면, 앞으로 AI 가 수만 자의 긴 문서나 복잡한 논문을 처리할 때 훨씬 빠르고 정확하게 답변해 줄 수 있을 것입니다.

1. 문제 정의 (Problem Statement)

블록 단위 확산 언어 모델 (Block-wise DLMs) 의 한계:

기존 자기회귀 (Autoregressive) 모델과 달리, 블록 단위 확산 언어 모델 (DLM) 은 블록 내 토큰을 임의의 순서로 생성하며 복잡한 추론 작업에 유망한 대안으로 부상했습니다.
그러나 긴 문맥 (Long-context) 환경에서 DLM 은 여전히 메모리 병목 (Memory-bound) 현상에 시달립니다. 각 디노이징 (Denoising) 단계에서 블록 내 모든 토큰이 캐시된 프리픽스 (Prefix) KV(키 - 값) 에 주의를 기울여야 하기 때문입니다.

기존 희소 어텐션 (Sparse Attention) 의 실패: KV Inflation 문제

자기회귀 모델용 희소 어텐션 기법 (예: QUEST) 을 DLM 에 단순히 적용할 때 KV Inflation(kv 팽창) 문제가 발생합니다.
원인: 블록 내의 서로 다른 쿼리 (Query) 토큰들이 서로 다른 프리픽스 KV 위치를 선택합니다. 메모리 병목 환경에서는 개별 쿼리가 아닌 블록 전체의 선택된 KV 위치의 합집합 (Union) 크기가 실제 메모리 트래픽과 지연 시간을 결정합니다.
결과: 각 쿼리가 적은 KV 만 참조하더라도, 쿼리 간 선택이 겹치지 않아 합집합 크기가 커지고, 이로 인해 희소성의 이점이 사라져 오히려 밀집 어텐션 (Dense Attention) 과 유사하거나 더 비효율적인 성능을 보입니다.

2. 방법론 (Methodology: LoSA)

저자들은 DLM 의 디노이징 과정에서 표현 변화의 국소성 (Locality of Representation Changes) 이라는 고유한 구조를 발견하고 이를 활용한 LoSA(Locality-aware Sparse Attention) 를 제안합니다.

핵심 관찰:

연속된 디노이징 단계 사이에서, 블록 내 토큰의 대부분 (Stable Tokens) 은 히든 상태 (Hidden State) 가 거의 변하지 않습니다.
오직 소수의 토큰 (Active Tokens) 만이 [MASK] 에서 실제 토큰으로 변경되거나 주변 토큰의 표현에 큰 변화를 일으킵니다.

LoSA 의 작동 원리:

로컬리티 프루닝 (Locality Pruning):
- 각 디노이징 단계에서 쿼리 벡터의 변화량 (MSE) 을 계산하여 토큰을 분류합니다.
- 변화가 큰 토큰은 Active Tokens, 변화가 적은 토큰은 Stable Tokens 로 구분합니다.
캐시 재사용 (Reuse for Stable Tokens):
- Stable Tokens 에 대해서는 이전 단계에서 계산된 프리픽스 어텐션 결과 (통계량 및 출력) 를 그대로 재사용합니다. 이는 해당 토큰들이 KV 캐시를 다시 로드할 필요가 없음을 의미합니다.
희소 어텐션 적용 (Sparse Attention for Active Tokens):
- Active Tokens 에 대해서만 희소 어텐션 (예: QUEST 알고리즘) 을 적용하여 프리픽스 KV 위치를 선택하고 계산합니다.
KV 합집합 축소:
- 전체 블록 크기 ( $B$ ) 대신 Active Tokens 집합 ( $|A|$ ) 에 대해서만 KV 인덱스를 선택하므로, 메모리에 로드해야 하는 KV 위치의 합집합 크기가 크게 감소합니다.
- 초기화 단계 (첫 번째 디노이징) 에는 밀집 어텐션을 수행하여 정확한 캐시를 구축한 후, 이후 단계부터 LoSA 를 적용합니다.

3. 주요 기여 (Key Contributions)

KV Inflation 문제의 규명: 블록 단위 DLM 에서 나쁜 희소 어텐션이 실패하는 근본 원인인 KV Inflation 문제를 최초로 식별하고 분석했습니다.
표현 변화 국소성 발견: 디노이징 단계 간 토큰 표현의 변화가 소수 토큰에 집중된다는 현상을 발견하고 이를 최적화에 활용했습니다.
LoSA 알고리즘 제안: Stable 토큰의 캐시 재사용과 Active 토큰의 희소 어텐션 결합을 통해 KV 캐시 로드량을 획기적으로 줄이는 새로운 아키텍처를 제안했습니다.
광범위한 검증: Trado-8B, SDAR-8B, Trado-4B 등 다양한 모델과 LongBench, 추론 벤치마크에서 LoSA 의 유효성을 입증했습니다.

4. 실험 결과 (Results)

정확도 (Accuracy):

LongBench: Trado-8B-Instruct 모델에서 공격적인 희소성 (Budget 128) 조건에서도 LoSA 는 QUEST 대비 평균 정확도 +10.43%p 향상 (41.97% vs 31.54%) 을 보였습니다.
밀집 어텐션 근접: 높은 희소성 수준에서도 Dense Baseline 과 유사한 정확도를 유지하며, 기존 희소 어텐션 방법 (SparseD, Sparse-dLLM) 보다 우월한 성능을 기록했습니다.

효율성 (Efficiency):

KV 캐시 밀도: LoSA 는 QUEST 대비 평균 1.54 배 낮은 KV 캐시 밀도를 유지하며 메모리 트래픽을 줄였습니다.
속도 향상 (Speedup): NVIDIA RTX A6000 GPU 에서 최대 4.14 배의 어텐션 속도 향상을 달성했습니다. (QUEST 는 3.26 배).
신규 하드웨어 호환성: RTX 5090 에서도 3.67 배의 속도 향상 추세를 유지하여 메모리 대역폭 의존성이 높은 작업임을 확인했습니다.

5. 의의 및 결론 (Significance)

DLM 의 실용성 확보: LoSA 는 블록 단위 확산 언어 모델이 긴 문맥 처리 시 겪는 메모리 병목 문제를 해결하여, 실제 배포 가능한 효율적인 추론을 가능하게 합니다.
새로운 최적화 패러다임: 단순히 KV 캐시를 줄이는 것을 넘어, 토큰의 동적 변화 특성을 분석하여 '어떤 토큰을 계산할지'와 '어떤 토큰을 캐시할지'를 지능적으로 분할하는 새로운 접근법을 제시했습니다.
미래 지향성: 이 연구는 확산 기반 언어 모델이 자기회귀 모델의 한계를 극복하고 긴 문맥 및 복잡한 추론 작업에서 경쟁력을 갖추기 위한 핵심 기술로 평가받습니다.

요약하자면, LoSA는 블록 단위 DLM 의 고유한 특성 (표현 변화의 국소성) 을 활용하여 KV Inflation 문제를 해결함으로써, 높은 정확도를 유지하면서 메모리 효율과 추론 속도를 극대화한 획기적인 방법론입니다.