Hymba는 2024년 NVIDIA에서 제안한 모델입니다.
[Hymba]
https://arxiv.org/abs/2411.13676
Hymba: A Hybrid-head Architecture for Small Language Models
We propose Hymba, a family of small language models featuring a hybrid-head parallel architecture that integrates transformer attention mechanisms with state space models (SSMs) for enhanced efficiency. Attention heads provide high-resolution recall, while
arxiv.org
<Hymba>
* Abstract
1. 저희는 Hymba라는 소형 언어 모델 계열을 제안하며, 이는 하이브리드-헤드 병렬 아키텍처를 특징으로 합니다.
2. 해당 아키텍처는 트랜스포머 어텐션 메커니즘과 상태 공간 모델(SSM, State Space Models)을 통합하여 효율성을 향상시킵니다. 어텐션 헤드는 고해상도 기억(recall)을 제공하며, SSM 헤드는 효율적인 문맥 요약(context summarization)을 가능하게 합니다.
3. 또한, 프롬프트 앞에 삽입되는 학습 가능한 메타 토큰(learnable meta tokens)을 도입하여 중요한 정보를 저장하고 어텐션 메커니즘에서 발생하는 강제적인 어텐션(forced-to-attend) 부담을 완화하였습니다.
4. 모델의 최적화를 위해 레이어 간 키-값(KV) 공유(cross-layer key-value sharing) 및 부분 슬라이딩 윈도우 어텐션(partial sliding window attention)을 적용하여 캐시 크기를 최소화하였습니다.
5. 특히, Hymba는 소형 언어 모델(LM) 중 최신(state-of-the-art) 성능을 달성하였습니다.
* Introduction
[기존 모델들의 한계]
1. 트랜스포머 모델: 이차적(quadratic) 연산 비용 + 높은 메모리 요구량
2. SSM(Mamba, Mamba-2 등) 모델: 메모리 회상(memory recall) 작업에서의 한계로 인해 일반적인 벤치마크 성능 저하
3. 기존의 어텐션 + SSM 레이어 적층 하이브리드 모델: 특정 작업에 적합하지 않은 레이어가 병목현상을 유발
[Hymba의 핵심 아이디어]
1. 하이브리드-헤드(hybrid-head): 하나의 레이어 내에서 어텐션 헤드와 SSM 헤드 통합
→ 동일한 입력을 병렬적이며 상호 보완적인 방식으로 처리할 수 있도록 설계
→ 어텐션의 고해상도 기억 기능과 SSM의 효율적인 문맥 요약 기능을 동시에 활용하여 효과적인 정보 처리
2. 입력 시퀀스 앞에 학습 가능한 메타 토큰(learnable meta tokens) 추가
→ 슬라이딩 윈도우 어텐션(sliding window attention)에서도 모든 이후 토큰과 상호작용
→ 세계 지식(world knowledge)의 압축 표현 역할을 수행
→ 소프트맥스 어텐션이 공백(nothing)에 어텐션을 할 수 없는 문제 완화
3. 레이어 간에 KV(Kay-Value) 캐시 공유
→ 일반적으로 어텐션 헤드 간에 KV 공유
* Road Map
[융합된 하이브리드 모듈(Fused hybrid modules)]
* 하이브리드-헤드 모듈 내에서 어텐션 헤드와 SSM 헤드를 병렬로 융합하는 방식은 순차적 적층(sequential stacking)보다 성능 우수
* 두 헤드가 동일한 정보를 동시에 처리하므로 추론(reasoning)과 기억력(recall) 정확도가 향상
[효율성 및 KV 캐시 최적화(Efficiency and KV cache optimization)]
* 어텐션 헤드는 작업 성능을 향상시키지만, KV 캐시 요구량을 증가 및 처리량(throughput) 감소
* 이를 완화하기 위해, 1. 하이브리드-헤드 모듈을 최적화하여 로컬 및 글로벌 어텐션을 결합 2. 레이어 간 KV 캐시 공유
→ 처리량은 3배 향상, KV 캐시는 약 4배 감소
[메타 토큰(Meta Tokens)]
* 메타 토큰의 주요 역할
- 어텐션 드레인(attention drain) 완화 – 백스톱 토큰(backstop tokens) 역할을 수행하여 어텐션을 효과적으로 재분배
- 세계 지식(world knowledge) 압축 표현 – 학습된 정보의 요약된 형태로 활용됨
[스케일링(Scaling)]
* 300M(3억) 개의 파라미터를 가진 모델에서 100B(1,000억) 개의 학습 토큰을 사용하여 제거 연구(ablation studies) 수행
* 최종 모델은 1.5T(1.5조) 개의 토큰으로 학습 → 350M(3억 5천만) 및 1.5B(15억) 개의 파라미터를 가진 모델로 확장
* Architecture
융합된 하이브리드 모듈
* 기본 트랜스포머의 다중 헤드 어텐션 구조(multi-head attention)에서 영감
→ SSM의 처리 효율성과 문맥 요약 능력 + 어텐션의 고해상도 기억 기능
* attention 헤드와 SSM 헤드가 동일한 정보를 각각 다른 방식으로 저장, 검색 및 처리 가능
* 두 operator의 강점을 모두 계승
[설계 수식]
* W_{in_proj}: X^~를 변환
* W^Q, W^K, W^V: 어텐션 헤드의 query, key, value 생성
* W^{SSM}: SSM 헤드의 입력 특징
* W^{G}: Gate 생성
* 어텐션 헤드 수식
* SSM 헤드 수식
* 출력 크기 정규화 및 융합
* 학습 가능한 벡터를 이용해 normalization 및 re-scailing
* 출력을 averaging한 후 최종 출력 프로젝션 적용
* beta_1, beta_2 는 learnable
[기억 관점]
* 하이브리드-헤드 모듈의 구성 요소는 인간 두뇌의 기능과 유사한 방식으로 해석 가능
* 어텐션 헤드
* 인간의 스냅샷 메모리(snapshot memory) 와 유사한 역할 수행
* 고해상도 기억(high recall resolution) 제공
* 특정 순간이나 사건의 세부적인 정보 저장
* SSM 헤드
* 인간의 희미해지는 기억(fading memory) 과 유사한 역할 수행
* 과거 사건의 세부 사항은 점차 잊어버리지만, 핵심 정보(core, gist)는 유지
효율성 및 KV cache 최적화
* 하이브리드-헤드 모듈은 기억력과 추론 능력을 향상시키지만, 어텐션 헤드의 KV 캐시 요구량이 증가하여 메모리 및 처리량 효율성 저하
→ 이를 해결하기 위해, KV 캐시 크기를 줄이면서도 기존과 유사한 작업 성능을 유지하는 것을 목표
[글로벌 및 로컬 어텐션 결합]
* 하이브리드-헤드 모듈은 이미 SSM 헤드를 통해 전역 문맥을 요약 → 적극적으로 글로벌 어텐션을 로컬 어텐션으로 대체 가능
[로컬 어텐션과 글로벌 어텐션 비율 조정]
* 모든 레이어에서 글로벌 어텐션을 SWA로 대체한 결과, recall-intensive 작업에서 성능이 20% 이상 감소하는 문제 발생
* 이에 대응하여, 일부 레이어에서 글로벌 어텐션을 점진적으로 다시 도입하였습니다.
* 첫 번째, 중간, 마지막 레이어(총 3개)에서만 글로벌 어텐션을 사용하면 recall-intensive 성능을 회복하면서도 상식 추론(commonsense reasoning) 정확도를 유지할 수 있음을 확인
→ 처리량 2.7배 증가, KV 캐시 크기 3.8배 감소 달성
[레이어 간 KV 공유]
* 최근 연구에서 인접한 레이어 간 KV 캐시가 높은 유사도를 가지므로 각 레이어에 개별적인 KV 캐시를 사용하는 것은 메모리 낭비 및 파라미터 중복을 초래한다고 보고됨
* 연속된 두 개의 레이어가 동일한 KV 캐시를 공유하는 Cross-layer KV Sharing 적용
* KV 메모리 사용량과 모델 파라미터를 동시에 감소 → 처리량 1.5배 증가, 기억력 유지하며 상식 추론 성능 0.60% 향상
메타 토큰
* 초기 토큰(initial tokens)이 의미론적으로 중요하지 않음에도 불구하고, 이후 토큰들로부터 높은 어텐션 점수를 받는 경향이 있음을 관찰
[메타 토큰의 도입]
* 학습 가능한 메타 토큰 집합 R = [r_1, r_2, ... , r_m] 도입하여 이를 초기 입력 토큰으로 사용
* 입력 시퀀스 X = [x_1, x_2, ... , x_n]가 주어졌을 때, 모델에 입력되는 최종 시퀀스는
[기억 관점의 해석]
* 메타 토큰 → 인간 두뇌의 메타메모리(metamemory)와 유사한 역할 수행
* 메타메모리: 인간이 어디에서 필요한 정보를 찾을 수 있는지 인식하는 능력을 의미합니다. 즉, 기억 자체를 조절하고 관리.
[메타 토큰의 역할]
1. 토큰 덮어쓰기(token overwriting) 방지
* 기존 연구에 따르면, 어텐션은 특정 토큰을 과도하게 참조(over-attend)하거나 덮어쓰는(overwrite) 경향이 존재
* 학습 가능한 토큰을 추가하면, 보다 대표성 있는(representative) 특징 맵(feature maps) 형성 가능
2. 강제적 어텐션(forced-to-attend) 문제 해결을 위한 exit tokens
* 입력 앞에 토큰을 추가하면, 소프트맥스 함수의 분모를 수정하여 어텐션의 형태에 영향
* 메타 토큰은 학습 가능한 형태로 만들어져, 최적의 소프트맥스 형태를 학습할 수 있도록 한다.
3. KV 캐시 및 SSM 상태 초기화
* 초기 토큰을 학습하는 것은 "학습된 프롬프트 튜닝(learned prompt tuning)" 또는 "학습된 초기화(learned initialization)"의 한 형태
* 추론 시, 메타 토큰은 고정된 상태로 존재하며, 키(keys)와 값(values)은 사전에 계산되어 저장 가능
[메타 토큰에 의한 기억력과 상식 추론 정확도 향상]
* 메타 토큰 도입 이후 어텐션 및 SSM 헤드의 엔트로피가 전반적으로 감소
* 추론(reasoning) 및 기억력(recall) 성능이 향상된 결과와 결합하여 볼 때, 메타 토큰이 어텐션 및 SSM 헤드가 작업 성능에 가장 중요한 토큰 하위 집합(subset of important tokens)에 집중하는 데 도움을 줄 수 있음을 시사
Attention 패턴
[Hymba's attention map]
* Hymba의 어텐션 맵(attention map) 내 요소들을 다음 네 가지 유형으로 분류하였다:
1. ‘Meta’ (메타 토큰 어텐션)
* 모든 실제(real) 토큰이 메타 토큰에 주는 어텐션 점수
* 모델이 메타 토큰에 얼마나 집중하는지를 반영
* 어텐션 맵에서 첫 번째 몇 개의 열(예: Hymba의 경우 128개 열)에 위치
2. ‘BOS’ (시퀀스 시작 토큰 어텐션)
* 모든 실제 토큰이 시퀀스 시작(Beginning-of-Sequence, BOS) 토큰에 주는 어텐션 점수
* 어텐션 맵에서 메타 토큰 다음의 첫 번째 열에 위치
3. ‘Self’ (자기 참조 어텐션)
* 모든 실제 토큰이 자기 자신에 주는 어텐션 점수
어텐션 맵에서 대각선(diagonal) 상에 위치
4. ‘Cross’ (다른 토큰으로의 어텐션)
모든 실제 토큰이 다른 실제 토큰들에 주는 어텐션 점수
어텐션 맵에서 비대각선(off-diagonal) 영역에 위치
[vs Transformer]
1. vanilla Transformer
* 어텐션 점수가 ‘BOS’ 토큰에 과도하게 집중
* ‘Self’ 어텐션 점수의 비율이 높음
2. Hymba
* 메타 토큰, 어텐션 헤드, SSM 헤드가 상호 보완적으로 작동하여 보다 균형 잡힌 어텐션 분포를 형성
* 메타 토큰이 ‘BOS’ 토큰의 어텐션 점수를 분산(offload) → 모델이 실제(real) 토큰에 더 집중할 수 있도록 유도
* SSM 헤드는 글로벌 문맥을 요약하면서 ‘Self’ 어텐션 점수에 더 집중 → 현재 토큰 중심으로 문맥을 정리
* 어텐션 헤드는 ‘Self’ 및 ‘BOS’ 토큰에 대한 어텐션을 줄이고, ‘Cross’ 어텐션 점수를 증가 → 다른 실제 토큰들 간의 관계를 더 많이 고려
* Model Evaluations
Benchmark with SOTA small LMs
* Hymba-1.5B 모델은 단 1.5T(1.5조) 개의 사전 학습 토큰만으로도 모든 2B(20억) 이하의 언어 모델(LM) 중 최고의 성능을 달성
Benchmark Different Architectures Under The Same Setting
* Hymba 모델은 대부분의 작업에서 다른 1B 모델보다 우수한 성능을 지속적으로 기록