Is RL fine-tuning harder than regression? A PDE learning approach for diffusion models

2025-09-02

AI Summary

This work addresses reinforcement learning (RL) fine-tuning of diffusion models to learn optimal control policies that steer a given diffusion process toward desired objectives. We propose a variational inequality framework grounded in the Hamilton–Jacobi–Bellman (HJB) equation, reformulating optimal control as a supervised regression task—thereby circumventing the high sample complexity inherent in conventional RL. Our method integrates PDE theory with universal function approximation, yielding a tractable algorithm that jointly learns the value function and control policy. We establish tight statistical error bounds that explicitly depend on the complexity of the hypothesis class and the approximation error of the value function, and we prove improved convergence rates. The key contribution is the first rigorous reduction of diffusion model fine-tuning to supervised learning—achieving both theoretical soundness and computational efficiency.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We study the problem of learning the optimal control policy for fine-tuning a given diffusion process, using general value function approximation. We develop a new class of algorithms by solving a variational inequality problem based on the Hamilton-Jacobi-Bellman (HJB) equations. We prove sharp statistical rates for the learned value function and control policy, depending on the complexity and approximation errors of the function class. In contrast to generic reinforcement learning problems, our approach shows that fine-tuning can be achieved via supervised regression, with faster statistical rate guarantees.

Problem

• Learning optimal control policy for diffusion model fine-tuning
• Solving variational inequality from Hamilton-Jacobi-Bellman equations
• Demonstrating fine-tuning via supervised regression with faster rates

Innovation

• Solving variational inequality from HJB equations
• Using value function approximation for control
• Achieving fine-tuning via supervised regression