LayerRoute: Input-Conditioned Adaptive Layer Skipping via LoRA Fine-Tuning for Agentic Language Models

기술 요약: LayerRoute – LoRA 미세 조정을 통한 입력 조건부 적응형 레이어 스킵(Layer Skipping)

문제 정의

현대의 에이전틱 AI 시스템은 두 가지 뚜렷한 단계 유형으로 구성된 이질적인 워크플로우를 통해 대규모 언어 모델(LLM)을 오케스트레이션합니다:

1. 구조화된 도구 호출(Structured tool calls): 낮은 퍼플렉시티(perplexity)를 가진 짧고 결정론적인 시퀀스 (예: 데이터베이스 쿼리).
2. 개방형 계획/추론(Open-ended planning/reasoning): 다단계 추론을 필요로 하는 길고 복잡하며 높은 퍼플렉시티를 가진 시퀀스.

현재의 추론 시스템은 모든 입력에 대해 동일한 계산 자원을 적용하여 모든 트랜스포머 레이어를 통과하도록 합니다. 이는 비효식적입니다. 왜냐하면 도구 호출의 경우 초기 잔차 스트림(residual stream) 상태에서 높은 확신도를 보이기 때문에, 깊은 레이어들이 최종 예측에 기여하는 바가 미미하다는 점이 시사되기 때문입니다. 본 논문은 모델 품질을 희생하지 않으면서 입력 복잡도에 따라 계산량을 동적으로 할당할 수 있는 추론 시스템의 필요성을 다룹니다.

방법론

저자들은 Qwen2.5-0.5B-Instruct 모델(24개 레이어, 은닉 크기 896)을 위해 설계된 경량 어댑터인 LayerRoute를 제안합니다. 이 아키텍처는 동결된 백본(backbone)에 각 트랜스포머 블록당 두 가지 주요 구성 요소를 추가합니다:

1. 레이어별 라우터(Per-Layer Routers): 각 블록 $i$는 소프트 시그모이드 값 $\sigma_i$를 출력하는 경량 선형 레이어($w_i^T \bar{h}_i + b_i$)를 포함합니다. 이는 임계값 0.5를 통해 하드 이진 게이트 $g_i \in \{0, 1\}$로 변환됩니다.

   • 순전파(Forward Pass): 만약 $g_i=1$이면, 해당 블록은 정상적으로 실행됩니다(LoRA 어댑터 활성화 상태). 만약 $g_i=0$이면, 은닉 상태는 변경 없이 그대로 통과하며, 해당 블록에 대한 계산량은 발생하지 않습니다.
   • 학습(Training): 비미분 가능한 하드 게이트는 **스트레이트스루 추정기(Straight-Through Estimator, STE)**를 통해 처리됩니다. 이를 통해 역전파 과정에서 그래디언트가 시그모이드를 통해 흐를 수 있게 하면서도, 추론 시에는 하드 결정을 유지할 수 있습니다.

2. LoRA 어댑터: 모든 24개 블록의 Query, Key, Value 및 Output 프로젝션($W_Q, W_K, W_V, W_O$)에 로우-랭크 어댑터(Rank 8)가 적용됩니다. 백본 가중치는 동결된 상태를 유지합니다.

학습 전략 및 혁신 요소:

• 결합 최적화(Joint Optimization): 라우터 가중치와 LoRA 행렬은 결합 손실 함수 $L = L_{LM} + \lambda \cdot \frac{1}{L}\sum \sigma(s_i)$를 사용하여 공동으로 학습됩니다.
• 게이트 정규화(Gate Regularization): 정규화 항($\lambda=1.0$)은 높은 소프트 게이트 값을 억제하여, 라우터가 모든 게이트를 열어두는 단순한 해(trivial solution)로 수렴하는 것을 방지합니다.
• 편향된 초기화(Biased Initialization): 대칭성을 깨뜨리고, 게이트가 스킵하기 전까지는 차별화할 수 없는 "닭과 달걀" 문제를 피하기 위해, 저자들은 중간 레이어 라우터(레이어 8–16)를 $-1.0$의 바이어스(처음에는 닫힘/스킵 상태)로, 외곽 레이어를$+1.0$의 바이어스(처음에는 열림 상태)로 초기화했습니다. 이는 언어 모델링 손실을 바탕으로 즉각적인 차별화를 강제합니다.

주요 기여

본 논문은 다섯 가지 구체적인 기여를 명시합니다:

1. 학습/추론 간의 불일치를 제거하기 위해 STE를 사용하는 레이어별 하드 게이티드 스킵 연결(per-layer hard-gated skip connections).
2. 라우팅 정책이 적응된 가중치와 함께 진화할 수 있도록, 어텐션 프로젝션과 LoRA 행렬을 결합 LoRA 미세 조정(Joint LoRA fine-tuning).
3. 전체 계산량을 유지하는 대신 스킵 가능한 블록을 발견하도록 강제하는 게이트 정규화(Gate regularization).
4. 첫 번째 학습 단계부터 그래디언트 기반의 차별화를 가능하게 하는 중간 레이어 라우터의 편향된 초기화(Biased initialization).
5. Qwen2.5-0.5B-Instruct에 대한 경험적 검증을 통해, 7분 미만 만에 학습된 도구 호출과 계획 단계 사이의 유의미한 스킵 차이를 입증함.

실험 결과

모델은 혼합 에이전틱 데이터셋(Hermes, Glaive, GSM8K, Turing)에 대해 3,000 스텝(A100 40GB에서 약 6.4분 소요) 동안 학습되었으며, 단 1.10M 개의 학습 가능한 파라미터(백본 494M의 0.22%)만을 사용했습니다.

• 스킵 차이(Skip Differential): LayerRoute는 12.91%의 스킵 차이를 달ر성했습니다. 도구 호출은 **15.25%**의 FLOPs를 스킵한 반면, 계획 단계는 **2.34%**만을 스킵했습니다.
• 품질 개선: 일반적인 효율성-품질 간의 트레이드오프와 달리, 게이티드 모델은 전체 레이어 베이스라인보다 개선된 퍼플렉시티를 보여주었습니다.
  • 도구 호출: 퍼플렉시티 변화량 −1.29.
  • 계획: 퍼플렉시티 변화량 −1.30.
  • 이러한 개선은 LoRA 적응이 스킵으로 인한 품질 비용을 보상하기 때문입니다.
• 절제 연구(Ablation Studies):
  • LayerRoute-NoReg: 스킵이 거의 발생하지 않는(0.1%) 상태로 붕괴되었으며, 이는 정규화가 필수적임을 확인시켜 줍니다.
  • LayerRoute-UniformInit: 스킵 차이가 역전되었습니다(계획 단계가 도구 호출보다 더 많이 스킵됨). 이는 편향된 초기화의 결정적 역할을 증명합니다.
  • LayerRoute-BCE: 명시적 레이블 대신 결합 최적화로 학습했을 때, 훨씬 낮은 스킵 차이(0.3%)를 보였습니다.

게이트 구조: 학습된 게이트는 안정적인 두 클러스터 구조로 수렴했습니다. 레이어 0–7 및 17–23은 안정적으로 열려 있는 상태를 유지했고, 레이어 8–16(중간 클러스터)은 도구 호출 시에는 확실히 스킵되었으나 계획 단계에서는 유지되었습니다.

의의 및 주장

본 논문은 에이전틱 단계 유형이 별도의 태스크 유형 레이블이나 별도의 분류 헤드 없이도 언어 모델링 목적 함수로부터 암시적으로 학습될 수 있음을 입증한다고 주장합니다.

• 효율성: 도구 호출에 대해 15.2%의 FLOPs 감소를 달성하면서도 품질을 유지하거나 오히려 향상시켰습니다.
• 적응성: 시스템은 짧고 구조화된 시퀀스와 길고 추론 집약적인 시퀀스를 성공적으로 구분하여, 필요한 곳에만 계산량을 적용합니다.
• 파라미터 효율성: 전체 적응 과정은 백본 파라미터의 0.22% 미만을 필요로 하며 몇 분 내에 완료될 수 있습니다.

인정된 한계점:
저자들은 중간 레이어 클러스터(8–16)가 내부적인 차별화가 없는 단일 스킵 블록을 형성하여 잠재적인 FLOPs 절감 효과를 제한한다는 점을 언급했습니다. 또한, 계획 단계에서의 스킵 비율이 2.34%로 낮아, 계획 중심의 워크로드에 대해서는 이점이 미미합니다. 현재 평가는 0.5B 모델 크기로 제한되어 있습니다.