기술 요약: CausalGate

문제 정의

현재의 LLM(대규모 언어 모델)용 적응형 추론 방식은 은닉 상태 유사성(hidden-state similarity)이나 활성화 크기(activation magnitude)와 같은 관찰 기반의 휴리스틱에 크게 의존하고 있습니다. 저자들은 이러한 상관관계 기반 지표가 의미론적 정확성에 필수적인 미묘하고 비선형적인 구조적 연산을 포착하는 데 실패할 수 있다고 주장합니다. 특정 모듈은 은닉 상태에서 기하학적 변위가 미미할지라도, 사실적 정확성이나 논리적 일관성을 위해 필요한 결정적인 "인과적 정제(causal refinement)"를 수행할 수 있습니다. 결과적으로, 휴리스틱 기반의 정적 프루닝(pruning)은 이러한 조용히 중요한 모듈들을 건너뛰게 되어, 연산량을 줄임에 따라 급격한 성능 저하를 초래할 위험이 있습니다. 또한, 기존의 동적 라우팅 프레임워크(예: CALM, Mixture-of-Depths)는 런타임 오버헤드와 분기 예측(branch prediction) 페널티를 유발하여 효율적인 하드웨어 배포를 저해합니다.

방법론: CausalGate

CausalGate는 영구적이고 오버헤드가 없는 모듈 선택 계층을 구축하기 위해 설계된 오프라인 개입 가이드(intervention-guided) 프레임워크입니다. 런타임 표현을 관찰하는 방식과 달리, CausalGate는 보정(calibration) 단계 동안 모델 구조에 능동적으로 개입하여 진정한 인과적 필요성을 결정합니다.

1. 개입 기반 인과적 중요도 추정

본 프레임워크는 각 Attention 및 MLP 서브 레이어를 독립적인 계산 단위로 취급합니다. 보정 과정 중:

• 개입(Intervention): 주어진 입력 시퀀스에 대해, 특정 모듈 $m_i$의 출력은 0으로 설정(zeroed out)되는 반면 나머지 모든 모듈은 변경되지 않은 상태를 유지합니다.
• 측정: 시스템은 원래 모델의 로짓 분포와 개입된 분포 사이의 쿨백-라이블러 발산($D_{KL}$)을 측정합니다.
• 가설: 큰 $D_{KL}$ 값은 해당 모듈이 최종 예측에 인과적으로 필수적임을 나타내는 반면, 작은 값은 해당 모듈을 안전하게 건너뛸 수 있음을 의미합니다. 이 과정은 세밀한 모듈 수준의 인과적 중요도 맵을 생성합니다.

2. 인과적 중요도 증류(Distillation)

원시 개입 점수는 훈련 반복 과정에서 변동될 수 있으므로, CausalGate는 안정적인 추론을 보장하기 위해 이러한 신호를 정적이고 학습 가능한 스칼라 게이트($g_i$)로 증류합니다.

• 대상 집계: 노이즈를 필터링하고 안정적인 글로벌 모듈 관련성 순위를 확립하기 위해, 개입 점수는 지수 이동 평균(EMA)을 사용하여 정규화 및 평활화됩니다.
• 손실 함수: 학습 목표는 다음 네 가지 구성 요소를 결합합니다:
  1. 언어 모델링 손실 ($L_{LM}$): 일반적인 모델 능력을 유지합니다.
  2. 인과적 회귀 손실 ($L_{MSE}$): 게이트가 EMA로 평활화된 인과적 중요도 타겟을 일치시키도록 학습합니다.
  3. 쌍별 랭킹 손실 ($L_{rank}$): 학습된 게이트가 개입 과정에서 발견된 상대적 인과 순위를 엄격하게 보존하도록 강제하는 미분 가능한 손실입니다.
  4. 희소성 정규화 항 ($L_{sparse}$): 프루닝을 극대화하기 위해 더 작은 게이트 활성화를 장려합니다.

3. 제로 오버헤드 추론

학습 후, 학습된 게이트는 결정론적이고 시퀀스 불가지론적인(sequence-agnostic) 모듈 중요도 순위를 제공합니다.

• 정적 마스킹: 목표 연산 감소 비율 $\rho$에 대해, 가장 높은 게이트 값을 가진 상위 $K$개의 모듈은 유지되고 나머지는 영구적으로 비활성화됩니다.
• 실행: 이는 런타임 라우팅 로직, 분기 예측 페널티, 또는 동적 제어 흐름이 없는 직선형(straight-line) 추론 경로를 생성하여 효율적인 하드웨어 실행을 가능하게 합니다.

주요 기여

1. 인과적 프루닝 패러다임: 본 논문은 관찰 기반 휴리스틱을 구조적 인과 개입(do-operations)에서 직접 도출된 그라운드 트루스(ground-truth) 중요도 계층으로 대체하는 구조적 프루닝 접근 방식을 도입합니다.
2. 서브 레이어 단위의 세밀함: 프레임워크는 단일 레이어 수준이 아닌 개별 Attention 및 MLP 서브 레이어 수준에서 작동합니다. 이를 통해 서로 다른 기능(컨텍스트 라우팅 vs 국소적 메모리)을 독립적으로 프루닝할 수 있어, 균일한 레이어 드롭핑보다 더 큰 유연성을 제공합니다.
3. 증류 프레임워크: 저자들은 EMA 평활화와 쌍별 랭킹 손실을 사용하여 복잡한 인과 계층을 정적 스칼라 게이트로 인코딩하는 최적화 전략을 제안하여 제로 오버헤드 추론을 가능하게 합니다.
4. 실증적 검증: 광범한 평가를 통해 CausalGate가 높은 연산 감소 수준에서도 언어 모델링 품질과 상식 추론 능력을 기존의 동적 라우팅 및 레이어 스킵 베이스라인보다 훨씬 더 잘 보존함을 입증했습니다.

실험 결과

저자들은 TinyLlama-1.1B, Qwen2.5-3B, Llama-3.1-8B를 대상으로 언어 모델링(WikiText-2, C4) 및 상식 추론(HellaSwag, PIQA, CSQA, WinoGrande) 벤치마크에서 CausalGate를 평가했습니다.

• 성능 유지: CausalGate는 테스트된 모든 연산 예산(5% ~ 40% 모듈 제거)에서 주요 베이스라인(CALM, MoD, GateSkip, AdaSkip 포함)보다 일관되게 우수한 성능을 보였습니다.
  • 20% 제거 예산에서의 TinyLlama-1.1B에 대해, CausalGate는 WikiText-2 퍼플렉시티(perplexity) 67.0을 달이트했으며, 이는 MoD(407.0)나 AdaSkip(1816.0)과 같은 베이스라인보다 현저히 낮은 수치입니다.
  • 40% 제거 예산에서 CausalGate는 WikiText-2 퍼플렉시티 1365.0을 유지했으나, 많은 베이스라인(예: Act-Norm은 20,825.4에 도달)은 성능이 붕괴되었습니다.
• 하드웨어 효율성: 정적인 프루닝 특성은 이론적인 연산 절감을 실제 하드웨어 레이턴시 감소로 전환했습니다. TinyLlama-1.1B에서 20% 모듈 제거는 1.20배의 속도 향상(레이턴시 2.100초에서 1.747초로 감소)과 처리량(throughput) 증가(60.9에서 73.3 tokens/s)를 가져왔습니다.
• 확장성: 이 방법은 더 큰 백본(Qwen2.5-3B 및 Llama-3.1-8B)에서도 효능을 입증하였으며, 10% 및 20% 제거율에서 베이스라인 대비 우월한 퍼플렉시티 점수를 유지했습니다.

의의 및 주장

본 논문은 CausalGate가 모델 압축의 경계를 런타임의 동적 결정에서 컴파일 타임의 정적 최적화로 근본적으로 전환한다고 주장합니다. 모듈의 중요성이 최종 로짓 분포에 미치는 인과적 영향에 의해 정의된다는 원칙을 실행함으로써, 이 프레임워크는 관찰 기반 대리 지표의 "추측"을 피합니다.

저자들은 자신들의 접근 방식이 인과 분석과 하드웨어 실행 효율성 사이의 가교 역할을 한다고 강조합니다. 시퀀스 불가지론적인 중요도 계층을 구축함으로써, CausalGate는 동적 라우팅 프레임워크와 관련된 운영 오버헤드를 제거하여, 국소적 메모리 경계 및 표준 하드웨어 루프와 매우 잘 결합되는 직선형 추론을 보장합니다. 결과는 공격적인 프루닝이나 동적 스키핑과 흔히 연관되는 심각한 성능 저하 없이도 상당한 연산 감소를 달성할 수 있음을 시사합니다.