🤖 AI Summary
This study addresses the disconnect between theoretical objectives and practical deployment in existing one-step generative models, where training often overlooks the implementation costs of parameter sharing. To bridge this gap, we propose TDAction, a framework that formulates training dynamics as an optimal control problem for the first time, achieving cost-aware optimal transport through local parameter sharing. Specifically, we derive a closed-form Batch Tangent Action-to-Go value function and introduce soft terminal control, stochastic tangent probing, and low-rank approximation techniques to enhance optimization efficiency. By preserving the distributional evolution process while significantly reducing implementation overhead, our method achieves state-of-the-art performance on ImageNet 256×256, yielding an FID below 1.1 without requiring distillation.
📝 Abstract
One-step generative models construct a static generator through iterative training-time transport. Existing transport objectives primarily assess distributional motion, although a neural generator needs to realize the requested sample displacements jointly through shared parameter updates. The training-time construction raises the question: \emph{once training becomes the iterative process that constructs the final one-step map, what to optimize: the next distributional move, or the route by which the finite generator learns the final map?} To address the question, we introduce \textbf{T}raining \textbf{D}ynamics \textbf{A}ction (\textbf{TDAction}), which selects transport targets according to local shared-parameter realization cost while retaining a prescribed level of distributional progress. We formulate the cost as a soft-terminal control problem and derive a closed-form Batch Tangent Action-to-Go value that accounts for parameter effort and terminal mismatch. The criterion captures cross-sample interactions omitted by independent pairwise costs; under isotropic mobility, the criterion agrees with quadratic Euclidean assignment for deterministic balanced couplings. Randomized tangent probes provide a low-rank implementation that constructs shared detached targets without adding an inference-time trajectory. Controlled studies examine the relationship between generator geometry, transport selection, and realized local action. On ImageNet $256\times256$, TDAction attains an FID below $1.1$ without distillation.