기술 요약: LaDi-RL

문제 정의

강화 학습 (RL) 은 특히 추론 궤적의 최적화를 통해 대규모 언어 모델 (LLM) 의 추론 능력을 향상시키는 주된 패러다임으로 자리 잡았습니다. 그러나 기존 방법들은 주로 이산적인 토큰 시퀀스 상의 정책을 최적화합니다. 이러한 접근 방식은 근본적인 불일치를 초래합니다. 즉, 중요한 추론 결정은 종종 의미론적이며 전역적이고 궤적 수준의 특성을 지닌 반면, 토큰 수준의 최적화는 지역적 섭동에 의존하기 때문입니다.

이러한 불일치는 **다양성 붕괴 (diversity collapse)**라는 현상을 야기합니다. 이는 모델의 정책이 좁은 범위의 추론 패턴에 확률 질량을 집중시켜 대체 전략을 억제하는 현상입니다. 토큰 수준의 탐색은 "therefore"를 "so"로 대체하는 것과 같은 어휘적 변형을 생성할 수는 있지만, 대수적 추론에서 기하학적 추론으로 전환하는 것과 같이 의미론적으로 구별되는 전략을 발견하는 데는 어려움을 겪습니다. 결과적으로 RL 로 튜닝된 모델들은 종종 pass@1(단일 샘플 정확도) 은 향상시키지만, 큰 $k$값에서 pass@k 성능을 유지하거나 향상시키지 못하여 다양한 올바른 궤적의 상실을 나타냅니다.

더욱이, 탐색을 연속적인 잠재 공간으로 이동시키면 **할당 문제 (credit-assignment challenge)**가 발생합니다. 잠재 공간 RL 에서 정책은 연속적인 잠재 공간에서 행동하지만, 보상은 잠재 궤적이 텍스트로 디코딩된 후에만 관찰됩니다. 단순한 롤아웃 전략은 잠재 추론의 품질과 텍스트 디코딩의 품질을 얽히게 만듭니다. 낮은 보상은 결함이 있는 잠재 궤적에서 비롯되었을 수도 있고, 단순히 건전한 잠재의 불완전한 텍스트적 실현에서 비롯되었을 수도 있어, 노이즈가 많고 오해의 소지가 있는 학습 신호를 생성합니다.

방법론: LaDi-RL

저자들은 확산 모델을 사용하여 엔트로피 붕괴를 방지하기 위해 연속적인 잠재 공간에서의 탐색을 재정의하는 **강화 학습을 통한 잠재 확산 추론 (Latent Diffusion Reasoning with Reinforcement Learning, LaDi-RL)**이라는 프레임워크를 제안합니다.

1. 잠재 확산 정책

추론을 토큰 단위로 생성하는 대신, LaDi-RL 은 변분 오토인코더 (VAE) 를 사용하여 전체 추론 궤적을 64 토큰과 같은 컴팩트한 연속적인 잠재 블록으로 인코딩합니다. 정책은 이러한 궤적 표현에 대한 확산 모델로 매개변수화됩니다.

• 메커니즘: 노이즈에서 시작하여 정책은 잠재 추론 궤적을 향해 반복적으로 노이즈를 제거합니다. 이 과정은 구조화된 탐색과 표현력 있는 분포 모델링을 가능하게 하여 자연스럽게 다중 모드 추론 전략을 지원합니다.
• 다양성 유도: 붕괴를 더 방지하기 위해 저자들은 반발 기반 유도 메커니즘을 도입합니다. 노이즈 제거 과정에서 인접한 잠재 궤적에 반발력이 가해져 잠재 공간에서 분리되도록 장려합니다. 이는 정책 엔트로피에만 의존하지 않고 명시적으로 다양성을 보존하는 지역적 기하학적 정규화 역할을 합니다.

2. 계층적 잠재 - 텍스트 롤아웃

할당 문제를 해결하기 위해 LaDi-RL 은 계층적 $N \times M$ 롤아웃 전략을 사용합니다:

• 잠재 샘플링 ($N$): 주어진 쿼리에 대해 확산 정책은 $N$개의 서로 다른 잠재 추론 궤적을 샘플링합니다.
• 텍스트 디코딩 ($M$): 샘플링된 각 잠재 궤적에 대해 텍스트 정책은 $M$개의 후보 답변 텍스트를 생성합니다.
• 디코더-주변화 보상: $M$개의 텍스트 완성의 보상을 집계 (평균) 하여 기본 잠재 궤적의 유용성을 추정합니다. 이 "디코더-주변화" 추정치는 텍스트 디코딩의 확률적 요소로부터 잠재 추론의 품질을 분리하여, 확산 정책을 최적화하기 위한 더 깨끗하고 분산이 낮은 보상 신호를 제공합니다.

3. 공동 최적화

이 프레임워크는 그룹 상대 정책 최적화 (GRPO) 를 사용하여 두 개의 이질적인 정책을 동시에 최적화합니다:

• 잠재 정책: 각 잠재 블록과 관련된 $M$개 텍스트의 평균 보상을 사용하여 최적화됩니다. 이 목적 함수는 확산 모델이 다양하고 고품질의 추론 모드를 탐색하도록 장려합니다.
• 텍스트 정책: 동일한 잠재 블록에서 생성된 $M$개 텍스트 간의 상대적 차이를 사용하여 최적화됩니다. 이는 텍스트 정책이 주어진 추론 전략에 대한 최적의 텍스트적 실현을 찾도록 강요하여, 새로운 전략의 탐색과 이를 분리합니다.
• 손실 함수: 잠재 손실과 텍스트 손실의 가중 합은 잠재 정책이 다양한 사고의 사슬을 탐색하는 동안 텍스트 정책은 신뢰할 수 있고 효과적인 생성기로 유지되도록 보장합니다.

주요 결과

저자들은 코드 생성 및 수학 추론 벤치마크에서 LaDi-RL 을 토큰 수준 RL 베이스라인 (예: Vanilla GRPO, 엔트로피 이득이 있는 GRPO) 및 기타 잠재 추론 방법과 비교하여 평가했습니다.

• 성능 향상: LaDi-RL 은 토큰 수준 RL 베이스라인에 비해 코드 생성에서 pass@1 절대 개선율 +9.4%, 수학 추론에서 **+5.7%**를 달성했습니다.
• 다양성 보존: 큰 $k$에서 pass@k 가 기본 모델보다 떨어지는 다양성 붕괴를 보이는 Vanilla GRPO 와 달리, LaDi-RL 은 모든 벤치마크에서 기본 모델의 pass@k 성능을 일관되게 능가했습니다. 예를 들어, AIME 2025 에서 pass@128 기준 절대 개선율 +13.4를 달성했습니다.
• 효율성: 디코딩 전에 추론을 잠재 블록 (64 토큰) 으로 압축함으로써, LaDi-RL 은 생성된 토큰 수를 크게 줄였습니다 (예: 코드 작업의 경우 표준 GRPO 의 약 3,447 토큰에서 LaDi-RL 의 약 422 토큰으로), 결과적으로 엔드 - 투 - 엔드 추론 시간을 33% 단축했습니다.
• 애블레이션 연구:
  • 궤적 대 토큰 수준: 잠재 궤적 공간에서의 탐색은 토큰 수준 (0.18–0.22) 에 비해 훨씬 높은 의미론적 다양성 (0.52) 과 더 강력한 성능을 보여주었습니다.
  • 확산 대 VAE: 확산 기반 정책은 단일 단계 VAE 샘플러보다 우수한 성능을 발휘하여 복잡한 잠재 추론 분포를 학습하기 위한 반복적 노이즈 제거의 중요성을 입증했습니다.
  • 계층적 롤아웃: $N \times M$ 구조가 결정적이었습니다. $N$(잠재 다양성) 을 증가시키는 것은 의미론적 다양성을 주도하는 반면, $M$(텍스트 평균화) 을 증가시키는 것은 더 깨끗한 보상 신호를 제공했습니다.

중요성과 주장

이 논문은 LaDi-RL 이 LLM 추론 향상을 위한 토큰 수준 RL 의 원칙적인 대안을 제공한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

1. 다양성 붕괴 완화: 지역적 토큰 섭동에서 전역적 궤적 수준 확산으로 탐색을 전환함으로써, 이 방법은 다양한 해결책 전략 세트를 유지하여 정책이 비최적 지역 최소값으로 수렴하는 것을 방지합니다.
2. 할당 문제 해결: 계층적 롤아웃과 디코더-주변화 보상 추정은 추론 품질을 디코딩 노이즈에서 효과적으로 분리하여 잠재 정책의 안정적인 최적화를 가능하게 합니다.
3. 효율성과 품질: 이 접근 방식은 잠재 압축을 통해 계산 오버헤드를 크게 줄이면서도 정확도 (pass@1) 와 다양성 (pass@k) 모두에서 최첨단 성능을 달성합니다.

저자들은 잠재 확산이 구조화된 탐색을 위한 강력한 메커니즘을 제공하여, 토큰 수준 방법들이 놓치는 의미론적으로 구별되는 추론 전략을 모델이 발견할 수 있게 함으로써, 해결책의 다양성을 희생하지 않고 더 높은 추론 한계를 열어준다고 결론지었습니다.