ReMix: LLM 파인튜닝을 위한 Mixture-of-LoRAs의 강화학습 기반 라우팅

기존의 여러 LoRA를 섞어 쓰는 방식은 학습 과정에서 특정 LoRA 하나에만 의존하게 되는 라우팅 가중치 붕괴 현상 때문에 효율이 떨어졌다. ReMix는 강화학습을 도입해 모든 활성화된 LoRA가 골고루 기여하게 만듦으로써, 더 적은 파라미터로도 수학 연산과 코드 생성 등 복잡한 작업에서 뛰어난 성능을 보여준다.

LoRA는 모델 전체를 학습시키는 대신 작은 행렬만 추가해 효율을 높이지만, 복잡한 작업에서는 용량이 부족할 수 있다. 이를 해결하기 위해 여러 개의 LoRA를 두고 입력에 따라 선택하는 Mixture-of-LoRAs가 등장했으나, Softmax 기반의 학습 방식은 승자 독식 현상을 일으켜 결국 하나의 LoRA만 쓰게 되는 한계가 있었다.

ReMix는 이 문제를 가중치를 학습하지 않는 것으로 해결한다. 선택된 k개의 LoRA에 무조건 똑같은 가중치를 부여함으로써, 특정 LoRA가 지배하는 것을 원천 차단한다. 하지만 가중치가 고정되면 역전파를 통해 라우터를 학습시킬 수 없다는 새로운 문제가 발생한다.

연구진은 이를 위해 라우팅을 행동으로, 손실 함수를 보상으로 간주하는 강화학습 프레임워크를 도입했다. 특히 RLOO라는 기법을 통해 여러 샘플 간의 차이를 비교하며 라우터를 업데이트함으로써, 미분 없이도 어떤 입력을 어떤 LoRA로 보내야 할지 정확하게 학습할 수 있게 되었다.

ReMix는 입력 x에 대해 라우터 P를 거쳐 확률 분포 q를 생성한다. 학습 시에는 이 분포에서 k개의 LoRA를 중복 없이 샘플링하여 활성화한다. 활성화된 모든 LoRA에는 동일한 상수 가중치 ω를 부여하여 출력 y = Wx + ωΣBiAix를 계산한다.

라우터 학습을 위해 RLOO(Reinforce Leave-One-Out) 추정기를 사용한다. M개의 독립적인 선택 세트를 샘플링하고, 각 세트의 손실값에서 나머지 세트들의 평균 손실을 뺀 값을 기반으로 그래디언트를 계산한다. [입력 x → M번의 라우팅 샘플링 → 각 샘플의 SFT Loss 계산 → 평균 대비 성능 차이로 라우터 갱신 → 최적의 라우팅 경로 학습] 과정을 거친다.

추론 시에는 샘플링 대신 확률이 가장 높은 상위 k개를 선택하는 Top-k Selection을 수행한다. 이는 라우터가 충분히 학습되었을 때 이론적으로 최적의 선택임을 증명했다.

Llama 3 8B 모델을 기반으로 한 실험에서 ReMix는 GSM8K(수학) 65.66%, HumanEval(코드) 32.93%, ARC-c(상식) 83.73%를 기록했다. 이는 기존 SOTA인 MixLoRA나 HydraLoRA보다 평균 3.34% 높은 수치이며, 특히 활성화 파라미터가 훨씬 적은 상태에서 달성한 결과이다.

파라미터 효율성 측면에서 ReMix는 0.070B개의 활성화 파라미터만 사용하여, 0.675B개를 사용하는 VB-LoRA 대비 90%의 파라미터 절감을 실현하면서도 성능은 더 우수했다.

Ablation Study를 통해 RLOO 기법과 Top-k 추론 방식이 각각 성능 향상에 필수적임을 확인했다. RLOO를 제거할 경우 성능이 급격히 하락하여 강화학습 기반 라우팅의 중요성을 입증했다.

ReMix 아키텍처는 기존 MoE(Mixture of Experts)의 라우팅 가중치 붕괴 문제를 해결하기 위해 Non-learnable Constant Weights를 채택했다. 이는 Effective Support Size(ESS)를 항상 k로 유지하여 모델의 표현력을 강제로 보존한다.

수학적으로 라우터 파라미터 P에 대한 그래디언트가 정의되지 않으므로, 정책 그래디언트의 일종인 RLOO를 사용한다. RLOO는 M개의 샘플을 사용하여 베이스라인을 구축함으로써 그래디언트 추정의 분산을 획기적으로 줄인다.

이론적 분석(Theorem 1)을 통해 가우시안 초기화된 Softmax 라우터가 학습 초기부터 소수의 LoRA에 집중될 확률이 매우 높음을 증명했다. ReMix는 이러한 구조적 취약성을 강화학습 프레임워크로 우회한다.

구현 시 LoRA 파라미터(A, B)는 일반적인 역전파로 학습하고, 라우터(P)만 RLOO를 통해 학습하는 이중 최적화 구조를 가진다.