LeapAlign: Послетренировочное согласование потоковых моделей на любом шаге генерации путем построения двухшаговых траекторий
LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
April 16, 2026
Авторы: Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng
cs.AI
Аннотация
Данная работа посвящена проблеме согласования моделей согласования потоков (flow matching) с человеческими предпочтениями. Перспективным подходом является тонкая настройка путем прямой обратной передачи градиентов вознаграждения через дифференцируемый процесс генерации в flow matching. Однако обратное распространение по длинным траекториям приводит к запретительным затратам памяти и взрыву градиентов. Следовательно, методы прямого градиента не справляются с обновлением ранних шагов генерации, которые критически важны для определения глобальной структуры конечного изображения. Для решения этой проблемы мы представляем LeapAlign — метод тонкой настройки, который снижает вычислительные затраты и позволяет осуществлять прямую передачу градиента от вознаграждения к ранним шагам генерации. А именно, мы сокращаем длинную траекторию до всего двух шагов, проектируя два последовательных «прыжка», каждый из которых пропускает несколько шагов дискретизации ОДУ и предсказывает будущие латентные переменные за один шаг. Рандомизируя начальные и конечные моменты времени прыжков, LeapAlign обеспечивает эффективное и стабильное обновление модели на любом шаге генерации. Для лучшего использования таких сокращенных траекторий мы назначаем более высокие веса обучения тем траекториям, которые более согласованы с длинным путем генерации. Для дальнейшего повышения стабильности градиентов мы уменьшаем веса членов градиента с большой величиной, вместо их полного удаления, как это делалось в предыдущих работах. При тонкой настройке модели Flux метод LeapAlign последовательно превосходит современные методы на основе GRPO и прямого градиента по различным метрикам, демонстрируя превосходное качество изображений и соответствие изображения и текста.
English
This paper focuses on the alignment of flow matching models with human preferences. A promising way is fine-tuning by directly backpropagating reward gradients through the differentiable generation process of flow matching. However, backpropagating through long trajectories results in prohibitive memory costs and gradient explosion. Therefore, direct-gradient methods struggle to update early generation steps, which are crucial for determining the global structure of the final image. To address this issue, we introduce LeapAlign, a fine-tuning method that reduces computational cost and enables direct gradient propagation from reward to early generation steps. Specifically, we shorten the long trajectory into only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step. By randomizing the start and end timesteps of the leaps, LeapAlign leads to efficient and stable model updates at any generation step. To better use such shortened trajectories, we assign higher training weights to those that are more consistent with the long generation path. To further enhance gradient stability, we reduce the weights of gradient terms with large magnitude, instead of completely removing them as done in previous works. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods across various metrics, achieving superior image quality and image-text alignment.