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

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

왜 중요한가

핵심 기여

라우팅 가중치 붕괴 현상 규명

이론적 분석과 실험을 통해 기존 Mixture-of-LoRAs 모델이 학습 중 단 하나의 LoRA에만 가중치를 집중시켜 다중 LoRA의 이점을 잃는 현상을 증명했다.

ReMix 아키텍처 제안

활성화된 k개의 LoRA에 동일한 고정 가중치를 부여하여 모든 어댑터가 학습에 균등하게 참여하도록 강제하는 구조를 설계했다.

RLOO 기반 그래디언트 추정기 도입

미분 불가능한 고정 가중치 라우터를 학습시키기 위해 강화학습의 RLOO 기법을 적용하여 편향되지 않고 분산이 낮은 학습 알고리즘을 구축했다.

파라미터 효율성 극대화

기존 SOTA 모델 대비 활성화 파라미터 수를 최대 90% 줄이면서도 GSM8K, HumanEval 등 주요 벤치마크에서 더 높은 정확도를 달성했다.

핵심 아이디어 이해하기

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를 통해 학습하는 이중 최적화 구조를 가진다.

실무 활용

적은 연산 자원으로도 고성능 LLM을 유지해야 하는 온디바이스 AI나 멀티태스크 서빙 환경에 매우 적합하다. 특히 수학이나 코딩처럼 정교한 추론이 필요한 도메인에서 다중 어댑터의 능력을 극대화할 수 있다.

수학/과학 문제 풀이 전용 특화 모델 구축
다양한 프로그래밍 언어를 지원하는 코드 생성 에이전트
제한된 GPU 메모리 환경에서의 효율적인 멀티태스크 파인튜닝
실시간 사용자 피드백 기반의 동적 어댑터 선택 시스템

코드 공개 여부: 비공개

키워드

LoRA(저순위 어댑터)PEFT(파라미터 효율적 미세 조정)Mixture-of-LoRAs(LoRA 혼합 모델)Reinforcement Learning(강화학습)Routing Weight Collapse(라우팅 가중치 붕괴)

용어 해설

저순위 어댑터(LoRA): — 사전 학습된 모델의 가중치를 고정한 채, 작은 크기의 두 행렬만을 추가하여 학습시키는 기법이다. 전체 파라미터를 업데이트하지 않아도 되므로 메모리와 연산 비용을 획기적으로 줄이면서 모델을 특정 작업에 맞게 조정할 수 있다.
전문가 혼합(Mixture-of-Experts): — 입력 데이터에 따라 여러 개의 작은 신경망(전문가) 중 일부만을 선택적으로 활성화하여 처리하는 구조이다. 전체 모델 크기는 키우면서도 실제 연산량은 낮게 유지할 수 있어 대규모 언어 모델의 효율성을 높이는 데 사용된다.
라우팅 가중치 붕괴(Routing Weight Collapse): — 여러 전문가 모델을 섞어 쓸 때, 학습 과정에서 특정 전문가 하나에만 가중치가 쏠려 나머지 전문가들이 전혀 활용되지 못하는 현상이다. 이로 인해 다중 모델을 사용하는 이점이 사라지고 단일 모델과 다를 바 없는 성능을 내게 된다.
강화학습 기반 리브-원-아웃 추정기(RLOO): — 강화학습에서 여러 번의 시도(샘플)를 한 뒤, 특정 시도의 성과를 나머지 시도들의 평균과 비교하여 보상을 계산하는 방식이다. 그래디언트의 분산을 줄여 학습을 더 안정적으로 만들어주는 효과가 있다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

NoRA로 LoRA의 입력 투영 규모를 정규화하다

LLM 라우팅 개발부터 배포까지 통합 인프라

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

NoRA로 LoRA의 입력 투영 규모를 정규화하다

LLM 라우팅 개발부터 배포까지 통합 인프라