기술 요약: 재학습이 아닌 선택을 배우기: 하드 라우팅 기반의 추론 LoRA 혼합 (MoR-LoRA)

문제 정의

본 논문은 개인정보 보호 또는 규제 문제로 인해 원본 학습 데이터를 공유할 수 없는 상황에서, 다중 도메인 적응을 위해 독립적으로 학습된 저차원 적응(LoRA) 어댑터를 하나의 거대 언어 모델(LLM)로 구성하는 과정에서 발생하는 과제를 다룹니다.

기존의 표준 Mixture-of-Experts (MoE) 소프트 라우팅을 동결된(frozen) 독립 학습 LoRA 어댑터에 적용할 때 결정적인 문제가 발생합니다. LoRA 모듈은 단위 스케일 가산 업데이트($W_{updated} = W_{pretrained} + AB$)를 가정하며 학습됩니다. 그러나 표준 소프트 MoE 라우팅은 가중 합($W_{updated} = W_{pretrained} + \sum \omega_i A_i B_i$)을 사용하여 전문가를 결합하며, 이때 라우팅 가중치는 $\omega_i < 1$입니다. 이는 **스케일 불일치(scale mismatch)**를 유발합니다. 즉, 통합 과정에서 적용되는 유효 업데이트가 해당 전문가가 원래 학습되었던 단위 스적 업데이트와 달라지게 됩니다.

LoRA-Mixer와 같은 기존 솔루션은 보조 보존 손실(auxiliary preservation losses)을 사용하여 어댑터를 재학습함으로써 이를 해결하려고 시도합니다. 하지만 이 방식은 학습 가능한 파라미터 수를 증가시키고, 더 많은 데이터를 요구하며, 기존 전문가에 인코딩된 지식을 변경할 위험이 있습니다. 이러한 문제는 특히 Chain-of-Thought (CoT) 추론을 수행하는 인스트럭션 튜닝 모델에서 더욱 두드러지는데, 기존의 일반적인 미세 조정 방식은 내부 추론 행동을 저하시킬 수 있으며, 기존의 구성 방법들은 이러한 특화된 추론 흔적(reasoning traces)을 보존하는 데 실패하는 경우가 많습니다.

방법론

저자들은 동결된 추론 LoRA 전문가들을 **단위 스케일 하드 선택(unit-scale hard selection)**을 통해 구성하도록 설계된 2단계 프레임워크인 Hard-Routed MoR-LoRA를 제안합니다.

1단계: 추론 전문가 학습

도메인 특화 LoRA 어댑터는 **검증 가능한 피드백으로부터의 강화 학습(RLVF)**을 사용하여 독립적으로 학습됩니다. 참조 출력(reference outputs)을 모방하도록 유도하는 지도 미세 조정(SFT)과 달리, RLVF는 자동 검증 가능한 보상(예: 최종 정답 확인)을 사용하여 정확성을 최적화합니다. 이 패러다임은 베이스 모델의 정렬된 추론 패턴을 덮어쓰지 않으면서 구조화된 CoT 추론의 출현을 장려합니다. 이렇게 생성된 전문가들은 이후 동결됩니다.

2단계: 하드 라우팅 통합

동결된 전문가들은 **하드 라우티드 혼합(Hard-Routed Mixture)**을 통해 공유 모델로 통합됩니다.

• 하드 Top-1 라우팅: 각 토큰에 대해 라우터는 정확히 하나의 전문가를 선택합니다. 이는 선택된 전문가가 단위 스케일로 적용되도록 보장하여, 원래의 LoRA 가정($W_{updated} = W_{pretrained} + 1 \cdot A_j B_j$)을 유지합니다.
• **Straight-Through Estimator (STE): 이산적 선택(argmax)은 미분 불가능하므로, 저자들은 STE를 채택했습니다. 순전파(forward pass)에서는 하드 결정(hard decision)을 사용하고, 역전파(backward pass)에서는 소프트 확률을 통해 그래디언트를 전파합니다. 이를 통해 추론 시에는 이산적인 단위 스케일 동작을 유지하면서도 엔드투엔드(end-to-end) 라우터 학습을 가능하게 합니다.
• 파라미터 효율성: 경량화된 공유 선형 라우터와 작은 어텐션 LoRA 모듈만 학습됩니다. 사전 학습된 전문가들은 동결된 상태로 유지됩니다.
• 지식 증류(Distillation): 라우터는 동결된 전문가로부터 증류된 작은 규모의 추론 흔리 데이터셋을 통해 SFT 방식으로 학습됩니다. 목적은 추론 능력 자체를 재학습하는 것이 아니라, 어떤 전문가를 선택할지를 배우는 것입니다.

주요 기여

1. 스케일 불일치 식별: 본 논문은 동결된 독립 학습 LoRA 어댑터에 소프트 MoE 라우팅을 적용할 경우, 어댑터의 유효 함수를 변화시키는 스케일 불일치가 발생함을 식별했습니다.
2. 하드 단위 스케일 라우팅: 전문가를 재학습 없이 모듈식으로 구성할 수 있도록, Straight-Through Estimator를 사용한 하드 Top-1 라우팅을 제안하여 단위 스케일 적용을 보존했습니다.
3. 2단계 프레임워크: 추론 습득(RLVF를 통한)과 전문가 선택(하드 라우팅을 통한)을 분리하여, 통합 과정 중 전문가를 동결 상태로 유지하는 새로운 파이프라인을 제시했습니다.
4. 실증적 검증: 이 접근 방식이 전문가의 행동을 보존하면서도 소프트 라우팅 베이스라인보다 훨씬 적은 수의 학습 가능한 파라미터를 요구함을 입증했습니다.

실험 결과

실험은 LLaMA-3, Gemma-3, Qwen-2.5 제품군을 사용하여 5개의 벤치마크(GSM8K, ARC-Challenge, MedQA, BoolQ, CoLA)와 다양한 모델 규모(1B, 3B, 8B)에 대해 수행되었습니다.

• RLVF의 우수성: RLVF로 학습된 전문가는 유능한 모델(3B 및 8B)의 추론 중심 작업에서 SFT 학습 전문가보다 성능이 크게 뛰어났던 반면, SFT는 대규모 인스트럭션 튜닝 모델의 추론 성능을 저하시키기도 했습니다.
• 성능 대비 파라미터 수: Hard-Routed MoR-LoRA는 소프트 라우팅 베이스라인(LoRAMixer, MoLE)과 대등하거나 더 나은 성능을 달면서도 훨씬 적은 파라미터를 학습했습니다. 예를 들어, RLVF 전문가를 사용한 3B 모델의 경우, 제안된 방법은 평균 정확도 72.07에 도달하기 위해 약 73M 개의 파라미터를 학습한 반면, LoRAMixer는 65.67에 도달하기 위해 약 606M 개의 파라미터가 필요했습니다.
• 라우팅 행동 분석: 소프트 Top-2 라우팅 분석 결과, 종종 Near Top-1 선택으로 붕괴되는 현상(평균 지배적 가중치 ~0.71)이 나타났으며, 이는 하드 라우팅이 의도된 동작을 더 직접적이고 효율적으로 포착함을 시사합니다.
• 데이터 효율성: 통합 단계(2단계)는 전체 데이터셋으로 학습하는 것과 유사한 성능을 내기 위해 데이터셋당 약 1,000개의 샘플만을 필요로 했으며, 이는 라우터가 도메인 지식이 아닌 선택 패턴을 학습함을 나타냅니다.
• 혼합 도메인 능력: 혼합 도메인 평가(예: 수학과 독해 이해가 모두 포함된 프롬프트)에서 토큰 레벨 하드 라우팅이 프롬프트 레벨 라우팅보다 우수한 성능을 보였으며, 이는 단일 입력 내에서 서로 다른 전문가를 선택하는 유연성을 입증합니다.

의의 및 주장

본 논문은 모듈형 적응이 각 전문가가 무엇을 아는지 수정하는 것이 아니라, 각 전문가를 언제 사용할지를 학습함으로써 달성될 수 있다고 주장합니다.

저자들은 동결된 사전 학습 LoRA 전문가의 경우, 하드 단위 스케일 라우팅이 단순하고 효율적인 추상화를 제공한다고 논합니다. 통합을 주로 라우팅 문제로 취급함으로써, 전문가를 재학습하는 데 따르는 계산 비용과 잠재적인 지식 저하 문제를 피할 수 있습니다. 결과는 구성의 주요 이점이 전문가의 파라미터를 광범위하게 수정하는 것이 아니라, 컨텍스트에 적합한 동결된 전문가를 선택하는 데 있음을 시사합니다.

본 연구는 프롬프트 레벨 라우팅이 깨끗한 단일 도메인 입력에는 충분하지만, 복잡한 혼합 도메인 입력을 위해서는 토큰 레벨 하드 라우팅이 필수적인 유연성을 제공한다는 점을 강조합니다. 저자들은 2단계 과정에서 라우팅 정책을 배우기 위한 추가적인 레이블링된 데이터가 필요하며, 현재로서는 추가적인 라우팅 데이터 없이 새로운 도메인을 진정한 제로샷(zero-shot)으로 통합하는 것을 지원하지 않는다는 점에서 이 방법이 완전히 학습이 필요 없는(training-free) 방식은 아니라고 겸허히 밝히고 있습니다.