[2026-1] 박승원 - Hymba: A Hybrid-head Architecture for Small Language Models

논문 정보

https://arxiv.org/pdf/2411.13676

저자: Xin Dong*, Yonggan Fu* et al.

소속: NVIDIA

 

핵심 아이디어

Hymba는 Transformer의 attention mechanism과 Mamba 계열의 State Space Model을 결합한 hybrid architecture이다.

 

기존 hybrid 모델들이 attention layer와 SSM layer를 번갈아 쌓는 방식이었다면, Hymba는 한 layer 안에서 attention head와 SSM head를 병렬로 배치한다. 이를 통해 attention의 정확한 recall 능력과 SSM의 효율적인 context summarization 능력을 함께 활용한다. 또한 meta tokens, sliding window attention, cross-layer KV sharing을 통해 cache size와 inference 속도도 개선한다.

 

Introduction

Transformer는 self-attention을 통해 입력 sequence 내 token 간 관계를 직접 계산하며 강력한 성능을 보여왔다. 특히 attention은 이전 token을 직접 참조할 수 있기 때문에, 긴 문맥 안에서 특정 정보를 다시 찾아야 하는 recall task에 강하다.

 

하지만 Transformer는 문장 길이가 길어질수록 계산량과 메모리 사용량이 커진다. 특히 추론 과정에서는 이전 토큰들의 key와 value를 저장하는 KV cache가 필요하기 때문에, 긴 context에서는 메모리 부담이 커진다.

 

반면 Mamba와 같은 SSM은 과거 정보를 recurrent state 형태로 압축하여 전달하기 때문에 훨씬 효율적이다. 하지만 정보를 압축하는 구조상 특정 token을 정확히 다시 찾는 능력은 attention보다 약할 수 있다.

 

이러한 trade-off를 해결하기 위해 Hymba를 제안한다. Hymba는 attention과 SSM을 layer 단위로 번갈아 쌓는 대신, 같은 layer 안에서 attention head와 SSM head를 병렬로 사용하는 hybrid-head architecture를 도입한다.

 

Transformer

강점: 긴 문맥 속 특정 정보 탐색(recall task)

약점: Long context에서 메모리 부담 커짐

 

SSM

강점: 과거 정보 압축으로 메모리 효율적

약점: Fine-grained token retrieval 약함.

 

Hymba: Hybrid-Head Architecture

기존 하이브리드 모델 구조: Layer-wise stacking

e.g.)

Attention Layer
SSM Layer
Attention Layer
SSM Layer

 

한계: 특정 layer가 하나의 방식으로만 정보를 처리하기 때문에 병목이 생길 수 있다. 예를 들어 SSM layer에서는 정확한 recall이 어렵고, attention layer에서는 효율적인 context summarization이 어렵다. 

 

Hymba

이 한계를 보완해서 layer-wise stacking 방식이 아닌, attention과 SSM heads를 같은 layer 안에서 parallel 하게 연결함. 이렇게 하면 모든 헤드가 같은 input을 처리하게 됨.

 

e.g.) 

In a layer:

Attention head || SSM head || Attention head || SSM head

 

Efficiency and KV cache optimization

Attention 헤드가 성능을 향상시키지만, 긴 문맥에 경우에는 KV cache에 필요한 메모리가 더 많아짐. 이런 이슈를 해결하기 위해서 local과 global attention을 결합하고, 여러 레이어가 KV cache를 공유하는 구조를 택함. 이 방식은 처리량을 3배로 늘리고, cache를 4배 이상 감소시킨다.

 

Meta Tokens

Meta Tokens는 입력 앞에 붙이는 128개의 학습가능한 임베딩이다. 이 토큰들은 학습된 cache 초기화값처럼 작동해서 모델이 관련 정보에 더 잘 집중하도록 함. 또한 meta tokens는 backstop tokens로써 attention이 실제 중요한 토큰이 아니라 별 의미 없는 토큰에 과도하게 쏠리는 attention drain 현상을 완화시키는 역할을 수행함. 세계 지식(World knowledge)의 압축이라고 볼 수 있음.

 

A Fused Hybrid Head Module

SSM의 weak recall performance를 보완하기 위해, attention과 SSM heads를 같은 layer안에서 병렬적으로 연결.

 

Design formulation

한 layer 안에서 attention head와 ssm head를 병렬로 계산한 뒤에 normalize시키고 learnable scale로 맞춘 다음 합침.

 

$\tilde{X} = [\text{Meta Tokens}, \text{X}]$

$W_{\text{in\_proj}}
=
[W_Q, W_K, W_V, W_{SSM}, W_G]$

$W_{SSM}$은 SSM에 들어갈 feature, $W_G$는 SSM의 output gate.

Attention head: $Y_{\text{attn}} = \text{softmax}(QK^T) W_V \tilde{X} = M_{attn}\tilde{X}$ where $M_{\text{attn}} = \text{softmax}(QK^T)W_V$

풀어서 말하자면 입력 token들을 서로 비교해서, 현재 입력에 따라 달라지는 mixing matrix $M_{attn}$를 만들고, 그 Matrix로 입력 정보를 섞는 연산.

 

SSM head로는 Mamba모델을 채택.

$\alpha_{i,j}
=
C_i
\left(
\prod_{k=j+1}^{i}
\exp(A\Delta_k)
\right)
B_j\Delta_j$

의미: j번째 토큰의 입력이 i번째 토큰의 output에 얼마나 영향을 주는지. 오래된 토큰일수록 더 많은 transition을 거쳐 현재 output에 도달.

SSM head: $Y_{\text{ssm}} = G \odot \alpha(A,B,C,\Delta) W_{\text{SSM}}\tilde{X}$

gate G는 SSM이 만든 output을 선택적으로 통과시킴. 입력에 따라서 필요한 정보만 선택적으로 내보내는 역할을 수행.

 

 

 

SSM head의 출력의 크기가 attention head의 것보다 커서, normalize 시키고, 학습가능한 벡터(베타1과 베타 2)를 사용해서 rescale함. 

Total fusion formulation: $Y =
W_{\text{out\_proj}}
\left(
\beta_1 \text{norm}(M_{\text{attn}}\tilde{X})
+
\beta_2 \text{norm}(M_{\text{ssm}}\tilde{X})
\right)$

 

메모리 측면에서의 해석

인간의 뇌와 비교했을 때 Attention heads는 마치 스냅 사진처럼 높은 recall 해상도를 가지고 있음. 반면, SSM heads는 맥락을 요약하기에 점차 사라져가는 기억과 유사함. 

Hymba에서는 요약된 전체적 맥락이 로컬 정보를 기억하기 위한 스냅샷 메모리를 더 많이 할당할 수 있게 해주면서도 recall 능력을 유지할 수 있게 한다. 

-> 전역적 컨텍스트(요약되어서 저장됨, 용량 작음.) + 스냅샷 메모리(전역적 컨텍스트 용량이 작으므로 여기에 더 많이 할당할 수 있음.)

이는 global attention을 줄이고 local attention을 사용함으로써 이룰 수 있었음.

 

Head importance analysis

(1) Head의 상대적 중요성은 task에 따라, input에 따라 달라짐. 고로 여러가지의 input을 받을 때 수행하는 역할이 다르다는 것을  추론할 수 있음.

(2) 첫번째 layer의 SSM head는 language modeling에 중대한 역할을 수행함.

(3) 하나의 attention/ssm head를 없애는 건, 평균 accuracy를 0.24%/1.1% 감소시킴.

 

KV Cache Optimization

연구진들은 recall과 reasoning 성능을 올리기만 한 것 뿐 아니라 KV cache optimization을 통해 메모리 사용량을 감축하고 처리량을 증가시킴.

 

Combine global and local attention

Local attention인 Sliding Window Attention(SWA)을 사용함으로써 더 효율적인데다가, SSM head가 전역적인 정보를 요약해주기 때문에 더 공격적으로 global attention을 대체할 수 있음.(그러나 완전한 대체는 성능 하락) SWA만 쓰면 recall 성능이 20% 넘게 떨어짐.

Global attention은 초기 layer, 중간 layer, 마지막 layer 3곳에만 넣어도 충분함.

 

2.7배의 처리량 향상과 3.8배의 캐시 절약을 이뤄냄.

 

Cross-layer KV sharing

각 layer마다 고유한 KV cache를 하지 않고, 인접한 layer끼리의 KV cache는 동일하게 관리. 

연속되는 2개의 layer는 같은 kv cache를 공유한다. 

이렇게 되면 kv memory와 model parameter 수를 줄일 수 있음.

1.15배의 처리량 증가를 일으킴.

 

Meta Tokens

의미적으로 큰 의미를 갖지 않더라도 초기의 토큰에 지나치게 많은 attention이 쏠리는 것을 연구들이 발견. 예컨대, Llama3.2-3B 모델에서는 BOS 토큰에 50% 이상의 attention이 쏠림. 

 

이렇게 의미 없이 attention을 낭비하기 보다는 학습 가능한 meta token들을 앞에 넣어서 한 개의 초기 토큰에만 쏠리던 attention을 여러 meta token으로 분산시키고, 각 토큰마다 담고 있는 지식을 다르게 만들어 다른 task나 다른 domain에서는 다른 meta token이 활성화 되는 방식으로 attention을 가이드할 수 있게 됨.

 

또한 추론 시에는 meta token이 고정되므로 미리 계산해둘 수 있음. -> learned KV cache initialization

 

실제로 meta token을 추가한 뒤로 head들의 entropy가 낮아짐.(더 적은 수의 토큰에 집중하게 됨)

 

Attention heat map을 봤을 때도, 일반 transformer는 BOS와 Self에 attention이 많이 쏠리는 경향을 보이지만, Hymba는 더 균형있게 attention이 분산됨. 

 

Hymba Model Family

- Hymba-125M/350M/1.5B model을 학습

- DCLM-Baseline-1.0(web에서 추출한 고품질 pretraining 데이터셋), SmoLM-Corpus(고품질 교육용/합성 텍스트 데이터셋), 내부 고성능 데이터셋

- Warmup-Stable-Decay learning-rate scheduler 사용(처음에 warmup, 일정 구간 stable, 마지막 decay), data annealing 사용(학습 후반부에 더 품질 좋은 데이터의 비중을 높임)

- Maximum sequence length도 2K로 쓰다가, 마지막 100B tokens 구간에서는 8k까지 늘림. 긴 문맥 처리 능력을 강화함.

 

Model Evaluations

Baselines(quadratic attention을 사용한 모델): (small language models) Llama 3.2, SmolLM, SmolLM2, AMD-OLMo, StableLM, OLMo, Cosmo, Phi-1.5, H2O-Danube, OpenELM, MiniCPM, etc, (hybrid models) Rene etc.

 

Benchmark settings: (1) 학습한 Hymba를 SOTA small LM들과 직접 비교 (2) 서로 다른 아키텍처를 같은 데이터셋, 같은 수의 layer, model size, training recipes를 통일시킨 후 비교.(성능 차이가 구조에서 비롯된 것인지)

 

Benchmark tasks: Commonsense reasoning, recall-intensive tasks, instruction tuned 모델에서는 수학이나 function calling, role playing과 같은 downstream tasks도 비교.

1.5T pretraining tokens만 사용했음에도 불구하고, Hymba-1.5B 모델이 2B 아래 모델들 중 최고의 성능을 보임.

sub-2B 모델들 중 가장 강력한 성능을 보였던 SmolLM2 모델이 11T tokens를 사용해서 학습한 것에 비해 Hymba는 1.5T 토큰으로 학습했고, SmolLM의 성능도 능가했음.

Llama-3.2-3B 보다도 평균 성능이 더 높음.

 

다른 아키텍처들을 같은 레시피와 데이터로 처음부터 학습한 결과 Hymba 아키텍처의 성능이 대체적으로 가장 좋았음.

Mamba가 Recall intensive tasks에서 약한 성능을 보이는 가운데, attention 헤드도 사용한 Hymba 모델은 recall task에서 잘했음.

 

긴 컨텍스트 속에서 필요한 정보를 recall하는 needle in the haystack 문제도 Mamba2나 Llama 모델보다 더 잘 푸는 모습을 보였음.

Mamba2는 컨텍스트 끝에 있을 때 정보를 잘 찾지만, 시작이나 중간에 있는 경우에는 잘 못찾음. Llama3는 오히려 중간에 있는 needle을 잘 못찾음.

 

Instruction-tuned Model

Hymba-1.5B base model을 instruction model로 post training.

1. Instruction 데이터로 Full Fine tuning -> 2. Direct Preference Optimization으로 사람이 선호하는 답변을 더 내놓을 수 있게 조정.

 

비교 모델들: Llama-3.2-1B-Instruct, OpenELM-1.1B-Instruct, Qwen2.5-1.5B-Instruct, SmolLM-1.7B-Instruct

 

평가 downstream tasks:

MMLU 5-shot	다양한 지식/추론 능력
IFEval	instruction following 능력
GSM8K 5-shot	초등/중등 수준 수학 추론
GPQA 0-shot	어려운 과학·전문 지식 추론
BFCLv2	function calling 능력

 

Hymba 모델이 평균적으로 가장 높은 점수를 받았음. 이는 수학, 추론, function calling 문제를 잘 푸는 것을 증명.

 

DoRA를 이용한 role-playing fine-tuning에서는 7B 모델보다도 좋은 일반화 성능을 보여줬음.

 

Conclusion

Hymba는 작은 언어 모델을 위한 새로운 hybrid-head architecture를 제안. 핵심은 한 layer 안에서 attention head와 SSM head를 함께 사용하는 것.

 

Attention은 긴 문맥에서 특정 정보를 정확히 찾아오는 데 강하고, SSM은 긴 문맥을 효율적으로 요약하는 데 강하다. Hymba는 이 둘을 결합해 recall 성능과 추론 효율을 동시에 얻으려 한다.

 

또한 입력 앞에 학습가능한 meta token을 붙여, BOS token에 낭비되던 attention을 줄이고 모델이 중요한 정보에 더 잘 집중하도록 만든다. 이 meta token은 attention과 SSM 모두에서 일종의 learned cache처럼 작동한다.

 

실험 결과 Hymba는 다양한 small LM benchmark에서 높은 정확도와 좋은 throughput, 낮은 KV cache 사용량을 보였다.

 

Hymba는 단순히 새로운 모델 하나를 제안한 것이 아니라, 앞으로 efficient language model을 설계할 때 attention과 SSM을 어떻게 결합할 것인가에 대한 중요한 방향성을 보여주는 연구라고 볼 수 있음.