기술 요약: 예측기 게이팅 기반 뱅크 단위 희소성을 통한 지속적 LLM 업사이클링 (Continual LLM Upcycling via Predictor-Gated Bank-Wise Sparsity)

문제 정의

대규모 언어 모델(LLM), 특히 디코더 전용 트랜스포머(decoder-only Transformers)의 계산 비용은 파라미터 크기와 밀집된 피드포워드 네트워크(FFN)에 필요한 토큰당 계산량에 의해 점점 더 지배되고 있습니다. 현대적인 SwigLU 아키텍처에서 FFN은 은닉 상태(hidden state)를 큰 중간 차원으로 확장하고, 이를 처리한 후 다시 투영합니다. 긴 문맥(long-context) 학습 및 추론 과정에서, 이 밀집된 계산은 특정 토큰에 대해 일부 중간 채널만이 중요함에도 불구하고 모든 토큰에 대해 적용됩니다.

기존 솔루션들은 다음과 같은 한계에 직면해 있습니다:

• Mixture-of-Experts (MoE): 토큰을 특정 전문가 풀로 라우팅하여 계산량을 줄이지만, 별도의 파라미터 풀, 복잡한 부하 분산 요구 사항, 그리고 all-to-all 통신 오버헤드를 유발합니다.
• 사후 희소성 (Post-hoc Sparsity): 네이티브 FFN 활성화로부터 희소 마스크를 도출하는 방식은 지배적인 게이트(gate)와 업-프로젝션(up-projection) 계산을 건너뛰기에는 너무 늦게 도달하므로, 하드웨어 수준의 가속을 제공하지 못합니다.
• 문맥적 희소성 (Contextual Sparsity): 이전 연구들은 LLM 계산이 입력 의존적임을 보여주었으나, 훈련이 없는(training-free) 방식은 복잡한 생성 작업에서 취약할 수 있습니다.

본 논문은 **지속적 학습(continual training, upcycling)**을 통해 사전 훈련된 밀집(dense) LLM을 채널 희소 모델로 변환하면서도 성능 저하를 최소화하며, 라우팅 신호가 FFN 계산 전에 가용하도록 하여 하드웨어 지향적인 희소 경로를 만드는 문제를 다룹니다.

방법론

제안된 접근 방식인 **예측기 게이팅 뱅크 단위 희소성(Predictor-Gated Bank-Wise Sparsity)**은 특정 지속적 학습 레시피를 통해 Qwen2.5-8B 모델을 희소 변형 모델로 변환합니다.

1. 아키텍처: 예측기 게이팅 희소 SwigLU (Predictor-Gated Sparse SwiGLU)

네이티브 게이트 출력을 사용하여 희소성을 결정하는 대신, 저자들은 FFN 계산 전에 토큰 은닉 상태($x_{l,t}$)에 작용하는 **저차원 예측기(low-rank predictor, $P_l$)**를 도입했습니다.

• 예측기 (Predictor): 은닉 상태를 FFN 중간 채널에 대한 로짓(logits)으로 매핑하는 경량 2계층 모듈($R^d \to R^{448} \to R^{d_i}$)입니다.
• 라우팅 메커니즘: 예측기는 로짓을 출력하며, 이는 표준 밀집 SwigLU 게이트 $\sigma(u_{l,t})$를 대체하여 라우팅 신호를 생성합니다.
• 뱅크 단위 Top-k 규칙 (Bank-wise Top-k Rule): 중간 차원($d_i = 14336$)을 크기가 64인 고정된 뱅크들로 분할합니다. 각 뱅크에 대해 하드 top-k 규칙($k=16$)을 적용합니다. 이는 64채널 뱅크당 16개의 활성 채널을 유지하여, FFN 중간 활성화 폭($d_s = 3584$)에서 4배의 희소성을 달 achieves 합니다.
• 실행 (Execution):
  • 순전파 (Forward Pass): 하드 이진 마스크가 적용됩니다. 비활성 채널은 정확히 0이 됩니다.
  • 역전파 (Backward Pass): 미분 불가능한 top-k 선택은 선택 임계값 주변의 **부드러운 대리 그래디언트(soft surrogate gradient, sigmoid 형태의 로컬 그래디언트)**를 통해 처리되어, 언어 모델링 손실(loss)이 예측기 파라미터를 업데이트할 수 있도록 합니다.

2. 학습 레시피 (Dense-to-Sparse Continual Training)

이 프로세스는 8K 컨텍스트로 훈련된 Qwen2.5-8B 모델을 32K 희소 모델로 변환하기 위해 네 가지 뚜렷한 단계를 거칩니다:

1. Dense 8K 단계: 표준 Qwen2.5-8B 아키텍처를 사용하여 처음부터 학습합니다.
2. Dense 32K 단계: 8K 체크포인트에서 재개하여 컨텍스트를 32K로 확장하고, 약 3T 토큰 동안 밀집 학습을 계속합니다.
3. 예측기 웜업 (Predictor Warm-up): 네이티브 게이트 신호를 모방하는 보조 브랜치로서 저차원 예측기를 도입합니다. 이 브랜치는 메인 손실에 영향을 주지 않습니다.
4. 예측기 게이팅 밀집 FFN (Predictor-Gated Dense FFN): 예측기(sigmoid 사용)를 통해 메인 FFN 게이트를 라우팅하여 예측기가 손실에 영향을 미치게 하되, 아직 하드 top-k 선택은 적용하지 않습니다.
5. 예측기 게이팅 Top-k 희소 FFN (Predictor-Gated Top-k Sparse FFN): 4배 희소성을 강제하기 위해 뱅크 단위 top-k 선택을 활성화합니다.

참고: 희소 브랜치는 32K 단계의 후반부(약 3T 토큰 이후)에 도입되어, 밀집 모델이 희소 모델로 변환되는 실제적인 업사이클링 시나리오를 반영합니다.

3. 하드웨어적 함의

학습은 기능적 시뮬레이션으로서 GPU 상에서 수행됩니다. 현재의 GPU 구현은 마스킹 전의 밀집 게이트 및 업-프로젝션을 실체화하므로 즉각적인 벽시계 시간(wall-clock) 가속에는 제한이 있지만, 알고리즘적 원리는 전용 희소 하드웨어가 비활성 채널에 대한 모든 세 가지 FFN 단계(게이트, 업, 다운 프로젝션)를 건너뛰도록 하여 이론적으로 유효 중간 폭을 4배 줄일 수 있게 합니다.

주요 기여

1. 예측기 게이팅 채널 희소성 패턴: 저차원 예측기가 FFN 계산 전에 채널 라우팅 로짓을 생성하고, 하드 뱅크 단위 top-k 규칙과 결래된 새로운 아키텍처입니다. 이는 전용 커널이나 하드웨어가 사용될 경우 비활성 게이트, 업, 다운 채널을 건너뛰는 것을 가능하게 합니다.
2. Dense-to-Sparse 지속적 학습 레시피: Qwen2.5-8B 밀집 모델(8K 학습, 32K 지속 학습)을 성능 저하를 최소화하면서 예측기 게이팅 4배 희소 FFN 모델로 변환하는 실용적인 워크플로를 제공합니다.
3. 학습 교훈 및 실패 모드:
   • 채널 균형 (Channel Balancing): 본 논문은 MoE 스타일의 전문가 균형(expert balancing) 편향이 뱅크 단위 채널 라우팅으로 깔끔하게 전이되지 않음을 입증합니다. 명시적으로 채널 균형을 강제하는 것(온라인 업데이트를 통해)이 균형을 맞추지 않는 것보다 해롭거나 중립적이라는 것을 발견했습니다.
   • 계층 국소적 긴 문맥 실패 (Layer-Local Long-Context Failure): RULER-CWE 평가 결과, 짧은 문맥 벤치마크에서는 보이지 않았던 희소 모델의 급격한 성능 하락(12K–16K 구간)이 발견되었습니다. 저자들은 이를 계층 국소적 실패(특히 7번째 레이어)로 식과하였으며, 영향을 받는 길이 범위에서 성능을 상당히 회복시키는 단일 레이어 복구 알고리즘(밀집 폴백 개입)을 제안했습니다.

결과

희소 모델은 두 가지 베이스라인과 비교 평가되었습니다: Dense Companion (동일 예산의 완전 밀집 모델) 및 Dense-4x 베이스라인 (사후에 활성화 크기에 따라 4배 희소성을 적용한 모델).

• 성능 근접성: 광범로한 벤치마크(사실적 QA, 코드, 추론, 수학, 종합) 전반에서 예측기 게이팅 희소 모델은 나이브한 Dense-4x 베이스라인보다 밀집 컴패니언 모델에 일관되게 더 가까운 성능을 보였습니다.
• 구체적 이점:
  • 희소 모델은 MBPP, MultiPL-E, MMLU에서 밀집 베이스라인과 대등하거나 약간 상회하는 성능을 보였습니다.
  • Chinese SimpleQA, MBPP+, MBXP, KOR-Bench에서 나타나는 Dense-4x의 성능 저하를 실질적으로 줄였습니다.
• 긴 문맥 성능:
  • RULER 벤치마크(4K, 8K, 16K)에서 희소 모델은 밀집 모델과 유사한 성능을 보였습니다.
  • 32K에서 희소 모델(80.4)은 Dense-4x 베이스라인(73.5)보다 우수했으며, 밀집 컴패니언(79.2)을 약간 상회했습니다.
  • RULER-CWE 분석: 희소 모델은 12K에서 16K 사이에서 Recall(재현율)의 "절벽" 현상을 보였습니다. 제안된 단일 레이어(Layer 7) 복구를 적용했을 때, 평균 Recall은 8.7K–12.2K 범위에서 0.856에서 0.947로, 12.8K–16.3K 범위에서 0.583에서 0.802로 개선되었습니다.

의의 및 주장

본 논문은 희소 업사이클링이 손실이 없는(lossless) 것은 아니지만, 밀집 체크포인트으로부터 하드웨어 지향적인 희소 모델을 구축하기 위한 실행 가능한 경로임을 주장합니다.

• 학습 중심 접근법: 사후 압축과 달리, 이 방법은 채널 희소성을 지속적 학습 중에 최적화되는 언어 모델링 경로의 필수적인 부분으로 취급합니다.
• 하드웨어 정렬: 라우팅 예측기를 FFN 계산 이전에 배치하고 하드 마스크를 사용함으로써, 이 방법은 단순한 이론적 희소성을 넘어 진정한 가속을 위한 전용 희소 하드웨어의 요구 사항에 부합합니다.
• 진단적 가치: 본 연구는 광범위한 짧은 문맥 벤치마크가 놓칠 수 있는 실패 모드(계층 국소적 절벽 등)를 식별하는 데 긴 문맥 평가가 얼마나 중요한지를 강조합니다. 또한, 지속적 학습 설정에서 채널 수준 라우팅에 대한 MoE 균형 전략의 전이 가능성에 대한 부정적인 결과(negative result)를 제공합니다.

저자들은 희소 모델이 모든 시나리오에서 밀집 모델을 완벽하게 대체할 수는 없지만, 관리 가능한 성능 트레이드오프를 통해 밀집 모델을 효율적인 희소 변형 모델로 업사이클링하는 실용적인 레시피를 제공한다고 결론지었습니다.