Dynamic Latent Routing(DLR): 동적 잠재 코드 기반 라우팅으로 LLM 파인튜닝을 단일 단계에서 수행

본 논문은 time-varying rewards를 가지는 dynamic MDP에서 서브-정책의 시간적 연결을 통해 최적 goal-reaching 정책을 얻는 General Dijkstra Search(GDS)를 제시한다. 이를 LLM 파인튜딩에 적용한 Dynamic Latent Routing(DLR)은 discrete latent codes와 routing policy를 residual stream에 주입하는 단일-stage 학습으로, low-data 상황에서도 SFT를 상회하는 성능을 보이고 코드가 실제로 서브태스크에 대응하는 구조적 라우팅을 형성함을 보여준다.

단계별 구성 요소와 동작 원리: (1) Dynamic MDP에서 reward가 시간에 따라 변한다는 설정을 도입하고, 정책 합성(Concatenation)으로 최적 정책을 구성한다. (2) General Dijkstra Search(GDS)는 중간 목표에 대해 각각 최적의 서브-정책을 구성하고 이를 시간적으로 연결해 최적 목표 도달 정책을 발견한다는 이론을 제공한다. (3) DL R은 이 아이디어를 LLM에 적용해 discrete latent codes를 생성하고, routing head를 통해 chunk 단위로 코드를 선택하며 residual stream에 steering vector를 더한다. (4) Loss 구성은 Generalist(−log pθ(x)) + Information Gain(−log pθ(x|a) sg pθ(x)) + Policy Optimization(αpolicy log pθ(a|x)) + Marginal Entropy Regularization(αreg KL(pθ(a)||pbi-zipf(a)))의 합으로 정의된다. (5) 학습 루프은 N개의 후보 코드 시퀀스 중 가장 조건부 likelihood가 큰 것을 선택하고 LDLR, 코드북, 라우팅 헤드를 함께 업데이트하는 single-stage 루프이다. 이로써 입력 의존적 코드 라우팅이 학습과 탐색을 하나의 패스에서 조정하고, 코드를 통해 서브태스크를 해석 가능하게 만든다.

전체 접근: (1) 입력 시퀀스 x를 chunk 단위로 분해하고 m(t) = floor((t−1)/K)로 코드 a_m을 매핑한다. (2) 코드북에서 a_m에 대응하는 라우팅 벡터 α_eam를 얻고, h(l∗)에 α_eam을 더해 각 토큰의 hidden state에 주입한다. (3) 라우팅 로짓은 Wrt h(l∗)_mK+1를 통해 계산하고 sg(·)로 경로를 고정하지 않는 형태로 업데이트한다. (4) LDLR의 네 가지 항목으로 loss를 구성하고 stop-gradient를 통해 코드-레벨과 LM-레벨의 그래디언트를 분리한다. (5) 학습 중에는 pθ(a|x)를 이용한 policy-head를 통해 코드를 샘플링하고 가장 높은 pθ(x|a)를 가지는 a를 선택한다. (6) 코드북의 다양성은 bigram-Zipfian prior(pbi-zipf)로 유지되며, pθ(a)와 pbi-zipf(a)의 KL 발산으로 정규화를 수행한다. (7) 실험 설정은 24개 모델-데이터 구성, 2×H100, codebook 크기 C=32, abstraction ratio K=4, rollouts N=8, steered layer l∗ 등이다.

주요 결과 요약: (i) DLR은 24개model-dataset 설정에서 SFT 대비 평균 +6.6pp 개선을 보였고, GSM8K에서 +10.2pp, ScienceQA에서 +18.8pp의 최대 개선을 기록했다. (ii) 코드북 다양성은 cos 유사도 0.010.28 수준이며 41100%의 코드가 최소 한 번 이상 사용된다. (iii) 6-deep carry-cascade에서 DLR의 이점이 가장 두드러지며, 전반적으로 코드의 토픽-특화성을 확인할 수 있다. (iv) ablation에서 αpolicy를 0으로 두면 평균 약 -9.8pp 감소 등 정책-최적화 항이 중요함을 확인했다. (v) six-digit arithmetic 사례에서 abstraction codes는 carry/subtask를 외부로 읽을 수 있게 해주며, 코드 교환으로 일부 오답을 수정하는 효과를 보여준다.

A. Dynamic MDP 및 정책 구성: DMDP 정의와 시간-인덱스 값 함수 Vπ_t(s), Qπ_t(s,a) 정의. B. GDS 이론: 정책 연결을 통한 가치 분해 및 최적성 보장(Thm. 12, Thm. 13). C. DLR 구조: chunk 수준 인젝션(l∗)에서의 코드-기저와 routing/logits 계산, stop-gradient의 역할. D. Loss 구성: LDLR = −log pθ(x) + −log pθ(x|a) sg pθ(x) + αpolicy log pθ(a|x) + αreg KL(pθ(a) || pbi-zipf(a)). E. 학습: 샘플링된 N개의 코드 시퀀스 중 하나를 선택하고, LDLR로 정책-코드-모델을 함께 학습. F. 실험: 4개 벤치마크, 6개 모델, 24개 설정; 2×H100, codebook 크기 32, abstraction ratio 4, rollouts 8, steered layer 다양성 탐색. G. 해석: SciQA에서 코드-토픽 purity, 코드-subtask heatmap, carry-state 분류와의 연계, 자동 해석 인터뷰. H. Arithmetic Case Study: 2L/1H/128d 모델에서 서브태스크-전용 코드 t21, t23 등과 carry-상태 분류 STn=0/U/1를 코드로 재현.

한계로 GDS의 이론은 시간-가변 보상을 외부 환경으로 가정하는 유한한 Dynamic MDP에 적용되며, DLR의 보상은 내부 모델 파라미터 θ에 의존하는 확률 로그우 존재로 인해 이론적 보장을 넘는다. 실험은 저데이터(post-training) 시나리오에 한정되며, 6×4 구성이 두 개의 오픈-웨이트 계열(Qwen3, Llama-3.2) 및 0.6B~8B 사이에 한정된다. 계산 자원은 2×H100에서 수행되며, 더 큰 데이터-환경에서의 확장성은 추후 연구로 남겨 둔다.