FlashREINFORCE란 무엇인가요?
FlashREINFORCE는 에이전트형 언어 모델의 비동기 학습을 위한 비평자 없는 강화학습 프레임워크입니다. 가변 길이 추론, 도구 호출 또는 환경 상호작용으로 인해 궤적이 불규칙한 시간에 종료되는 장기 작업을 위해 설계되었습니다.
FlashREINFORCE는 모든 프롬프트에 대해 여러 형제 롤아웃을 생성하는 대신 프롬프트당 하나의 롤아웃을 사용합니다. 따라서 B개의 궤적을 포함하는 배치는 B개의 서로 다른 프롬프트를 처리할 수 있으며, 완료된 궤적은 동일한 프롬프트에 대한 여러 응답을 기다리지 않고 학습기에 전달될 수 있습니다.
이 저장소는 독립 실행형 PyTorch 기준 손실 함수, CPU 옵티마이저 업데이트 예제, 테스트, R1 및 Qwen2.5-Math용 Molt 실행 도구, 추론·Python 도구·혼합 전문가 모델·ALFWorld를 위한 실험 설정을 제공합니다. 전체 비동기 학습 실행 도구는 Molt를 사용합니다.
강화학습은 REINFORCE를 해야 합니다.
단일 롤아웃을 사용하는 이유
그룹 상대적 방법은 동일한 프롬프트에서 여러 롤아웃을 필요로 합니다. 비동기 에이전트 작업에서는 이러한 요구 사항이 동기화 장벽을 만들고, 고정된 롤아웃 예산으로 처리할 수 있는 서로 다른 프롬프트 수를 줄일 수 있습니다.
단일 롤아웃 학습은 형제 롤아웃 요구 사항을 제거하지만, 세 가지 안정성 문제를 발생시킵니다:
- 프롬프트 수준 기준선을 구성할 형제 그룹이 없습니다.
- 비동기적으로 수집된 궤적은 더 오래된 행동 정책으로 생성되었을 수 있습니다.
- 모든 토큰에 배치 수준의 동일한 가중치를 부여하면, 길고 실패한 궤적이 업데이트를 지배할 수 있습니다.
FlashREINFORCE는 One-Batch REINFORCE, Sequence Trust Region, Sample-Mean Optimization을 사용하여 이러한 문제를 해결합니다.
핵심 구성 요소
1. One-Batch REINFORCE
각 새 배치에서 학습기는 서로 독립적인 프롬프트에 대한 스칼라 보상을 중심화합니다. 궤적 i의 보상이 R_i라면, 어드밴티지는 궤적 보상에서 배치 평균을 뺀 값입니다:
mean_reward = sum(R_j for j in batch) / B
A_i = R_i - mean_reward평균보다 높은 궤적은 양의 어드밴티지를 받고, 평균보다 낮은 궤적은 음의 어드밴티지를 받습니다. 이를 통해 학습된 비평자를 훈련하지 않고도 부호가 있는 피드백을 제공합니다. 학습기는 배치에 대해 옵티마이저 업데이트를 한 번 수행한 후 해당 배치를 폐기합니다.
2. Sequence Trust Region
모든 궤적은 각 샘플링된 토큰을 실제로 생성한 행동 정책의 확률을 유지해야 합니다. 학습기는 현재 정책의 확률을 다시 계산하고 토큰 중요도 비율을 구성합니다:
rho_i_t = exp(current_log_probability_i_t
- behavior_log_probability_i_t)이 비율은 저장된 이력에서 샘플링된 행동의 분포를 보정합니다. 하지만 이력 자체가 더 오래된 행동 정책에서 생성되었다는 사실은 보정하지 않으므로, FlashREINFORCE는 전체 궤적에 걸쳐 누적된 편차도 평가합니다.
각 샘플링된 행동에 대해 행동 확률 p와 현재 확률 q를 사용하여 베르누이 KL 근사값을 계산합니다:
d_i_t = p * log(p / q)
+ (1 - p) * log((1 - p) / (1 - q))토큰 값은 궤적 전체에서 평균을 냅니다. 평균 근사값이 설정된 신뢰 임계값 이하인 경우에만 궤적을 승인합니다:
mean_drift_i = sum(d_i_t for t in trajectory_i) / T_i
mask_i = 1 if mean_drift_i <= delta else 0거부된 궤적은 토큰 손실에 기여하지 않습니다. 따라서 신뢰성 판단은 개별 토큰이 아니라 저장된 전체 이력에 적용됩니다.
3. Sample-Mean Optimization
토큰 수준의 배치 평균은 긴 궤적에 자동으로 더 큰 영향력을 부여합니다. 긴 실패 롤아웃이 유용한 중간 단계를 포함하면서도 음의 보상으로 끝나는 경우 이는 문제가 될 수 있습니다.
FlashREINFORCE는 먼저 각 궤적 내에서 토큰 기여도를 평균 낸 다음, 궤적 전체에서 평균을 냅니다:
trajectory_term_i = (
mask_i * advantage_i / trajectory_length_i
* sum(importance_ratio_i_t for each token t)
)
objective = sum(trajectory_term_i for i in batch) / B따라서 각 궤적은 응답 길이와 관계없이 1/B의 가중치를 받습니다. 실제 손실에서는 분리된 중요도 비율을 사용합니다. 기본 방법은 학습된 비평자, 비율 클리핑 또는 참조 모델 순방향 패스를 사용하지 않습니다.
설치 및 로컬 설정
FlashREINFORCE 저장소를 복제하거나 다운로드하고, 루트 디렉터리로 이동한 뒤 테스트 의존성과 함께 편집 가능한 모드로 설치합니다:
pip install -e '.[test]'다음으로 포함된 CPU 예제를 실행합니다:
python examples/toy_update.py이 예제는 기준 목적 함수가 CPU에서 옵티마이저 업데이트를 생성할 수 있는지 확인합니다. 논문의 벤치마크 결과를 재현하는 것이 아니라 기능을 검증하기 위한 예제입니다.
마지막으로 테스트 모음을 실행합니다:
python -m pytest -q예제와 테스트가 성공적으로 실행되면 로컬 기준 구현과 해당 의존성을 사용할 수 있음을 확인할 수 있습니다.
기준 구현 살펴보기
핵심 기준 목적 함수는 flashreinforce/loss.py에 있습니다. 이 함수는 논문과 README에 설명된 다음 작업을 다룹니다:
- 배치 전체의 보상 중심화.
- 토큰 수준 중요도 샘플링.
- 시퀀스 수준 신뢰 승인.
- 샘플 평균 최적화.
- 실패 토큰의 선택적 엔트로피 기반 필터링.
examples/toy_update.py에서 기준 손실이 옵티마이저 업데이트에 어떻게 참여하는지 확인하는 것으로 시작하세요. 그런 다음 flashreinforce/loss.py를 살펴보며 예제의 텐서와 마스크가 목적 함수의 개별 단계와 어떻게 연결되는지 확인하세요.
CPU 예제는 완전한 분산 롤아웃 인프라가 필요하지 않으므로 가장 안전한 기본 사용 경로입니다.
단일 패스 비동기 업데이트 이해하기
프로덕션 FlashREINFORCE 업데이트는 다음 6단계로 진행됩니다.
- 수집: 롤아웃 워커는 프롬프트마다 완료된 궤적 하나를 독립적으로 제출합니다. 각 궤적에는 보상과 토큰 샘플링에 실제로 사용된 행동 정책 확률이 포함됩니다.
- 배치: 학습기는 다음 B개의 완료된 궤적을 가져와 보상을 중심화하고 부호가 있는 어드밴티지를 얻습니다.
- 보정: 학습기는 현재 정책의 로그 확률을 다시 계산하고 토큰 중요도 비율을 구성합니다.
- 신뢰성 확인: 각 궤적에 대해 샘플링된 행동의 평균 KL 프록시를 계산하고, 그 값이 설정된 임계값 이내인 경우에만 전체 궤적을 허용합니다.
- 최적화: 허용된 각 궤적 내부에서 평균을 구하고, 배치 전체에서 다시 평균을 구한 뒤 옵티마이저 업데이트를 한 번 수행합니다.
- 폐기: 수집된 배치를 추가 학습기 업데이트에 재사용하지 않습니다.
마지막 폐기 단계는 이 방법의 중요한 부분입니다. FlashREINFORCE는 동일한 행동 정책 스냅샷에서 여러 순차 미니배치 업데이트를 실행하는 대신, 새로 수집한 각 배치마다 전체 배치 업데이트를 한 번 수행합니다.
Molt 학습 구성 미리 보기
저장소에는 Molt의 비동기 구현을 통해 학습을 시작하는 scripts/train_molt.py가 포함되어 있습니다. 전체 학습을 할당하기 전에 드라이 런 모드를 사용해 학습 플래그를 미리 확인하세요.
python scripts/train_molt.py --dry-runGPU 설정, 데이터 준비, 학습 명령, 논문과 코드 간 매핑, 알려진 재현 제한 사항은 docs/training.md를 참조하세요. examples/README.md 가이드에는 작업 프리셋, 절제 실험, 기능 검사가 포함된 네이티브 트레이너 실행기가 있습니다.
R1 및 Qwen2.5-Math 실행기는 저장소에 고정되어 있습니다. 도구 사용 및 ALFWorld 설정에는 호환 가능한 에이전트와 환경도 필요하므로, 해당 실험을 시작하기 전에 런타임 요구 사항을 검토하세요.
각 궤적이 보존해야 하는 데이터
업데이트는 생성 중 캡처된 정보에 의존합니다. 개념적인 학습 흐름에는 최소한 완료된 토큰 시퀀스, 스칼라 보상, 그리고 샘플링된 모든 행동에 대응하는 행동 확률이 필요합니다. 그런 다음 학습기는 저장된 행동과 이력을 평가하여 해당하는 현재 정책 확률을 얻습니다.
기록된 샘플링 확률을 나중에 재구성한 확률로 대체하지 마세요. 나중에 모델을 평가해도 추론 엔진이 실제 샘플링에 사용한 확률이 재현된다는 보장은 없습니다.
고급 안정성 팁
샘플링 지지집합 보존
정확한 중요도 보정은 대상 정책이 행동 정책의 지지집합 바깥에 확률 질량을 할당하지 않는다고 가정합니다. 공격적인 top-k 또는 top-p 절단은 이 가정을 위반할 수 있습니다. 따라서 지지집합 요구 사항을 고려해 샘플링 설정을 선택해야 합니다.
토큰 중요도 샘플링을 전체 궤적 보정으로 간주하지 않기
토큰 중요도 비율은 저장된 이력에 조건부인 샘플링된 행동을 보정합니다. 하지만 해당 이력의 분포는 보정하지 않습니다. 누적된 행동 정책과 학습기 간 차이가 지나치게 큰 궤적을 걸러내려면 Sequence Trust Region을 사용하세요.
업데이트를 최신 상태로 유지하기
수집한 각 배치마다 옵티마이저 업데이트를 한 번 수행한 다음 해당 배치를 폐기하세요. 오래된 배치를 여러 학습기 업데이트에 재사용하면 문서화된 단일 패스 방식에서 벗어나며 정책 불일치가 커질 수 있습니다.
궤적 수준 가중치 보존
배치 전체에서 궤적을 평균 내기 전에 각 궤적 내부의 토큰 기여도를 평균 내세요. 유효한 모든 토큰을 하나의 배치 전체 평균으로 평탄화하면 응답 길이에 따른 승수가 다시 생기고 긴 궤적에 과도한 가중치가 부여될 수 있습니다.
신뢰 마스크를 올바르게 해석하기
Sequence Trust Region은 전체 궤적에 대해 하나의 허용 여부를 결정합니다. 궤적이 임계값을 초과하면 해당 궤적의 모든 토큰 손실이 해당 업데이트에서 제거됩니다.
검증과 재현을 구분하기
CPU 예제는 옵티마이저 업데이트만 검증합니다. 연구 결과를 재현하려면 학습 및 예제 가이드에 문서화된 GPU, 데이터, 모델, 에이전트, 환경, 비동기 런타임 구성이 필요합니다.
보고된 적용 분야
보고된 실험은 수학적 추론, 다중 턴 Python 도구 사용, 전문가 혼합 학습, 대화형 ALFWorld 의사 결정을 다룹니다. README에는 DeepSeek-R1-Distill-Qwen-1.5B를 사용한 6,000회 업데이트까지의 안정적인 긴 연쇄 사고 학습, Qwen2.5-Math-1.5B의 수학적 추론 결과, Qwen2.5-7B-Instruct 및 Qwen3-30B-A3B의 Python 도구 실험, Qwen2.5-7B-Instruct의 ALFWorld 실험이 보고되어 있습니다.
이러한 결과는 프레임워크의 의도된 범위를 보여 주지만, 소규모 CPU 업데이트 예제에서 동일한 결과가 나올 것으로 기대해서는 안 됩니다.
유용한 리소스
결론
FlashREINFORCE는 단일 롤아웃 프롬프트 범위와 간결한 비평가 없는 업데이트를 결합합니다. 배치 중심 보상은 부호가 있는 어드밴티지를 제공하고, 토큰 중요도 샘플링은 오래된 행동을 처리하며, Sequence Trust Region은 누적된 궤적 드리프트를 제어합니다. 또한 샘플 평균 최적화는 응답 길이가 궤적 가중치를 결정하는 것을 방지합니다.
편집 가능한 설치, CPU 업데이트 예제 및 테스트부터 시작하세요. 해당 작업이 완료되면 --dry-run을 사용해 Molt 런처를 미리 실행하고, 모델별 GPU 및 데이터 준비에 대해서는 학습 가이드를 따르세요. 이 프로젝트는 Apache License 2.0에 따라 제공됩니다.