ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning

이 논문은 기존 Mixture-of-LoRAs 모델에서 발생하는 라우팅 가중치의 불균형 문제를 해결하기 위해, 강화 학습 기반의 RLOO 기법을 활용한 비학습형 라우팅 가중치를 도입하여 모든 LoRA 가 균등하게 활성화되도록 하는 'ReMix'를 제안하고, 이를 통해 제한된 활성화 파라미터 수로 최첨단 성능을 달성함을 보여줍니다.

핵심

🍽️ 비유: "요리사 팀"과 "지배인"의 이야기

거대한 AI 모델을 새로운 요리 (작업) 를 하도록 가르친다고 상상해 보세요. 이때 LoRA라는 것은 특화된 요리사들입니다.

• 예: A 요리사는 파스타를, B 요리사는 스테이크를, C 요리사는 디저트를 아주 잘합니다.
• 우리는 이 요리사들 중 몇 명을 뽑아서 (예: 4 명) 함께 일하게 하면 더 좋은 요리를 만들 수 있다고 생각합니다. 이것이 MixLoRA라는 기존 기술입니다.

하지만 여기서 치명적인 문제가 발생했습니다.

📉 문제: "지배인"이 한 명만 고집하는 현상 (라우팅 붕괴)

기존 기술에서는 **지배인 (라우터)**이 각 요리사에게 "이 손님을 당신에게 맡길게"라고 말해주는 **가중치 (점수)**를 매겼습니다.

• 문제는 지배인이 학습을 하다 보면, 한 명의 요리사 (예: 파스타 요리사) 에게만 99% 의 점수를 주고, 나머지 3 명은 0.3% 정도의 점수만 준다는 것입니다.
• 결과? 나머지 3 명의 요리사는 아예 일도 안 하고 구경만 합니다.
• 결론: 4 명을 뽑아서 일하게 했다고 해도, 실제로는 1 명만 일하는 것과 똑같은 결과가 나옵니다. 돈과 시간을 들여 4 명을 고용했는데, 1 명만 쓰는 꼴이 된 것입니다.

저자들은 이 현상을 **"라우팅 가중치 붕괴 (Routing Weight Collapse)"**라고 불렀습니다.

---

💡 해결책: ReMix (재믹스) - "공정한 배분"과 "강화 학습"

이 문제를 해결하기 위해 제안한 ReMix는 두 가지 핵심 아이디어를 사용합니다.

1. 고정된 공정한 배분 (Non-learnable Weights)

기존의 "지배인"이 점수를 매기는 방식 (학습 가능한 가중치) 을 버렸습니다. 대신, 활성화된 요리사들에게는 무조건 똑같은 점수 (예: 25%씩) 를 주는 고정된 규칙을 적용했습니다.

• 비유: "오늘은 파스타, 스테이크, 디저트, 수프 요리사 모두에게 똑같은 일을 분배해. 누구도 우대하지도, 무시하지도 않아."
• 이렇게 하면 4 명을 뽑았을 때 4 명 모두 100% 능력을 발휘하게 됩니다.

2. 강화 학습 (RL) 과 RLOO 기술

그런데 문제가 생겼습니다. "점수를 고정해버리면, AI 가 어떻게 배우지? 점수를 바꿀 수 없는데?"

• 해결책: 점수를 바꾸는 대신, **"어떤 요리사 조합을 뽑을지"**를 학습하게 했습니다.
• 비유: 지배인은 "누구에게 일을 줄지"만 결정하고, 일은 모두 똑같이 시킵니다. 그리고 "어떤 조합을 뽑았을 때 요리가 가장 맛있었는지 (손실 함수)"를 기준으로 다음에 누구를 뽑을지 학습합니다.
• RLOO 기술: 이 학습 과정에서 '실패한 경우'를 제외하고 '성공한 경우'만 집중적으로 분석하는 고급 통계 기법을 써서, 학습 속도를 높이고 정확도를 올렸습니다.

---

🚀 왜 ReMix 가 더 좋은가요?

1. 진짜로 여러 명이 일합니다: 4 명을 뽑으면 4 명이 모두 기여합니다. (기존 방법은 1 명만 기여)
2. 더 똑똑해집니다: 다양한 요리사들의 능력을 모두 활용하므로, 복잡한 문제 (수학 문제, 코드 작성 등) 를 더 잘 풉니다.
3. 효율적입니다: 더 적은 파라미터 (인력) 로 더 좋은 결과를 냅니다.

📊 실험 결과

논문에서는 수학 문제 (GSM8K), 코드 작성 (HumanEval), 일반 상식 (ARC-c) 테스트를 진행했습니다.

• 결과: ReMix 는 기존 최고의 방법들보다 평균적으로 더 높은 점수를 받았습니다.
• 특히, 학습 시간을 조금 더 들였을 때 성능이 훨씬 더 좋아지는 것을 확인했습니다. (기존 방법들은 학습 시간을 늘려도 성능이 안 올라갔지만, ReMix 는 더 많은 계산을 할수록 더 똑똑해졌습니다.)

🎯 한 줄 요약

"기존 AI 학습법은 한 명만 일하게 해서 낭비가 심했는데, ReMix 는 모든 전문가를 공평하게 일하게 하고, '누구를 뽑을지'만 지능적으로 학습시켜서 훨씬 더 똑똑하고 효율적인 AI 를 만들었습니다."

기술 요약

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 파라미터 효율적 파인튜닝 (PEFT) 기법 중 하나인 LoRA (Low-Rank Adapter) 를 기반으로 한 Mixture-of-LoRAs 모델이 존재합니다. 이는 각 레이어 입력을 여러 개의 전문화된 LoRA 서브셋으로 라우팅하여 표현력을 높이는 것을 목표로 합니다.

그러나 기존 Mixture-of-LoRAs 모델은 다음과 같은 치명적인 한계를 가지고 있습니다.

• 라우팅 가중치 붕괴 (Routing Weight Collapse): 학습 가능한 라우터 (Router) 가 각 LoRA 에 할당하는 가중치가 극도로 불균형하게 됩니다. 이론적 및 실험적 분석 결과, 파인튜닝이 진행됨에 따라 특정 입력에 대해 단 하나의 LoRA 가 거의 1 에 가까운 가중치를 차지하고, 나머지 $k-1$개의 LoRA 는 가중치가 0 에 수렴하여 사실상 비활성화되는 현상이 발생합니다.
• 표현력 저하: $k$개의 LoRA 를 활성화한다고 설정했더라도, 실제로는 $k=1$과 유사한 성능만 발휘하게 되어 Mixure-of-LoRAs 모델이 가진 잠재적인 표현력이 낭비됩니다.
• 기존 방법의 한계: 기존에는 라우팅 가중치를 학습 가능한 연속 값으로 두어 엔드 - 투 - 엔드 학습을 시도했으나, 이로 인해 위와 같은 붕괴 현상이 필연적으로 발생했습니다.

2. 제안 방법: ReMix (Methodology)

저자들은 라우팅 가중치 붕괴 문제를 해결하기 위해 ReMix (Reinforcement Routing for Mixture-of-LoRAs) 를 제안합니다. 핵심 아이디어는 학습 가능한 가중치 대신 고정된 (비학습) 가중치를 사용하여 모든 활성화된 LoRA 가 균등하게 기여하도록 하는 것입니다.

2.1 아키텍처: 비학습 라우팅 가중치 (Non-Learnable Weights)

• 균등 가중치 할당: 라우터가 $k$개의 LoRA 를 선택하면, 선택된 모든 LoRA 에 대해 상수 (Constant) 인 가중치 $\omega$를 부여합니다.
  • 수식: $\pi_i^{(l)} = \omega$ (선택된 경우), $0$ (선택되지 않은 경우).
• 효과: 이를 통해 활성화된 LoRA 의 유효 지지 크기 (Effective Support Size, ESS) 를 항상 $k$로 유지시켜, 어떤 하나의 LoRA 가 지배하는 것을 방지합니다.
• 추론 비용: 기존 Mixture-of-LoRAs 방법과 동일한 추론 오버헤드를 가지며, 추가적인 비용이 발생하지 않습니다.

2.2 파인튜닝: 강화학습 기반 RLOO 그라디언트 추정 (RL-based RLOO Estimator)

비학습 가중치를 사용하면 역전파 (Backpropagation) 를 통해 라우터 파라미터를 직접 학습할 수 없습니다. 이를 해결하기 위해 라우터 학습 문제를 강화학습 (RL) 문제로 재정의합니다.

• 문제 설정:
  • 정책 (Policy): 라우터가 LoRA 서브셋을 선택하는 확률 분포.
  • 보상 (Reward): 파인튜닝 손실 (SFT Loss) 의 음수 ($-L$).
• 그라디언트 추정:
  • 라우터 파라미터에 대한 그라디언트를 추정하기 위해 RLOO (Reinforce Leave-One-Out) 기법을 적용합니다.
  • $M$개의 선택 (Selection) 을 샘플링하고, 각 선택의 손실과 평균 손실의 차이를 이용하여 편향되지 않은 (Unbiased) 그라디언트 추정치를 계산합니다.
  • 이는 분산을 줄이고 안정적인 학습을 가능하게 하며, 계산 자원을 늘려 성능을 향상시킬 수 있는 확장성 (Scalability) 을 제공합니다.

2.3 추론: Top-k 선택 (Top-k Selection)

• 학습이 충분히 잘 이루어진 상태에서는, 라우터가 출력하는 확률 분포 $q^{(l)}$에서 확률이 가장 높은 Top-k 개의 LoRA를 선택하는 것이 최적의 전략임을 이론적으로 증명 (Theorem 2) 합니다.
• 따라서 추론 단계에서는 무작위 샘플링 대신 Top-k 선택을 사용하여 정확도를 극대화합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 통찰 (Routing Weight Collapse): 기존 Mixture-of-LoRAs 라우터가 학습 과정에서 하나의 LoRA 로 수렴하여 다른 LoRA 들을 무용지물로 만든다는 것을 이론적으로 증명하고 실험적으로 확인했습니다.
2. 간단하고 효과적인 라우터 설계: 학습 가능한 가중치 대신 고정된 균등 가중치를 사용하여 라우팅 불균형을 해결하고, 추가 추론 비용 없이 모든 활성화된 LoRA 를 효과적으로 활용합니다.
3. 강화학습 기반 학습 프레임워크: 비미분 가능한 라우터를 학습하기 위해 RLOO 기반의 편향되지 않은 그라디언트 추정기를 제안했습니다. 이는 대규모 계산 자원을 활용하여 성능을 극대화할 수 있게 합니다.
4. 압도적인 실험 결과: 다양한 벤치마크에서 기존 최첨단 PEFT 방법들을 압도하는 성능을 보여주었습니다.

4. 실험 결과 (Results)

저자들은 Llama 3 8B 모델을 기반으로 GSM8K (수학 추론), HumanEval (코드 생성), ARC-c (지식 회상) 등의 벤치마크에서 실험을 수행했습니다.

• 성능: ReMix 는 모든 베이스라인 (LoRA, DoRA, MixLoRA, HydraLoRA 등) 을 능가했습니다.
  • 평균 정확도 향상: 기존 최강 기법 대비 2.82%p 향상.
  • HumanEval: Pass@1 32.93 (기존 최강 대비 1.83%p 향상).
  • GSM8K: 정확도 65.66 (기존 최강 대비 3.19%p 향상).
• 파라미터 효율성: 약 0.070B의 학습 가능한 파라미터로 최상위 성능을 달성했습니다.
  • VB-LoRA 대비 파라미터 90% 감소, MixLoRA 대비 31% 감소하면서도 더 높은 성능을 기록했습니다.
• 활성화된 LoRA 의 다양성: $k$개의 LoRA 를 활성화했을 때, ReMix 는 $k$가 증가함에 따라 정확도가 지속적으로 향상되었습니다. 이는 ReMix 가 매번 동일한 LoRA 서브셋을 사용하는 것이 아니라, 입력에 따라 다양하고 적절한 LoRA 서브셋을 선택하고 있음을 의미합니다.
• 학습 확장성 (Scaling): 학습 시 샘플링 수 ($M$) 를 늘리면 성능이 지속적으로 향상되었습니다. 이는 기존 결정론적 학습 방식 (MixLoRA 등) 과 달리 ReMix 가 추가적인 계산 자원을 효과적으로 활용함을 보여줍니다.

5. 의의 및 결론 (Significance)

이 논문은 Mixture-of-LoRAs 모델의 핵심 병목 현상이었던 '라우팅 가중치 붕괴' 문제를 근본적으로 해결했습니다. 학습 가능한 가중치에 의존하던 기존 패러다임을 탈피하여, 균등한 가중치 할당과 강화학습 기반의 라우팅을 결합함으로써 다음과 같은 의의를 가집니다.

1. 효율성과 표현력의 균형: 적은 파라미터 수로 모델의 표현력을 극대화할 수 있는 새로운 PEFT 패러다임을 제시했습니다.
2. 확장성: 강화학습 기반의 그라디언트 추정을 통해 계산 자원을 늘릴수록 성능이 향상되는 스케일링 법칙을 따르도록 하여, 향후 더 큰 모델이나 복잡한 작업에 적용할 수 있는 가능성을 열었습니다.
3. 실용성: 추론 단계에서의 추가 비용 없이 구현 가능하여 실제 산업 환경에서의 적용 가능성이 높습니다.

결론적으로 ReMix 는 파라미터 효율적 파인튜닝 분야에서 새로운 State-of-the-Art (SOTA) 를 달성하며, LoRA 기반의 혼합 모델이 가진 잠재력을 온전히 발휘할 수 있게 하는 중요한 기여를 했습니다.