기술 요약: ASTOR – 코드 LLM 을 위한 유틸리티 기반 다중 작업 강화 학습

문제 정의

검증 가능한 보상을 활용한 강화 학습 (RL) 은 코딩 작업에 대한 대형 언어 모델 (LLM) 의 사후 학습에서 효과성이 입증되었습니다. 그러나 현재의 패러다임은 코드 생성용 모델, 단위 테스트 생성용 모델 등 작업별 RL 전문 모델을 별도로 훈련하는 방식에 의존합니다. 이 접근법은 작업 수에 비례하여 선형적으로 증가하는 메모리 및 연산 비용을 초래합니다. 반면, 단일 전문 모델에 의존할 경우 해당 모델이 특정 학습 도메인 밖에서 적용될 때 성능이 급격히 저하됩니다.

다중 작업 강화 학습 (MTRL) 은 모든 작업에 걸쳐 단일 모델을 훈련함으로써 통합된 해결책을 제시하지만, 기존 방법들은 이질적인 코딩 작업을 균일하게 처리함으로써 두 가지 치명적인 한계를 겪고 있습니다:

1. 데이터 스케줄링 ("무엇을 학습할 것인가"): 현재 방법들은 혼합된 작업 풀에서 샘플링하거나 고정된 커리큘럼을 따르며, 작업 간 상호작용을 무시합니다. 한 작업 (예: 코드 생성) 에 대한 훈련이 다른 작업 (예: 단위 테스트 생성) 에 긍정적 또는 부정적 영향을 미칠 수 있는 점을 고려하지 못합니다.
2. 정책 최적화 ("어떻게 학습할 것인가"): 기존 접근법들은 모든 작업에 대해 균일한 업데이트 전략과 단일 KL 정규화 계수를 적용합니다. 이는 코딩 작업마다 최적화 요구 사항이 상이하기 때문에 비최적입니다. 문법적으로 엄격한 작업 (예: 단위 테스트 생성) 은 실행 가능성을 유지하기 위해 보수적인 업데이트가 필요한 반면, 의미적으로 유연한 작업 (예: 커밋 메시지 생성) 은 더 넓은 탐색을 통해 혜택을 받기 때문입니다.

방법론: ASTOR

저자들은 ASTOR(uTility-driven coORdination 을 통한 multi-tASk 코드 강화 학습 프레임워크) 를 제안합니다. 이는 통합된 신호인 작업 유틸리티 (Task Utility) 를 중심으로 두 개의 결합된 모듈을 통해 이러한 과제를 해결하는 프레임워크입니다. 작업 유틸리티는 작업의 학습 잠재력과 작업 간 시너지를 모두 포착합니다.

1. 계층적 유틸리티 라우팅 데이터 스케줄링

이 모듈은 추정된 작업 유틸리티에 기반하여 훈련 예산을 동적으로 할당함으로써 "무엇을 학습할 것인가"를 다룹니다.

• 이중 수준 유틸리티 추정:
  • 작업 수준: 유틸리티는 두 가지 구성 요소로 정의됩니다:
    • 학습 잠재력: 롤아웃 보상의 분산으로 측정됩니다. 높은 분산은 성공과 실패가 공존하는 이질적인 결과 집합을 나타내며, 이는 실행 가능한 최적화 신호를 시사합니다.
    • 작업 간 시너지: 현재 작업과 다른 작업 간의 기울기 코사인 유사도로 측정됩니다. 양의 유사도는 건설적인 정합을 나타내고, 음의 유사도는 충돌을 시사합니다.
  • 프롬프트 수준: 작업 내에서 유틸리티는 현재 학습 잠재력 (보상 분산) 과 역사적 진행 상황 (평균 보상 개선) 에 의해 결정됩니다.
• 계층적 할당:
  • 작업 할당량 할당: 글로벌 훈련 배치 크기는 집계된 작업 유틸리티에 대한 온도 스케일 소프트맥스를 사용하여 작업 간에 분할됩니다. 이를 통해 잠재력과 시너지가 높은 작업은 더 많은 훈련 데이터를 받습니다.
  • 프롬프트 우선순위 지정: 각 작업의 할당된 할당량 내에서 프롬프트는 프롬프트 수준 유틸리티에 기반한 가중치 샘플링을 통해 선택되며, 모델이 적극적으로 개선 중인 샘플이 우선시됩니다.

2. 적응형 유틸리티 보정 정책 최적화

이 모듈은 각 작업의 현재 유틸리티에 기반하여 KL 정규화 계수 ($\beta_k$) 를 동적으로 보정함으로써 "어떻게 학습할 것인가"를 다룹니다.

• 메커니즘: 작업 $k$에 대한 KL 계수는 작업별 기본값 ($\beta_{base,k}$) 과 작업 유틸리티에 조건화된 동적 승수로 분해됩니다:
  $$\beta_k(t) = \beta_{base,k} \cdot (1 + \lambda \cdot U(t)(T_k))$$
• 적응 논리:
  • 높은 유틸리티를 가진 작업 (강한 학습 신호) 은 기울기 신호가 집중되었을 때 과도한 업데이트를 방지하기 위해 KL 제약을 강화하는 더 큰 $\beta_k$를 받습니다.
  • 낮은 유틸리티를 가진 작업은 정책이 개선할 여지가 있는 곳에서 더 공격적인 업데이트를 허용하는 더 작은 계수를 받습니다.
• 보상 시스템: 이 프레임워크는 소프트웨어 개발 워크플로우를 아우르는 네 가지 이질적인 작업을 지원합니다: 코드 I/O 예측 (실행 정확도), 코드 생성 (통과율), 단위 테스트 생성 (커버리지), 커밋 메시지 생성 (텍스트 유사도). 각 작업은 검증 가능하고 실행 기반의 보상 신호를 활용합니다.

주요 기여

1. 코드 LLM 을 위한 MTRL 에 대한 최초의 연구: 이 논문은 "무엇을 학습할 것인가"와 "어떻게 학습할 것인가"라는 두 가지 결합된 결정을 중심으로 코드에 대한 다중 작업 RL 을 공식화하고, 두 가지를 지배하는 통합 신호로서 작업 유틸리티를 도입합니다.
2. ASTOR 프레임워크: 다음과 같은 기능을 갖춘 조정 프레임워크입니다:
   • 계층적 유틸리티 라우팅 데이터 스케줄링: 학습 잠재력과 시너지에 기반하여 작업 간 예산을 할당하고 프롬프트에 우선순위를 부여합니다.
   • 적응형 유틸리티 보정 정책 최적화: 이질적인 작업의 훈련 상태에 맞춰 작업별 KL 정규화를 동적으로 조정합니다.
3. 실증적 검증: 단일 ASTOR 모델이 작업별 전문 모델 및 기존 MTRL 베이스라인 모두를 능가함을 보여주는 광범위한 실험.

실험 결과

저자들은 네 가지 대표적인 코딩 작업에 걸쳐 두 개의 기본 모델 (Qwen2.5-Coder-7B 및 Qwen3-8B) 에서 ASTOR 을 평가했습니다.

• 단일 작업 학습 (STL) 대비: ASTOR 은 강력한 다중 작업 성능을 달성하여 작업별 전문 모델을 능가하거나 이에 필적합니다. Qwen2.5-Coder-7B 에서는 전체 평균 성능에서 최상의 단일 작업 모델을 9.0% 상회합니다. Qwen3-8B 에서는 개선 폭이 9.5% 입니다.
• MTRL 베이스라인 대비: ASTOR 은 세 가지 베이스라인 (균일 혼합인 Joint Learning, 고정 순차적 순서인 Curriculum Learning, 사후 통합인 Model Merging) 을 크게 능가합니다.
  • 모델 및 지표 전반에 걸쳐 가장 강력한 MTRL 베이스라인 (Joint Learning) 대비 개선 폭은 7.5% 에서 19.2% 사이입니다.
  • 구체적으로 Qwen2.5-Coder-7B 에서 ASTOR 은 Joint Learning 보다 12.8% 더 우수한 성능을 보입니다.
• Ablation 연구: 데이터 스케줄링 모듈 (작업 할당 또는 프롬프트 우선순위 지정) 이나 정책 최적화 모듈 (동적 또는 작업별 정규화) 중 하나를 제거하면 성능이 크게 저하되어, 이질적인 작업을 조정하기 위해 두 구성 요소 모두 필수적임을 확인했습니다.

중요성과 주장

이 논문은 ASTOR 이 다중 작업 코드 RL 에 대한 유틸리티 기반 조정의 효과성을 검증한다고 주장합니다. 학습 잠재력과 시너지에 기반하여 자원을 할당하고 작업별 필요에 따라 업데이트를 제약함으로써 작업을 이질적으로 처리함으로써, 단일 모델이 여러 전문 모델을 배포하는 데 따른 선형적 확장 비용 없이 다양한 역량을 달성할 수 있습니다.

저자들은 이 연구를 효율적이고 통합된 AI 지원 코딩 시스템으로 나아가는 한 걸음으로 위치시킵니다. 이러한 시스템은 단일 모델 내에서 이해, 구현, 테스트, 문서화를 포함한 전체 소프트웨어 개발 워크플로우를 처리할 수 있으며, 모든 단계에서 성능을 유지하거나 향상시키면서 배포 장벽을 낮춥니다. 이 연구는 작업/언어 커버리지의 한계와 기본 KL 계수의 수동 튜닝에 대해 인정하며, 이를 향후 연구 과제로 제안합니다.