タイムステップ重み付け:ELBOベースフローマッチング強化学習の有効性を左右する隠れた鍵
Timestep Weighting: A Hidden Key to Effective ELBO-Based Flow-Matching RL
ELBOベースの強化学習におけるタイムステップ重み付けが性能に与える影響を分析し、報酬地形と学習段階に応じた適応的な重み付け手法を提案した。
著者: Qinwei Ma, Jingzhe Shi, Simin Fan, Ling Li, Mengdi Wang, Alex Lamb
分類: cs.LG, cs.AI, cs.CV
原文アブストラクト
ELBO-based reinforcement learning offers a sampler-agnostic approach to fine-tuning flow matching models with reward feedback. Timestep weighting in ELBO-based RL has large impact on performance, and it also provides a unified view (as we show in this work) to understand prediction losses heuristically chosen in prior work, yet it remains under-researched and is often chosen to inherit pretrain configs. We investigate impacts and dynamics of timestep weighting in ELBO-based RL. We show that effective weighting depends on both the reward landscape and stage of learning. (1) Through experiments on controlled CIFAR image generation, complemented by robotics, we investigate how weighting impacts reward-driven updates across noise levels. (2) Through gradient analysis, we reveal distinct patterns of cross-noise coordination across tasks and their evolution during training. These findings motivate the hypothesis that useful weighting depends on the gap between the policy's current behavior and the behavior favored by the reward. (3) Guided by this analysis, we study simple static weighting, budgeted profile selection, and dynamic schedules that improve performance beyond conventional target choices. Our results establish timestep weighting as an important design choice for flow-matching RL and motivate further research into methods that choose and adapt it throughout learning.