Timestep weighting improves reinforcement learning for flow matching models
Timestep Weighting: A Hidden Key to Effective ELBO-Based Flow-Matching RL
Machine LearningArtificial IntelligenceComputer Vision and Pattern Recognition
Summary
Getting computers to learn from rewards can be tricky, especially when using methods involving multiple steps or times. This paper shows that how you weight or emphasize learning at different timesteps greatly affects the success of training. The authors studied this weighting in various tasks like image generation and robotics, finding that the best way to weight depends on how far the current model is from the ideal behavior. They also found ways to choose and adjust these weights dynamically to improve learning outcomes.
What this means in practice
- •For robotics engineers: Improve robot learning tasks by adjusting training emphasis over time to better match rewarded behaviors.
- •For machine learning practitioners: Enhance training of flow matching models in generative tasks by choosing timestep weightings aligned with reward landscapes.
Authors
Qinwei Ma, Jingzhe Shi, Simin Fan, Ling Li, Mengdi Wang, Alex Lamb
Abstract
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.