📰 딥시크, AI의 긴 문맥 처리를 바꾸다

기존 Sparse Attention 연구는 크게 세 가지 범주로 나눌 수 있습니다.

1. Inference 전용 방식

   훈련은 Full Attention으로 진행하고 추론 시에만 희소 패턴을 적용합니다. 이 방식은 훈련 중에는 전혀 희소 구조를 경험하지 않기 때문에 모델이 구조에 최적화되지 못하고, 추론 시 성능 저하가 발생할 수 있습니다.

2. Fixed Pattern 기반

   ‘Fixed(고정된)’라는 표현에서 알 수 있듯이, 고정된 주기와 범위로만 토큰을 선택해 연산합니다. 구현이 간단하고 일정한 속도 향상을 보장하지만, 상황별 중요 정보에 맞춰 유연하게 적응하지 못해 장기 의존성 정보 활용이 제한됩니다.

3. Query-Aware Selection

   쿼리의 특성에 따라 중요한 토큰을 동적으로 골라내는 방식입니다. 하지만 비연속적인 메모리 접근으로 GPU 효율이 떨어지고, 선택 과정이 비미분 연산일 경우 훈련 시 직접 학습이 어렵습니다.

기존 방식들은 분명 각자 장점이 있지만, 속도·성능·훈련 가능성을 모두 만족시키기는 어려웠습니다. 이에 연구진은 아이디어를 떠올립니다:

희소 패턴을 처음부터 배우고, 하드웨어에 맞게 설계하자.

NSA는 이를 위해 <계층적 희소 전략(Hierarchical Sparse Strategy)>을 도입했습니다.

NSA는 기존의 Sparse Attention 방식과 달리, 하나의 쿼리에 대해 세 가지 서로 다른 경로를 동시에 처리하도록 설계되었습니다. 이렇게 병렬로 경로를 운용하는 접근은 단일 경로나 고정 패턴만을 사용하는 기존 방식과 차별화되며, 각 경로가 전역·세부·로컬 정보 처리를 분담함으로써 보다 효율적이고 균형 잡힌 정보 활용이 가능합니다.

1. Token Compression (토큰 압축)

2. Token Selection (토큰 선택)

3. Sliding Window (슬라이딩 윈도우)

이후 각 경로의 기여도를 계산하여, 상황에 따라 전역, 세부, 그리고 로컬 정보의 비중을 조정하는데요. 계층적 희소 전략을 통해 모델은 꼭 필요한 정보만 골라 계산하기 때문에 불필요한 연산과 메모리가 크게 줄어듭니다. 성능 저하 없이 처리량을 줄일 수 있게 되지요.

연산 구조에 대한 전략을 세웠으니, 이제 하드웨어를 준비할 차례입니다. NSA는 최신 LLM에서 사용하는 Grouped-Query Attention(GQA) 구조를 차용했습니다. 좀 더 알아볼까요?

GQA는 여러 Query 헤드가 각자 다른 Key-Value 캐시(cache)를 쓰는 대신, 하나의 KV 캐시를 공유하는 구조입니다. KV 캐시는 이전 토큰의 Key와 Value를 GPU 메모리에 저장해, 매 토큰마다 전체 문맥을 다시 계산하지 않고 필요한 부분만 추가 연산을 수행해 속도를 높이는데요.

NSA는 이 표준 구조에 하드웨어 최적화를 더합니다. KV 데이터를 그룹 단위로 한 번만 로드하고, 해당 그룹의 모든 헤드가 이를 재사용하도록 만들어 불필요한 중복 로드를 줄이지요. 덕분에 긴 문맥에서도 처리 효율을 높일 수 있습니다. 📈

이렇게 줄어든 데이터 접근을 GPU에서 효율적으로 수행하도록, 연구진은 ‘Triton 기반 최적화’를 시킵니다. 여기서 말하는 Triton은 GPU용 고성능 연산 코드를 파이썬과 유사한 문법으로 작성할 수 있도록 설계된 오픈소스 언어이자 프레임워크인데요. 특히 딥러닝 모델의 연산을 하드웨어 특성에 맞춰 세밀하게 제어하고 최적화할 수 있습니다. 연구진이 Triton으로 최적화 시킨 세 가지 기법을 간단하게 보겠습니다.

• Group-Centric Loading: 동일 그룹의 쿼리를 한 번에 로드하여 메모리 접근 효율을 극대화
• Shared KV Fetching: 읽어온 KV 블록을 그룹 내 모든 헤드가 공유해 불필요한 중복 로드를 제거
• Balanced Loop Scheduling: GPU에 작업 부하를 균등하게 배분하여 처리 병목을 방지

실제로 Triton을 기반으로 한 NSA의 성능은 어땠을까요?

문맥 길이에 따른 연산 지연(latency)을 비교한 결과, NSA는 모든 구간에서 FlashAttention-2보다 짧은 처리 시간을 보였습니다. 특히 문맥 길이가 길어질수록 NSA의 속도 우위가 더 뚜렷하게 나타나는데요. 이렇게 계층적 희소 전략을 통해 불필요한 연산을 줄이고, GQA를 통해 남은 연산을 GPU에서 최대 효율로 처리함으로써 NSA는 기존 방식 대비 속도 향상과 성능 유지를 동시에 달성했습니다.

NSA 성능을 검증하기 위해 연구진은 27B 파라미터 규모의 대형 언어 모델을 사용해 260B 토큰을 학습시켰고, 최대 64k 토큰 길이의 문맥을 처리할 수 있도록 설정했습니다. 이 조건에서 여러 벤치마크를 평가한 결과는 아래와 같았습니다!

• 일반 벤치마크: Full Attention과 동등~우위
• LongBench: Full 대비 평균 +0.032 점 향상
• 수학 추론(AIME24): 더 높은 정답률
  
  

게다가 우측 그래프를 보면 긴 문맥 처리 환경에서 NSA가 모두 앞질렀음을 알 수 있습니다.

올해 ACL 최우수 논문 4편 중 2편이 중국 연구진의 작품이었습니다. 절대적인 수는 많지 않지만, 전 세계 석학들이 모여 치열하게 경쟁하는 자리에서 상위 50%를 차지했다는 건 의미 있는 성과입니다. 미국과 중국이 AI 분야의 양대산맥으로 자리매김하고 있는 가운데, 우리나라도 AI 기술 자립을 향한 본격적인 행보를 시작했습니다. <독자 AI 파운데이션 모델> 프로젝트가 그 출발점인데요.

셀렉트스타 역시 SK텔레콤이 주관하는 컨소시엄에 참여 기업으로 함께하게 되었습니다!👏🏼