더 작고 강해진 Kanana-2 SLM 개발기

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

Kanana-2 시리즈는 온디바이스 제약을 겨냥해 Dense 3B 모델을 시작점으로 단계적 프루닝과 Teacher 기반 디스틸레이션을 결합하여 3B·1.3B·0.9B SLM을 생성했고, 히든 차원 프루닝에 PCA 회전을 적용해 주요 표현을 보존하면서 구조를 효율적으로 축소했다. 온디바이스 추론 병목을 줄이기 위해 Sliding Window Attention 레이어와 Full Attention 레이어를 3:1로 교차 배치한 하이브리드 구조를 적용해 32K 문맥에서 KV Cache를 최대 72.7% 절감했고 Window size 1024는 Full Attention과 대등한 성능을 보였다. Post-training 단계에서는 Staged SFT로 도메인별 Expert를 얻고 SCE 방식으로 병합한 뒤 DAPO 기반 RL을 적용해 응답 다양성과 안정성을 확보했으며, 임베딩은 Bidirectional+Mean Pool 초기화와 HN-only 손실, Temperature=0.02 설정으로 한국어·영어에서 우수한 성능과 차원 효율을 달성했다. 전체 파이프라인은 Proxy Token Scale 기반의 LR 스케일링, PCA 기반 프루닝, Cascade/Elastic 훈련 옵션을 통해 성능과 비용의 균형을 맞추는 방향으로 설계되어 실무적 재현 가능성과 온디바이스 실용성을 동시에 목표로 한다.

빠른 이해

새로운 점

PCA 기반 히든 차원 프루닝과 SWA 하이브리드 배치의 결합으로 소형 온디바이스 모델에서 성능 손실을 최소화하면서 KV Cache를 크게 절감한 점

핵심 메커니즘

입력으로는 Dense 3B Base 모델과 Teacher rollout·캘리브레이션 데이터가 주어지고, 처리 과정에서는 PCA 회전을 통한 히든 차원 축소와 단계적(Cascade) 프루닝·디스틸레이션, SWA 하이브리드 배치로 KV Cache를 제한하며 Post-training에서 DAPO 기반 RL로 정책을 정교화한 뒤 출력으로는 3B·1.3B·0.9B SLM과 고효율 임베딩 모델이 산출된다.

핵심 수치

Pre-training 토큰량: Stage-1 7.5T + Stage-2 2T
SWA 하이브리드(Kanana-2-1.3B, 32K): KV Cache 최대 -72.7%- SWA:Full=3:1 대비 Full-only
소형 모델 Distillation 토큰: 각 단계별 약 300B- 2B→1.3B→0.9B 각 단계에 대해 300B 토큰으로 Distillation 수행

섹션별 상세

프로젝트 개요와 목표

Kanana LLM 팀은 온디바이스 환경에서 동작하는 효율적인 Small Language Model을 목표로 Kanana-2 시리즈(3B, 1.3B, 0.9B)를 개발했다. 이번 시리즈는 Dense 3B 모델을 시작점으로 하여 단계적 프루닝과 디스틸레이션, Long Context 학습, 그리고 온디바이스 최적화 기법을 결합해 성능과 메모리 효율을 동시에 개선하는 방향으로 설계되었다. 공개 대상은 Kanana-2-3B와 Kanana-2-1.3B의 Base 및 Instruct 모델로, 연구와 서비스 적용을 위한 재현 가능성과 실용성을 모두 고려했다.

사전학습 전략과 하이퍼파라미터 탐색

Kanana-2-3B는 TPU v5e 클러스터에서 from-scratch로 2단계(7.5T 토큰 + 2T 토큰) 사전학습을 수행한 뒤 GPU로 옮겨 Distillation을 진행했다. 제한된 예산으로도 본 학습 스케일에 맞는 학습률을 얻기 위해 100B 토큰의 Proxy Token Scale에서 다양한 LR 후보를 탐색하고, LR 스케일링 법칙 LR*(D_target) ≈ LR*(D_proxy) · (D_target/D_proxy)^{-β} (β=0.32, D_proxy=100B, D_target=7.5T)을 적용해 타깃 학습률을 산출했다. 이러한 접근은 대규모 환경에서 직접 탐색할 수 없는 하이퍼파라미터를 안정적으로 이전하는 실용적 방안으로 작동했다.

근거

사전학습은 TPU 기반으로 Stage-1 7.5T 토큰, Stage-2 2T 토큰의 2단계로 진행되었고 Proxy D_proxy=100B와 β=0.32를 사용해 학습률을 스케일링했다. — TPU 기반 Pre-Training 섹션의 토큰 수와 Token Horizon LR 스케일링 공식

Teacher 기반 Distillation과 Long Context 학습

Megatron-LM 기반 워크플로에서 Kanana-2-30B-A3B-Instruct-2601을 주된 Teacher로 사용해 학생 모델을 디스틸레이션했고, Instruct 계열 Teacher가 수학·코드 같은 특정 영역에서 더 나은 학생 성능을 유도함을 실험으로 확인했다. Long Context 확장을 위해 YaRN을 적용해 4K에서 32K로 문맥 길이를 늘리고, Mid-training 단계에서 LR decay와 추가 데이터로 학습을 안정화했다. Long Context에서는 일부 한국어 장문 데이터의 노이즈를 줄이고 수학·코드·영어 데이터 비중을 늘려 RULER-32K 성능을 개선하는 데이터 재구성 전략을 적용했다.

프루닝 및 구조 축소 설계

소형 모델(1.3B, 0.9B)은 3B Base에서 단계적으로 Hidden dimension과 MLP intermediate dimension을 중심으로 폭(Width) 프루닝을 적용해 구조를 축소했다. 레이어 수와 쿼리 헤드는 유지하되 Hidden/Intermediate 차원을 줄이는 폭 기반 전략이 레이어 수를 줄이는 방식보다 동일 파라미터 예산에서 더 우수한 성능을 보였고, Attention head 축소는 성능 저하를 크게 유발해 최소화했다. 프루닝 구성 후보들을 비교한 뒤 Cascade 방식으로 3B→2B→1.3B→0.9B 순으로 단계적 디스틸레이션을 수행해 초기 성능과 수렴 속도를 개선했다.

PCA 기반 히든 차원 프루닝 파이프라인

히든 차원 프루닝은 레이어별 RMSNorm 입력 활성화 통계를 캘리브레이션 데이터로 수집하고, PCA로 전역 회전 행렬을 계산해 임베딩·어텐션·MLP 투영에 일관되게 회전을 적용한 뒤 차원을 축소하는 절차로 진행되었다. 회전된 모델에서 다시 활성화를 수집해 Attention head 및 MLP intermediate dimension의 중요도를 산출하고 순차적 프루닝을 수행했으며, Layer-wise 집계 방식을 채택해 레이어별 특성을 반영했다. 동일한 파라미터 예산에서 PCA 기반은 기존 인덱스 기반 방식보다 초기 및 전체 학습 구간에서 일관되게 높은 성능을 보였고, 특히 Attention capacity가 제한된 후보에서 우위가 두드러졌다.

python

def prune(model, target_config):
    # Activations are cached during calibration forward passes.
    # Norm inputs are collected before rotation, while attention and MLP
    # activations are collected after rotation.
    # Hidden dimension pruning
    norm_inputs = collect_norm_inputs(model)
    rotation = PCA( norm_inputs, n_components=target_config.hidden_dim, )
    model = apply_rotation( model, rotation, target_dim=target_config.hidden_dim, )
    # Attention head pruning
    head_importance = sum( score_query_heads(layer.attn) for layer in model.layers )
    heads_to_keep = select_topk_heads_per_kv_group( head_importance, k=target_config.query_heads, )
    for layer in model.layers:
        layer.attn = prune_query_heads( layer.attn, heads_to_keep, )
    # MLP intermediate dimension pruning
    for layer in model.layers:
        intermediate = ( silu(layer.mlp.gate_proj.output_x) * layer.mlp.up_proj.output_x )
        importance = aggregate_intermediate_importance(intermediate)
        dims_to_keep = topk( importance, k=target_config.intermediate_dim, )
        layer.mlp = prune_intermediate_dims( layer.mlp, dims_to_keep, )
    return model

이 코드는 PCA 기반 회전을 적용한 뒤 Attention 헤드와 MLP 중간 차원을 순차적으로 프루닝하는 파이프라인의 핵심 로직을 예시한 파이썬 형태의 의사코드이다.

온디바이스 최적화와 Sliding Window Attention

온디바이스 환경의 핵심 병목은 시퀀스 길이에 비례해 증가하는 KV Cache의 메모리 읽기량이며, 이를 완화하기 위해 Sliding Window Attention(SWA)을 도입해 각 토큰의 어텐션 범위를 고정된 윈도우로 제한했다. SWA 레이어와 Full Attention 레이어를 3:1 비율로 교차 배치한 하이브리드 구조를 채택해 전체 레이어의 3/4에서 KV Cache 크기와 토큰당 읽기량을 상수로 유지했고, 1.3B 모델 기준으로 32K 문맥에서 KV Cache를 최대 72.7% 절감했다. Window size는 성능과 메모리 절감의 트레이드오프를 결정하며, 실험에서 Window 1024는 대부분 벤치마크에서 Full Attention과 유사한 성능을 유지해 온디바이스 장문 처리에서 실용적 해법으로 검증되었다.

Post-Training, RL 및 임베딩 설계

3B 모델은 Staged SFT로 도메인별 Expert를 얻은 뒤 SCE 방식으로 가중치를 병합하고 Calibration을 거쳐 Instruct 모델을 확보했으며, RL 단계에서는 GRPO 대신 DAPO를 채용해 Dynamic Sampling, Clip-Higher, Token-level PG Loss와 Overlong Penalty를 결합해 학습 안정성과 다양성을 확보했다. 1.3B/0.9B의 Post-training은 Teacher rollout 기반의 Off-Policy Distillation으로 시작해 Forward KL을 사용한 On-Policy Distillation으로 이어지며 RL 전의 탐색 여지를 유지하는 설계를 택했다. 임베딩은 Decoder-only 모델로부터 전이해 Bidirectional Attention+Mean Pool 초기화를 채택하고 HN-only 손실과 Temperature=0.02, Hard negative 9개 설정으로 학습해 한국어·영어 모두에서 높은 성능과 차원 효율을 달성했다.

근거

SWA:Full=3:1 하이브리드 구조는 Kanana-2-1.3B 모델에서 32K 문맥 기준 KV Cache를 최대 72.7% 절감했다. — KV Cache 크기 비교 표(Architecture별 Length별 MiB 수치)

용어 해설

PCA 기반 히든 차원 프루닝(PCA-based Hidden Dimension Pruning): — 레이어별 활성화 통계를 모아 전역 회전 행렬을 계산하고, 토큰 임베딩·어텐션·MLP 투영에 일관되게 회전을 적용한 뒤 차원을 축소하여 주요 표현을 보존하면서 파라미터를 줄이는 방식으로, 프루닝으로 인한 정보 손실을 줄이고 압축 모델의 성능 회복 속도를 개선한다.
슬라이딩 윈도우 어텐션(Sliding Window Attention): — 각 토큰이 고정된 윈도우 범위 내에서만 어텐션을 수행하도록 하여 KV Cache 길이를 윈도우 크기로 제한함으로써 토큰당 메모리 읽기량을 시퀀스 길이와 무관하게 일정하게 유지해 온디바이스 추론 시 메모리와 레이턴시 병목을 완화하는 지역적 어텐션 기법이다.
단계적(카스케이드) 프루닝 및 디스틸레이션(Cascade Pruning & Distillation): — 대형 모델에서 작은 모델로 한 번에 압축하지 않고 중간 규모의 서브모델을 차례로 생성하면서 각 단계에서 프루닝과 교사 모델 기반 디스틸레이션을 반복하여 최종 소형 모델의 초기 성능을 높이고 학습 안정성을 확보하는 압축 파이프라인이다.
Decoupled Clip and Dynamic Sampling Policy Optimization(DAPO): — 동적 샘플링으로 유효 샘플만 배치에 포함시키고 클립 범위를 확장해 정책 변화로 인한 엔트로피 붕괴를 방지하며 토큰 단위의 정책 그래디언트와 과긴 출력에 대한 패널티를 결합해 RL 학습의 효율과 안정성을 개선하는 최신 강화학습 최적화 기법이다.
토큰 지평선 학습률 스케일링(Token Horizon LR Scaling): — 작은 토큰 규모(Proxy)에서 탐색한 최적 학습률을 타깃 토큰 규모로 변환하기 위한 스케일링 법칙으로, LR*(D_target) ≈ LR*(D_proxy) · (D_target / D_proxy)^{-β} 공식을 사용해 긴 토큰 지평선에서 안정적인 학습률 설정을 확보한다.

기술

TPU v5e
Megatron-LM
YaRN
PCA
Sliding Window Attention
DAPO

언급된 리소스

문서원문 링크