Optimal Self-Distillation for Rectified Flow
via Linear Probing
July 16, 2026
Modern generative models are increasingly trained using model-generated signals, creating both opportunities for self-improvement and risks of collapse. We study optimal self-distillation (SD) for rectified flow (RF): given a suboptimal teacher velocity field, can a student trained on a mixture of true RF velocities and teacher velocities provably improve the teacher? For linear RF with ridge regularization on fixed interpolation pairs, we prove an exact affine path identity, derive the optimal mixing coefficient in closed form, and show strict improvement in integrated velocity risk whenever the teacher risk is nonstationary along the regularization path. The optimal coefficient obeys a sign rule: positive mixing corrects under-regularized teachers, while negative mixing corrects over-regularized teachers. We also give one-shot generalized cross-validation (GCV) and validation tuning procedure that avoids grid search over mixing weights and repeated refitting. Combining this theorem with RF Wasserstein convergence bounds, we show that optimal self-distillation improves the velocity estimation terms controlling continuous-time and finite-step generation error. Experiments with Gaussian models, Gaussian mixtures, and image data show that optimal self-distillation improves velocity risk, mode recovery, and finite-step generation relative to both the teacher and pure distillation.
Synthetic data and model-generated supervision are increasingly common in modern AI pipelines, with seemingly opposing consequences. Recursively training generative models on their own outputs can lead to model collapse, loss of diversity, or distributional drift [1]–[3]. Yet model-generated signals can also be useful: in supervised learning, self-distillation often improves a model by retraining on a mixture of ground-truth labels and teacher predictions. We ask whether this positive self-improvement phenomenon can be isolated and analyzed in a generative model.
We focus on rectified flow (RF) within the broader flow matching framework [4]–[6], which offers a simple approach to generative modeling with strong empirical performance [7]–[9]. Let \(X_0\sim \rho_0\) be source noise and \(X_1\sim \rho_1\) be data, and define \[\label{eq:linear-interp} X_t :=(1-t)X_0+tX_1, \qquad u=X_1-X_0.\tag{1}\] RF learns a time-conditioned velocity field \(v(t,x)\) by regressing \(u\) on \((t,X_t)\) with squared loss, thereby approximating a suitable velocity field \(v^\star(t,x)\). This problem has the same geometry as supervised regression, but the learned field subsequently defines an ODE sampler. RF is therefore a natural setting in which to ask whether a student trained on mixed teacher-generated and true velocity targets can provably improve a suboptimal RF teacher.
The answer is yes for a linear RF model trained on fixed interpolants. We adopt linear probing [10], [11], an efficient adaptation strategy that freezes a pretrained representation or feature map and updates only a lightweight linear head \(W\). We keep the interpolation covariates fixed, fit a teacher \(\widehat W_\lambda\) with ridge regularization on the true velocity targets, evaluate it at those interpolation points, and train a student on mixed targets (see Figure 2). For a true velocity \(u\), a teacher prediction \(\widehat u\), and a mixing coefficient \(\xi\in\mathbb{R}\), the mixed target is \(u^{(\xi)}=(1-\xi)u+\xi\widehat u\). Unlike recursive Reflow, which generates new noise–data pairs and changes the endpoint coupling, keeping the interpolants fixed preserves the ridge geometry: the self-distilled student lies on an affine path between the teacher and a pure-distilled refit. Consequently, the integrated RF risk is an exact quadratic in \(\xi\), yielding an optimum in closed form and a theorem guaranteeing strict improvement.
Optimal self-distillation for RF velocity learning. Section 3 introduces self-distillation of RF velocity targets on fixed interpolants; Figure 2 and Algorithm 4 summarize the procedure. For linear RF with ridge regularization, we prove an exact affine path identity linking the teacher, pure-distilled refit, and self-distilled student (Proposition 1). This identity makes the integrated RF risk an exact quadratic in \(\xi\), yielding an optimum in closed form and strict improvement whenever the teacher risk is nonstationary along the regularization path (Theorem 2). We also propose one-shot tuning procedure based on GCV or validation that avoids grid search and repeated refitting, providing an efficient adaptation strategy for RF models based on fixed features or linear probing.
Consequences for generation error. Section 4 connects the improvement in velocity risk to generative performance. RF regression risk decomposes into excess velocity approximation error and an irreducible term, so the strict reduction in RF risk also reduces the velocity error quantity entering RF sampling bounds. Combining this identity with Wasserstein convergence results for RF, we show that optimal SD improves the continuous-time generation error upper bound (Proposition 3) and the velocity estimation term in finite-step Euler bounds (Theorem 4).
Empirical validation. Section 5 validates the theory using Gaussian models, nonlinear Gaussian mixtures, random features, and controlled RF experiments on real images. Across these settings, optimal SD improves velocity risk and finite-step generation, one-shot tuning closely tracks the oracle choice, and unconstrained negative mixing is essential for correcting over-regularized teachers; see Figures 1, 3, and 5, as well as Table 1. A time-conditioned U-Net experiment on CIFAR-10 tests the same correction beyond linear probing with fixed features by comparing the teacher, pure distillation, clipped SD, and unconstrained SD. The supplement reports metric definitions, sensitivity to teacher quality and output scaling, and a fine-tuning baseline with matched compute.
Rectified flow and flow matching. Flow matching and stochastic interpolants learn velocity fields by squared regression against conditional velocity targets [4], [6]. Couplings inspired by optimal transport and careful path design can produce straighter flows, stabilize training, and accelerate sampling [12]–[14]; complementary deterministic sampling error bounds depend on velocity field approximation error and flow regularity [15]. RF specializes this framework to linear interpolants and uses the learned ODE to transport source noise to the data distribution [5]. A Wasserstein analysis of RF further shows that terminal sampling error is controlled by integrated velocity estimation error and an Euler discretization term governed by straightness [16]. We use this regression view of estimation error to analyze optimal self-distillation of the learned velocity field.
Self-consuming generative models and collapse. Model collapse under recursive synthetic data training has been observed and studied across a range of settings [1], [2], [17]–[20]. In RF, recursive Reflow repeatedly changes the endpoint coupling; proposed safeguards include mixing real and synthetic endpoint pairs [3] and conic reflow, which learns a spherical linear interpolation between real data and its inverse noise [21]. Our complementary setting keeps the endpoint coupling and interpolation covariates fixed, isolating a one-shot target mixing mechanism that is analytically tractable and computationally efficient.
Self-distillation and optimal mixing. Self-distillation retrains a student on the same examples using a mixture of ground-truth labels and teacher predictions. Classical distillation uses soft teacher predictions to train a student, often transferring teacher performance to a smaller or retrained model [22]–[26]. Theory explains self-distillation gains through implicit regularization, bias–variance tradeoffs, teacher mimicry, and repeated distillation effects [27]–[29]. Related work also studies uncertainty-aware Bayesian knowledge distillation [30]; [31] survey distillation methods for diffusion models. For ridge regression, [32] show that unconstrained optimal mixing can strictly improve any nonstationary teacher along the ridge path and that the optimal coefficient can be negative in over-regularized regimes. We transfer this structural geometry from supervised ridge regression to RF velocity matching.
Let \(\rho_0\) be the source distribution, typically \(\mathcal{N}(0,I_d)\), and \(\rho_1\) the data distribution. We sample endpoint pairs \((X_0,X_1)\) from a coupling \(\pi\) of \((\rho_0,\rho_1)\), usually the independent coupling \(\pi=\rho_0\otimes\rho_1\). For \(t\in[0,1]\), define the linear interpolant and pairwise velocity as in 1 . RF learns a time-dependent velocity field by minimizing \[\label{eq:rf-loss} \mathcal{L}_{\mathrm{RF}}(v) =\mathbb{E}_{T\sim\nu,\,(X_0,X_1)\sim\pi} \big[\left\lVert u-v(T,X_T)\right\rVert^2\big],\tag{2}\] where \(\nu\) is usually uniform on \([0,1]\). The population minimizer is \(v^\star(t,x)=\mathbb{E}[u\mid X_t=x]\). If \(Z_t\) solves \[\label{eq:true-rf-ode} \frac{\mathrm{d}Z_t}{\mathrm{d}t}=v^\star(t,Z_t),\qquad Z_0\sim\rho_0,\tag{3}\] then \(Z_t\) has the same marginal law as \(X_t\) for every \(t\in[0,1]\) [5]. The projection identity under squared loss gives \[\label{eq:irreducible} \mathbb{E}\left\lVert u-v(T,X_T)\right\rVert^2 = \mathbb{E}\left\lVert v^\star(T,X_T)-v(T,X_T)\right\rVert^2 + \mathbb{E}\left\lVert u-v^\star(T,X_T)\right\rVert^2.\tag{4}\] Thus RF regression risk equals integrated velocity approximation error plus an irreducible conditional variance term.
We study the linear (in parameters) model \[\label{eq:linear-model} v_W(t,x)=W^\top\phi(t,x),\qquad W\in\mathbb{R}^{m\times d},\tag{5}\] where \(\phi:[0,1]\times\mathbb{R}^d\to\mathbb{R}^m\) is a fixed time-conditioned feature map. For example, let \(b(t)=(b_1(t),\ldots,b_L(t))^\top\) be fixed temporal basis functions and set \(\phi(t,x)=b(t)\otimes x\). Writing \(W=(A_1^\top,\ldots,A_L^\top)^\top\) with \(A_\ell\in\mathbb{R}^{d\times d}\) gives \(v_W(t,x)=A_W(t)x\), where \(A_W(t):=\sum_{\ell=1}^L b_\ell(t)A_\ell\). This yields a time-dependent linear velocity field with parameters shared across time. More generally, the fixed map \(\phi\) may comprise random features, kernel or Nyström approximations, or frozen pretrained representations [33]; the affine path identity below applies to any fixed choice of \(\phi\). Once \(\phi\) is fixed, training the RF teacher reduces to estimating the shared coefficient matrix \(W\).
Given training triples \(\mathcal{D}_N=\{(X_{0,i},X_{1,i},T_i)\}_{i=1}^N\) with \(T_i\stackrel{\mathrm{iid}}{\sim}\nu\), define \(X_{T_i,i}=(1-T_i)X_{0,i}+T_iX_{1,i}\). Let \(u_i=X_{1,i}-X_{0,i}\) and \(\phi_i=\phi(T_i,X_{T_i,i})\). We use plain italic notation throughout: \(u\) and \(u_i\) are velocity vectors, while \(U\) denotes the stacked velocity target matrix. Specifically, \(\Phi=(\phi_1,\ldots,\phi_N)^\top\in\mathbb{R}^{N\times m}\) and \(U=(u_1,\ldots,u_N)^\top\in\mathbb{R}^{N\times d}\). The RF teacher fitted with ridge regularization is \[\label{eq:teacher} \widehat W_\lambda =\mathop{\mathrm{argmin}}_W \, \frac{1}{N}\left\lVert U-\Phi W\right\rVert_\mathrm{F}^2+\lambda\left\lVert W\right\rVert_\mathrm{F}^2 =\Big(\frac{\Phi^\top\Phi}{N}+\lambda I_m\Big)^{-1}\frac{\Phi^\top U}{N}.\tag{6}\] For any learned velocity field \(v\), define the conditional integrated RF risk \[\label{eq:conditional-risk} \mathcal{R}_\nu(v\mid\mathcal{D}_N) =\mathbb{E}\big[\left\lVert u-v(T,X_T)\right\rVert^2\mid \mathcal{D}_N\big],\tag{7}\] where the expectation is over an independent test triple with \(T\sim\nu\) and \((X_0,X_1)\sim\pi\). For a linear estimator \(\widehat W\), we use the shorthand \(\mathcal{R}_\nu(\widehat W\mid\mathcal{D}_N):=\mathcal{R}_\nu(v_{\widehat W}\mid\mathcal{D}_N)\). When \(\nu\) is uniform on \([0,1]\), we omit the subscript and write \(\mathcal{R}(v\mid\mathcal{D}_N)\).
The teacher’s fitted velocities at the original RF interpolation points form \(\widehat U_\lambda=\Phi\widehat W_\lambda\). Pure distillation retrains on these targets: \[\label{eq:pd} \widehat W_{\mathrm{pd},\lambda} =\mathop{\mathrm{argmin}}_{W} \, \frac{1}{N}\left\lVert \widehat U_\lambda-\Phi W\right\rVert_\mathrm{F}^2+\lambda\left\lVert W\right\rVert_\mathrm{F}^2.\tag{8}\] For \(\xi\in\mathbb{R}\), define the mixed velocity target matrix \[U^{(\xi)}=(1-\xi)U+\xi\widehat U_\lambda, \label{eq:mixed-target}\tag{9}\] and let \(\widehat W_{\mathrm{sd},\lambda,\xi}\) be the ridge RF fit on \((\Phi,U^{(\xi)})\). For compactness, throughout the linear theory we write \[v_\lambda:=v_{\widehat W_\lambda},\qquad v_{\mathrm{pd},\lambda}:=v_{\widehat W_{\mathrm{pd},\lambda}},\qquad v_{\mathrm{sd},\lambda,\xi}:=v_{\widehat W_{\mathrm{sd},\lambda,\xi}}.\]
Figure 3: The sign rule in Gaussian RF. (a) In a correctly specified Gaussian RF model, one-shot GCV closely tracks oracle mixing and selects negative mixing in over-regularized regimes. (b) Constraining \(\xi\in[0,1]\) often returns the teacher when the optimal correction is negative.. a — One-shot GCV tuning., b — Unconstrained mixing is necessary.
Proposition 1 (Affine path). For every \(\lambda>0\) and \(\xi\in\mathbb{R}\), \[\widehat W_{\mathrm{sd},\lambda,\xi} = (1-\xi)\widehat W_\lambda+\xi\widehat W_{\mathrm{pd},\lambda}. \label{eq:affine-path}\qquad{(1)}\] Consequently, for every \((t,x)\), \[v_{\mathrm{sd},\lambda,\xi}(t,x) = (1-\xi)v_\lambda(t,x)+\xi v_{\mathrm{pd},\lambda}(t,x).\]
The affine family in Proposition 1 is reminiscent of classifier-free guidance (CFG) in diffusion and flow models [34], [35], but the endpoints and purpose differ. CFG typically extrapolates from an unconditional field toward a conditional field to strengthen conditioning. Here the path joins a teacher trained on true targets and its pure-distilled refit; when \(\xi<0\), the estimator moves away from the pure-distilled field to correct repeated shrinkage. Negative guidance has also been explored in diffusion models for suppressing unwanted features [36]. For notational simplicity, define \[\label{eq:r-definitions} R(\lambda)=\mathcal{R}(\widehat W_\lambda\mid\mathcal{D}_N),\quad R_{\mathrm{pd}}(\lambda)=\mathcal{R}(\widehat W_{\mathrm{pd},\lambda}\mid\mathcal{D}_N),\quad R_{\mathrm{sd}}(\lambda,\xi)=\mathcal{R}(\widehat W_{\mathrm{sd},\lambda,\xi}\mid\mathcal{D}_N).\tag{10}\] Let \((T,X_T,u)\) denote an independent test triple and define \[\begin{align} C(\lambda) &= \mathbb{E}\big[ \left\langle u-v_\lambda(T,X_T),u-v_{\mathrm{pd},\lambda}(T,X_T)\right\rangle \mid\mathcal{D}_N \big], \tag{11}\\ D(\lambda) &= \mathbb{E}\big[ \left\lVert v_\lambda(T,X_T)-v_{\mathrm{pd},\lambda}(T,X_T)\right\rVert^2 \mid\mathcal{D}_N \big]. \tag{12} \end{align}\] By expansion, \(D(\lambda)=R(\lambda)+R_{\mathrm{pd}}(\lambda)-2C(\lambda)\ge 0\).
Theorem 2 (Optimal self-distillation for integrated RF risk). Fix \(\lambda>0\) and assume \(D(\lambda)>0\). Then, for every \(\xi\in\mathbb{R}\), \[R_{\mathrm{sd}}(\lambda,\xi) = R(\lambda) - 2\xi\{R(\lambda)-C(\lambda)\} + \xi^2D(\lambda). \label{eq:risk-quadratic}\tag{13}\] Consequently, \[\xi^\star(\lambda) = \frac{R(\lambda)-C(\lambda)}{D(\lambda)}, \qquad R_{\mathrm{sd}}^\star(\lambda) = R(\lambda) - \frac{(R(\lambda)-C(\lambda))^2}{D(\lambda)}. \label{eq:optimal-xi}\tag{14}\] Moreover, if \(R(\lambda)\) is differentiable along the ridge path, then \[\xi^\star(\lambda) = -\frac{\lambda}{2}\frac{R'(\lambda)}{D(\lambda)}, \qquad R_{\mathrm{sd}}^\star(\lambda) = R(\lambda) - \frac{\lambda^2}{4}\frac{(R'(\lambda))^2}{D(\lambda)}. \label{eq:sign-rule}\tag{15}\] Thus, whenever \(R'(\lambda)\neq0\), \(R_{\mathrm{sd}}^\star(\lambda)<R(\lambda)\) and \(\mathrm{sign}(\xi^\star(\lambda))=-\mathrm{sign}(R'(\lambda))\).
The sign rule determines the correction direction: on the increasing side of the ridge risk curve, \(R'(\lambda)>0\) and \(\xi^\star(\lambda)<0\), so the estimator extrapolates away from the pure-distilled field; on the decreasing side, \(\xi^\star(\lambda)>0\). Although \(\xi<0\) makes 9 an affine rather than convex target mixture, fitting the student remains an ordinary ridge regression problem for any real \(\xi\). Appendix 10 gives the time-dependent extension, its exact gain, and numerical comparisons.
For the linear ridge model, define the smoother matrix \[S_\lambda = \Phi \Big(\frac{\Phi^\top\Phi}{N}+\lambda I_m\Big)^{-1} \frac{\Phi^\top}{N}.\] Then \(\widehat U_\lambda=S_\lambda U\) and \(\widehat U_{\mathrm{pd},\lambda}:=\Phi\widehat W_{\mathrm{pd},\lambda}=S_\lambda^2U\). Following generalized cross-validation (GCV) for ridge regression (see, e.g., [37], [38]), we estimate \(R\), \(R_{\mathrm{pd}}\), and \(C\) by \[\widehat R = \frac{\left\lVert U-S_\lambda U\right\rVert_\mathrm{F}^2/N}{\{1-{\sf Tr}(S_\lambda)/N\}^2}, \; \widehat R_{\mathrm{pd}} = \frac{\left\lVert U-S_\lambda^2U\right\rVert_\mathrm{F}^2/N}{\{1-{\sf Tr}(S_\lambda^2)/N\}^2}, \; \widehat C = \frac{ \left\langle U-S_\lambda U,U-S_\lambda^2U\right\rangle_\mathrm{F}/N }{ \{1-{\sf Tr}(S_\lambda)/N\} \{1-{\sf Tr}(S_\lambda^2)/N\} } ,\] where \(\left\langle A,B\right\rangle_\mathrm{F}={\sf Tr}(A^\top B)\). Plugging these estimates into 14 gives \[\widehat\xi_{\mathrm{GCV}} = \frac{\widehat R-\widehat C}{\widehat R+\widehat R_{\mathrm{pd}}-2\widehat C}. \label{eq:gcv-xi}\tag{16}\]
For neural or other nonlinear RF models, train a teacher field \(v\) and a pure-distilled field \(v_{\mathrm{pd}}\) once, then choose \(\xi\) by validation over the affine family \[v_\xi(t,x) = v(t,x)+\xi\{v_{\mathrm{pd}}(t,x)-v(t,x)\}. \label{eq:validation-mixture}\tag{17}\] This one-dimensional search avoids retraining a student for every candidate \(\xi\). Algorithm 4 summarizes the resulting one-shot procedure.
This section shows how the improvement in velocity risk from optimal self-distillation tightens standard Wasserstein generation error bounds, both in continuous time and for Euler sampling.
Let \(v\) be any learned velocity field and let \(\rho_v=\mathrm{Law}(Y_1^v)\), where \[\label{eq:ode-y-v} \frac{\mathrm{d}Y_t^v}{\mathrm{d}t}=v(t,Y_t^v),\qquad Y_0^v\sim\rho_0 .\tag{18}\] For each time \(t\), define the excess velocity approximation error \[\label{eq:excess-velocity-error} \mathcal{E}_t(v)=\mathbb{E}\big[\left\lVert v^\star(t,X_t)-v(t,X_t)\right\rVert^2\big],\qquad \mathcal{E}(v)=\int_0^1 \mathcal{E}_t(v)\,\mathrm{d}t.\tag{19}\] By 4 , \(\mathcal{R}(v\mid\mathcal{D}_N)\) and \(\mathcal{E}(v)\) differ only by the irreducible term \(\mathbb{E}\left\lVert u-v^\star(T,X_T)\right\rVert^2\), which is independent of \(v\). Thus, any decrease in RF regression risk yields the same decrease in \(\mathcal{E}(v)\).
The next proposition combines this identity with a standard RF Wasserstein bound.
Proposition 3 (Continuous-time generation bound improvement). Assume that \(v(t,\cdot)\) is \(L\)-Lipschitz in \(x\), uniformly over \(t\), and that the ODEs 3 and 18 have unique solutions. Then \[\mathcal{W}_2^2(\rho_v,\rho_1) \le \exp(1+2L)\,\mathcal{E}(v). \label{eq:continuous-w2-bound}\qquad{(2)}\] Consequently, for \(v_{\mathrm{sd}}^\star:=v_{\mathrm{sd},\lambda,\xi^\star(\lambda)}\), the teacher \(v_\lambda\), and the pure-distilled field \(v_{\mathrm{pd},\lambda}\), \[\mathcal{E}(v_{\mathrm{sd}}^\star)\le \min\{\mathcal{E}(v_\lambda),\mathcal{E}(v_{\mathrm{pd},\lambda})\},\] and, using a common Lipschitz bound \(L\) for these fields, the reduction relative to the teacher’s bound is \[\exp(1+2L)\{\mathcal{E}(v_\lambda)-\mathcal{E}(v_{\mathrm{sd}}^\star)\} =\exp(1+2L)\frac{\lambda^2}{4}\frac{(R'(\lambda))^2}{D(\lambda)}. \label{eq:bound-gain-cont}\qquad{(3)}\]
The inequality in ?? follows from [16]. Proposition 3 combines it with Theorem 2: optimal self-distillation weakly improves the continuous-time Wasserstein upper bound relative to both the teacher and pure distillation, with strict improvement over the teacher whenever \(R'(\lambda)\ne0\).
For finite-step RF sampling, the same idea applies to the gridwise velocity errors. Let \(t_i=i/K\), \(i=0,\ldots,K\), be the uniform Euler grid, and let \(\widehat\rho_{v,K}\) be the law of the output of \[\label{eq:emp-ode-disc} \widehat Y_{t_{i+1}}=\widehat Y_{t_i}+(t_{i+1}-t_i)v(t_i,\widehat Y_{t_i}), \qquad \widehat Y_0\sim\rho_0.\tag{20}\] At a fixed grid time \(t_i\), the self-distilled conditional risk has the decomposition \[\label{eq:r-ti} R_{\mathrm{sd},t_i}(\lambda,\xi)=R_{t_i}(\lambda)-2\xi A_{t_i}(\lambda)+\xi^2 D_{t_i}(\lambda), \qquad A_{t_i}(\lambda):= R_{t_i}(\lambda)-C_{t_i}(\lambda),\tag{21}\] where \(R_{\mathrm{sd},t_i}(\lambda,\xi)\), \(R_{t_i}(\lambda)\), \(C_{t_i}(\lambda)\), and \(D_{t_i}(\lambda)\) are defined as in 10 , 11 , 12 , and 13 , with expectations conditioned on \(T=t_i\). Define the gridwise errors \[\label{eq:error-definitions} \varepsilon_i^2(v):=\mathcal{E}_{t_i}(v),\qquad \bar\varepsilon_K^2(v):=\frac{1}{K}\sum_{i=0}^{K-1}\varepsilon_i^2(v).\tag{22}\] All sums below run over \(i=0,\ldots,K-1\). When \(\sum_iD_{t_i}(\lambda)>0\), the coefficient minimizing the average gridwise error and its gain are \[\begin{align} \xi_{\mathrm U,K}^\star(\lambda) &=\frac{\sum_i A_{t_i}(\lambda)}{\sum_iD_{t_i}(\lambda)},\\ \bar\varepsilon_K^2(v_\lambda)-\bar\varepsilon_K^2(v_{\mathrm{sd},\lambda,\xi_{\mathrm U,K}^\star}) &=\frac{\{\sum_i A_{t_i}(\lambda)\}^2}{K\sum_i D_{t_i}(\lambda)}. \end{align} \label{eq:grid-average-gain}\tag{23}\] Write \(v_{\mathrm{sd},K}^\star:=v_{\mathrm{sd},\lambda,\xi_{\mathrm U,K}^\star(\lambda)}\).
Let \(\mathcal{Z}^\star:=\{Z_t\}_{t\in[0,1]}\) denote the exact RF trajectory generated by 3 . Following [16], define its piecewise straightness by \[\gamma_{2,K}(\mathcal{Z}^\star) :=\max_{0\le i<K}\frac{1}{t_{i+1}-t_i} \int_{t_i}^{t_{i+1}}\mathbb{E}\left\lVert \frac{\mathrm{d}}{\mathrm{d}t}v^\star(t,Z_t)\right\rVert_2^2\,\mathrm{d}t.\]
Theorem 4 (Euler generation bound improvement). Under the assumptions of Proposition 3, the sampling distribution \(\widehat\rho_{v,K}\) produced by 20 satisfies, conditionally on \(\mathcal{D}_N\), \[\label{eq:euler-w2-bound} \mathcal{W}_2^2(\widehat\rho_{v,K},\rho_1) \le C_L\left\{\gamma_{2,K}(\mathcal{Z}^\star)/K^2 + \bar\varepsilon_K^2(v)\right\},\tag{24}\] where \(C_L=6\exp(1+\sqrt{1+12L^2})/(1+\sqrt{1+12L^2})\). If \(\sum_iD_{t_i}(\lambda)>0\), then for the gridwise optimally mixed field \(v_{\mathrm{sd},K}^\star\), the teacher \(v_\lambda\), and the pure-distilled field \(v_{\mathrm{pd},\lambda}\), \[\bar\varepsilon_K^2(v_{\mathrm{sd},K}^\star) \le \min\{\bar\varepsilon_K^2(v_\lambda),\bar\varepsilon_K^2(v_{\mathrm{pd},\lambda})\},\] and, using a common Lipschitz bound \(L\), the reduction in the velocity error term for the teacher is \[C_L\{\bar\varepsilon_K^2(v_\lambda)-\bar\varepsilon_K^2(v_{\mathrm{sd},K}^\star)\} =C_L\frac{\{\sum_i A_{t_i}(\lambda)\}^2}{K\sum_i D_{t_i}(\lambda)}. \label{eq:emp-bound-gain-cont}\tag{25}\]
Equation 23 is the gridwise analogue of ?? ; whenever \(\sum_i A_{t_i}(\lambda)\ne0\), it strictly reduces the velocity error term and hence the corresponding upper bound relative to the teacher. Appendix 8.2 proves the result; Appendix 8.2.2 gives the gridwise derivation. The straightness term in 24 is common to all learned estimators for the same target RF path [16]; self-distillation affects only the velocity approximation error on the grid. These results establish an improvement in the RF velocity error component of standard generation error bounds. To complement these general upper bounds, Appendix 8.3 gives exact Wasserstein and KL errors in a Gaussian linear model and shows explicitly how they depend on the integrated RF excess risk. Neither the bounds nor this Gaussian calculation imply that every downstream sample quality metric must strictly improve in every misspecified model: finite-step generation also depends on Lipschitz constants, straightness, ODE stability, discretization, and the expressivity of the learned field. Accordingly, our experiments report both velocity risks and generation metrics.
| Method | RF risk \(\downarrow\) | FID \(\downarrow\) | KID \(\downarrow\) |
|---|---|---|---|
| Teacher | 0.2435 | 284.26 | 0.3215 |
| Pure distillation | 0.3885 | 324.77 | 0.3569 |
| SD clipped \([0,1]\) | 0.2435 | 283.56 | 0.3196 |
| Optimal SD | 0.1803 | 30.08 | 0.0186 |
We evaluate the theory in Gaussian and nonlinear mixture models, then examine finite-step generation and RF experiments on real images. All experiments use fixed interpolants: the teacher and pure-distilled model are trained once, then \(\xi\) is selected by oracle test risk (as a benchmark), GCV, or validation. We report validation and test RF risk because it is the quantity controlled by theory and used in RF generation bounds.
In the correctly specified Gaussian setting \(X_1\sim\mathcal{N}(0,4I_d)\), the population velocity is \(v^\star(t,x)=g(t)x\) (see Appendix 8.3) and the feature map is exact. Figure 3 shows that GCV recovers most of the oracle gain. It also illustrates the central sign rule: in over-regularized regimes \(\xi^\star<0\), so SD constrained to \(\xi\in[0,1]\) cannot improve on the teacher. At \(\lambda\approx31.6\), for example, the teacher velocity risk is about \(1.67\), oracle SD gives \(1.14\), and one-shot SD gives \(1.22\).
We next use a four-component Gaussian mixture whose nonlinear population RF velocity is computable and exactly representable by our oracle feature map. Figure 5 shows that the same correction mechanism carries over to nonlinear generation. At \(\lambda\approx7.5\), the teacher risk is \(2.077\), pure distillation worsens it to \(2.639\), oracle SD reduces it to \(0.265\), and GCV-SD gives \(0.273\) with \(\widehat\xi\approx-6.46\). Sliced \(\mathcal{W}_2^2\) improves from about \(1.02\) for the teacher to \(0.093\) for GCV-SD, while mode mass and covariance errors also improve.
We evaluate Euler sampling across a range of numbers of function evaluations (NFEs). In the same Gaussian mixture setting, oracle and GCV-SD improve generation metrics across NFEs. At NFE \(=32\), sliced \(\mathcal{W}_2^2\) is about \(1.18\) for the teacher, \(0.096\) for oracle SD, and \(0.091\) for one-shot SD; see Section 9.6 for the full curves. These results are consistent with the generation bounds: improving the velocity estimation term can improve finite-step sampling.
Figure 1 reports controlled neural RF tests on real images. The teacher output is deliberately scaled down to emulate excessive shrinkage or miscalibration of the velocity scale. Pure distillation compounds the shrinkage, while unconstrained SD selects negative mixing and recovers recognizable samples. For Fashion-MNIST, validation selects \(\widehat\xi=-2.83\) and reduces RF risk from \(0.828\) to \(0.210\). The best Feature-FD improves from \(818.5\) to \(8.8\), and conditioning accuracy at NFE \(=64\) improves from \(0.107\) to \(0.895\). At NFE \(=32\), the same experiment improves Inception-FID from \(355.1\) to \(14.8\); see Table 2 in the supplement.
We further test whether negative mixing remains useful beyond linear probing with fixed features. Starting from a trained time-conditioned U-Net RF field \(v_T\) on \(32\times32\) CIFAR-10, we use \(v=a v_T\) with \(a=0.75\) as the teacher, thereby introducing a controlled miscalibration of the velocity scale. We then train a pure-distilled U-Net and tune \(\xi\) by validation over \(v_\xi=v+\xi(v_{\mathrm{pd}}-v)\). Validation selects \(\widehat\xi=-1.225\), so restricting \(\xi\) to \([0,1]\) returns the teacher. Table 1 and Figure 6 show that optimal SD improves RF risk from \(0.2435\) to \(0.1803\), FID from \(284.3\) to \(30.1\), and KID from \(0.3215\) to \(0.0186\) at NFE \(=32\), whereas pure distillation worsens the teacher. Table 5 in Appendix 9.3 shows the same pattern across output scales: as \(a\) increases toward one and the miscalibration becomes milder, the selected negative correction moves toward zero and the gain shrinks. The appendix also reports a fine-tuning baseline with matched compute.
In this paper, we identify a one-shot mechanism by which an RF model can improve from its own predictions. For linear RF trained on fixed interpolants, the self-distilled student follows an affine path from the teacher to its pure-distilled refit, and the geometry of squared risk yields an optimal mixing coefficient in closed form. The resulting student strictly improves every nonstationary teacher along the regularization path. Unlike recursive Reflow, our procedure does not replace endpoint pairs or change interpolation covariates; it optimizes only the mixture of true and teacher velocity targets.
The main guarantee concerns integrated RF velocity risk for a fixed teacher regularization level. It does not claim that the teacher is globally optimally tuned, nor that every downstream finite-step generation metric must improve under arbitrary misspecification. Rather, it shows that a suboptimal regularization level defines a teacher-to-pure-distillation direction along which velocity risk can be reduced. Through RF Wasserstein bounds, this risk reduction also tightens natural upper bounds on generation error.
Empirically, pure distillation can worsen over-regularized teachers because it repeats the same shrinkage. Optimal SD often corrects this by extrapolating in the opposite direction, yielding a negative mixing coefficient. This need for extrapolation is the main practical distinction from conventional convex distillation. In practice, one can tune \(\xi\) by GCV or validation; when the estimated gain is small, the selected coefficient remains close to zero.
The strict improvement theorem is exact only for linear RF with ridge regularization; the neural affine family in 17 is therefore an empirical extension, and a corresponding guarantee for end-to-end nonlinear training remains open. Our controlled experiments on real images are diagnostic stress tests of teacher miscalibration: the CIFAR-10 scale sweep in Appendix 9.3 traces the correction across miscalibration levels, while the Fashion-MNIST sensitivity analysis shows that gains diminish as the teacher improves. Finally, extending the method to recursive endpoint replacement would bring it closer to Reflow practice, but requires new tools because interpolation covariates change across rounds.
We thank Hien Dang, Alessandro Rinaldo, and Sujay Sanghavi for helpful conversations. Computing support is in part provided by the Texas Advanced Computing Center (TACC).
Supplement
This supplement accompanies the paper “Optimal Self-Distillation for Rectified Flow via Linear Probing”. It contains proofs, further theoretical remarks, and additional experimental details and plots. The supplement is organized as follows.
For fixed \(\lambda\), the ridge solution is linear in the response matrix. Define \[Q_\lambda=\left(\frac{\Phi^\top\Phi}{N}+\lambda I_m\right)^{-1}.\] Using the definition of the mixed target matrix, \[\begin{align} \widehat W_{\mathrm{sd},\lambda,\xi} &=Q_\lambda\frac{\Phi^\top U^{(\xi)}}{N}\\ &=(1-\xi)Q_\lambda\frac{\Phi^\top U}{N} +\xi Q_\lambda\frac{\Phi^\top\widehat U_\lambda}{N}\\ &=(1-\xi)\widehat W_\lambda+\xi\widehat W_{\mathrm{pd},\lambda}. \end{align}\] The prediction identity follows from the linearity of \(v_W(t,x)=W^\top\phi(t,x)\) in \(W\).
Let \[r_T=u-v_{\widehat W_\lambda}(T,X_T),\qquad r_P=u-v_{\widehat W_{\mathrm{pd},\lambda}}(T,X_T).\] By Proposition 1, the self-distilled residual is \((1-\xi)r_T+\xi r_P\). Expanding the squared norm and taking conditional expectation gives \[\begin{align} R_{\mathrm{sd}}(\lambda,\xi) &=(1-\xi)^2R(\lambda)+\xi^2R_{\mathrm{pd}}(\lambda)+2\xi(1-\xi)C(\lambda)\\ &=R(\lambda)-2\xi\{R(\lambda)-C(\lambda)\}+\xi^2D(\lambda). \end{align}\] Since \(D(\lambda)>0\), completing the square yields \[R_{\mathrm{sd}}(\lambda,\xi) =D(\lambda)\left\{\xi-\frac{R(\lambda)-C(\lambda)}{D(\lambda)}\right\}^2 +R(\lambda)-\frac{\{R(\lambda)-C(\lambda)\}^2}{D(\lambda)},\] which proves both identities in 14 .
It remains to prove the derivative identity. Set \(\widehat\Sigma=\Phi^\top\Phi/N\). Since \[\widehat W_\lambda=(\widehat\Sigma+\lambda I_m)^{-1}\frac{\Phi^\top U}{N},\] its derivative is \[\partial_\lambda\widehat W_\lambda=-(\widehat\Sigma+\lambda I_m)^{-1}\widehat W_\lambda.\] The pure-distilled coefficient is \[\widehat W_{\mathrm{pd},\lambda}=(\widehat\Sigma+\lambda I_m)^{-1}\widehat\Sigma\widehat W_\lambda.\] Using \((\widehat\Sigma+\lambda I_m)^{-1}\widehat\Sigma=I_m-\lambda(\widehat\Sigma+\lambda I_m)^{-1}\) gives \[\widehat W_\lambda-\widehat W_{\mathrm{pd},\lambda}=-\lambda\partial_\lambda\widehat W_\lambda.\] Because the velocity field is linear in its coefficient matrix, it follows that \[v_{\mathrm{pd},\lambda}(t,x)-v_\lambda(t,x) =\lambda\partial_\lambda v_\lambda(t,x).\] Consequently, \[\begin{align} R(\lambda)-C(\lambda) &=\mathbb{E}\left[ \left\langle u-v_\lambda(T,X_T),v_{\mathrm{pd},\lambda}(T,X_T)-v_\lambda(T,X_T)\right\rangle \mid\mathcal{D}_N \right]\\ &=\lambda\mathbb{E}\left[ \left\langle u-v_\lambda(T,X_T),\partial_\lambda v_\lambda(T,X_T)\right\rangle \mid\mathcal{D}_N \right]. \end{align}\] Differentiating the conditional risk gives \[R'(\lambda) =-2\mathbb{E}\left[ \left\langle u-v_\lambda(T,X_T),\partial_\lambda v_\lambda(T,X_T)\right\rangle \mid\mathcal{D}_N \right],\] Combining the last two displays gives \(R(\lambda)-C(\lambda)=-(\lambda/2)R'(\lambda)\). Substituting this identity into 14 proves 15 . Since \(\lambda>0\) and \(D(\lambda)>0\), the risk reduction is strict whenever \(R'(\lambda)\neq0\), and the sign of \(\xi^\star(\lambda)\) is the opposite of the sign of \(R'(\lambda)\).
The Wasserstein convergence theorem for rectified flow in [16] gives, under the stated Lipschitz and well-posedness assumptions, \[\mathcal{W}_2^2(\rho_v,\rho_1) \le \exp(1+2L)\int_0^1 \mathbb{E}\left\lVert v^\star(t,X_t)-v(t,X_t)\right\rVert^2\,\mathrm{d}t.\] The integral is \(\mathcal{E}(v)\). The orthogonal decomposition 4 gives \[\mathcal{R}(v\mid\mathcal{D}_N)=\mathcal{E}(v)+\mathfrak S, \qquad \mathfrak S:=\mathbb{E}\left\lVert u-v^\star(T,X_T)\right\rVert^2,\] where \(\mathfrak S\) is independent of \(v\). Since the optimized affine family contains both the teacher (\(\xi=0\)) and pure distillation (\(\xi=1\)), \[\mathcal{E}(v_{\mathrm{sd}}^\star) \le \min\{\mathcal{E}(v_\lambda),\mathcal{E}(v_{\mathrm{pd},\lambda})\}.\] Moreover, Theorem 2 gives \[\mathcal{E}(v_\lambda)-\mathcal{E}(v_{\mathrm{sd}}^\star)=R(\lambda)-R_{\mathrm{sd}}^\star(\lambda) =\frac{\lambda^2}{4}\frac{(R'(\lambda))^2}{D(\lambda)}.\] Multiplying this gain by the common factor \(\exp(1+2L)\) proves Proposition 3.
All expectations below are conditional on \(\mathcal{D}_N\), with the conditioning suppressed for readability. Couple the exact RF trajectory and the Euler sampler through the same initial point: \[\label{eq:og-rf} \frac{\mathrm{d}Z_t}{\mathrm{d}t}=v^\star(t,Z_t),\qquad Z_0\sim\rho_0.\tag{26}\] On the uniform grid \(t_i=i/K\) with \(h=K^{-1}\), the Euler iterates are \[\widehat Y_{t_{i+1}} =\widehat Y_{t_i}+h v(t_i,\widehat Y_{t_i}),\qquad \widehat Y_0=Z_0,\qquad i=0,\ldots,K-1.\] For \(t\in[t_i,t_{i+1}]\), define the continuous Euler interpolation \[\label{eq:interpolation-proc} \bar Y_t :=\widehat Y_{t_i}+(t-t_i)v(t_i,\widehat Y_{t_i}), \qquad i=0,\ldots,K-1.\tag{27}\] Thus \(\bar Y_{t_i}=\widehat Y_{t_i}\), \(\bar Y_{t_{i+1}}=\widehat Y_{t_{i+1}}\), and \(\mathrm{d}\bar Y_t/\mathrm{d}t=v(t_i,\widehat Y_{t_i})\). Let \[e_t:=\bar Y_t-Z_t,\qquad \delta_t:=v(t_i,\widehat Y_{t_i})-v^\star(t,Z_t).\] For any \(\alpha>0\), Young’s inequality gives \[\label{eq:ddt-decomp} \frac{\mathrm{d}}{\mathrm{d}t}\left\lVert e_t\right\rVert_2^2 =2\left\langle e_t,\delta_t\right\rangle \le \alpha\left\lVert e_t\right\rVert_2^2+\frac{1}{\alpha}\left\lVert \delta_t\right\rVert_2^2.\tag{28}\] Applying the integrating factor \(\exp\{-\alpha(t-t_i)\}\) and integrating over \([t_i,t_{i+1}]\) yields \[\begin{align} \left\lVert \widehat Y_{t_{i+1}}-Z_{t_{i+1}}\right\rVert_2^2 &\le e^{\alpha h}\left\lVert \widehat Y_{t_i}-Z_{t_i}\right\rVert_2^2\\ &\quad+\frac{e^{\alpha h}}{\alpha} \int_{t_i}^{t_{i+1}}\left\lVert \delta_t\right\rVert_2^2\,\mathrm{d}t. \end{align}\] Define \[\Delta_i:=e^{-\alpha t_i}\mathbb{E}\left\lVert \widehat Y_{t_i}-Z_{t_i}\right\rVert_2^2.\] Because \(\widehat Y_0=Z_0\), we have \(\Delta_0=0\). Taking expectations in the preceding inequality gives \[\Delta_{i+1} \le \Delta_i+\frac{e^{-\alpha t_i}}{\alpha} \int_{t_i}^{t_{i+1}}\mathbb{E}\left\lVert \delta_t\right\rVert_2^2\,\mathrm{d}t.\] Decompose the integral using \(\left\lVert a+b+c\right\rVert^2\le3\{\left\lVert a\right\rVert^2+\left\lVert b\right\rVert^2+\left\lVert c\right\rVert^2\}\): \[\begin{align} \int_{t_i}^{t_{i+1}}\mathbb{E}\left\lVert \delta_t\right\rVert_2^2\,\mathrm{d}t \le 3\{({\rm I})+({\rm II})+({\rm III})\}, \end{align}\] where \[\begin{align} ({\rm I}) &:=\int_{t_i}^{t_{i+1}} \mathbb{E}\left\lVert v^\star(t,Z_t)-v^\star(t_i,Z_{t_i})\right\rVert_2^2\,\mathrm{d}t,\\ ({\rm II}) &:=\int_{t_i}^{t_{i+1}} \mathbb{E}\left\lVert v^\star(t_i,Z_{t_i})-v(t_i,Z_{t_i})\right\rVert_2^2\,\mathrm{d}t,\\ ({\rm III}) &:=\int_{t_i}^{t_{i+1}} \mathbb{E}\left\lVert v(t_i,Z_{t_i})-v(t_i,\widehat Y_{t_i})\right\rVert_2^2\,\mathrm{d}t. \end{align}\]
We bound these three terms in turn. For \(({\rm I})\), the fundamental theorem of calculus and Cauchy–Schwarz give, for \(t\in[t_i,t_{i+1}]\), \[\label{eq:straightness-bound} \begin{align} \mathbb{E}\left\lVert v^\star(t,Z_t)-v^\star(t_i,Z_{t_i})\right\rVert_2^2 &\le (t-t_i)\int_{t_i}^{t} \mathbb{E}\left\lVert \frac{\mathrm{d}}{\mathrm{d}\tau}v^\star(\tau,Z_\tau)\right\rVert_2^2\,\mathrm{d}\tau\\ &\le h^2\gamma_i, \end{align}\tag{29}\] where \[\gamma_i:=\frac{1}{h}\int_{t_i}^{t_{i+1}} \mathbb{E}\left\lVert \frac{\mathrm{d}}{\mathrm{d}\tau}v^\star(\tau,Z_\tau)\right\rVert_2^2\,\mathrm{d}\tau.\] Therefore, \[({\rm I})\le h^3\gamma_i\le h^3\gamma_{2,K}(\mathcal{Z}^\star).\] Since \(Z_{t_i}\) has the same law as \(X_{t_i}\), 22 gives \[({\rm II})=h\varepsilon_i^2(v).\] Finally, the \(L\)-Lipschitz property of \(v(t_i,\cdot)\) implies \[({\rm III}) \le hL^2\mathbb{E}\left\lVert Z_{t_i}-\widehat Y_{t_i}\right\rVert_2^2 =hL^2e^{\alpha t_i}\Delta_i.\] Combining these bounds gives the recursion \[\label{eq:moment-recursion} \Delta_{i+1} \le \left(1+\frac{3L^2h}{\alpha}\right)\Delta_i +\frac{3e^{-\alpha t_i}}{\alpha} \left\{h^3\gamma_{2,K}(\mathcal{Z}^\star)+h\varepsilon_i^2(v)\right\}.\tag{30}\] Let \(q:=1+3L^2h/\alpha\). Dividing 30 by \(q^{i+1}\), summing over \(i\), and using \(e^{-\alpha t_i}/q^{i+1}\le1\) gives \[\frac{\Delta_K}{q^K} \le \frac{3}{\alpha}\left\{ \frac{\gamma_{2,K}(\mathcal{Z}^\star)}{K^2} +\bar\varepsilon_K^2(v) \right\}.\] Since \(q^K\le\exp(3L^2/\alpha)\) and \[\mathcal{W}_2^2(\widehat\rho_{v,K},\rho_1) \le \mathbb{E}\left\lVert \widehat Y_{t_K}-Z_{t_K}\right\rVert_2^2 =e^\alpha\Delta_K,\] we obtain, for every \(\alpha>0\), \[\label{eq:euler-alpha-bound} \mathcal{W}_2^2(\widehat\rho_{v,K},\rho_1) \le \frac{3\exp(3L^2/\alpha+\alpha)}{\alpha} \left\{\frac{\gamma_{2,K}(\mathcal{Z}^\star)}{K^2}+\bar\varepsilon_K^2(v)\right\}.\tag{31}\] The prefactor is minimized at \[\alpha^\star=\frac{1+\sqrt{1+12L^2}}{2}.\] Writing \(s:=\sqrt{1+12L^2}\), substitution into 31 gives the sharper constant \[\widetilde{C}_L:=\frac{6e^s}{1+s} \le \frac{6e^{1+s}}{1+s}=C_L.\] Thus 24 follows with the stated constant \(C_L\).
For a fixed grid point \(t_i\), the same quadratic expansion as in Theorem 2 gives \[R_{\mathrm{sd},t_i}(\lambda,\xi) =R_{t_i}(\lambda)-2\xi A_{t_i}(\lambda)+\xi^2D_{t_i}(\lambda).\] The time-specific irreducible term \[\mathfrak S_{t_i} :=\mathbb{E}\left\lVert u-v^\star(t_i,X_{t_i})\right\rVert_2^2\] is independent of \(v\), so the same quadratic and gain hold for \(\mathcal{E}_{t_i}\). If \(D_{t_i}(\lambda)>0\), minimizing over a time-specific coefficient gives \[\xi_{t_i}^\star(\lambda) =\frac{A_{t_i}(\lambda)}{D_{t_i}(\lambda)},\qquad \mathcal{E}_{t_i}(v_\lambda) -\mathcal{E}_{t_i}(v_{\mathrm{sd},\lambda,\xi_{t_i}^\star}) =\frac{A_{t_i}(\lambda)^2}{D_{t_i}(\lambda)}.\] Whenever \(R_{t_i}(\lambda)\) is differentiable, the fixed-time version of the derivative identity in Appendix 7 gives \[A_{t_i}(\lambda)=-\frac{\lambda}{2}\partial_\lambda R_{t_i}(\lambda).\]
For a shared scalar \(\xi\) across the grid, the average error is \[\bar\varepsilon_K^2(v_{\mathrm{sd},\lambda,\xi}) =\frac{1}{K}\sum_{i=0}^{K-1} \left\{R_{t_i}(\lambda)-2\xi A_{t_i}(\lambda) +\xi^2D_{t_i}(\lambda)-\mathfrak S_{t_i}\right\}.\] If \(\sum_iD_{t_i}(\lambda)>0\), completing the square gives \[\xi_{\mathrm U,K}^\star(\lambda) =\frac{\sum_iA_{t_i}(\lambda)}{\sum_iD_{t_i}(\lambda)},\] and \[\bar\varepsilon_K^2(v_\lambda)-\bar\varepsilon_K^2(v_{\mathrm{sd},K}^\star) =\frac{\{\sum_iA_{t_i}(\lambda)\}^2}{K\sum_iD_{t_i}(\lambda)}.\] This proves 23 . Because the shared affine family contains \(\xi=0\) and \(\xi=1\), its optimum is no worse than either the teacher or pure distillation. Combining this fact with the Wasserstein bound above and a common Lipschitz constant proves the remaining claims of Theorem 4.
This subsection makes explicit, in a Gaussian linear model, how velocity error translates into terminal distribution error. Suppose \[X_0\sim\mathcal{N}(0,I_d),\qquad X_1\sim\mathcal{N}(0,\Sigma),\qquad X_0\perp X_1,\] where \(\Sigma\succ0\). For \(u=X_1-X_0\) and \(X_t=(1-t)X_0+tX_1\), \[C_t\mathrel{\vcenter{:}}=\operatorname{Cov}(X_t)=(1-t)^2I_d+t^2\Sigma, \qquad \operatorname{Cov}(u,X_t)=t\Sigma-(1-t)I_d.\] The Gaussian conditioning formula therefore gives \[\label{eq:gaussian-velocity} v^\star(t,x)=A_t^\star x, \qquad A_t^\star=\{t\Sigma-(1-t)I_d\}C_t^{-1}.\tag{32}\]
Now consider a learned linear velocity \(v(t,x)=A_tx\) whose matrices are jointly diagonalizable with \(\Sigma\): \[\Sigma=V\mathrm{diag}(s_1,\ldots,s_d)V^\top, \qquad A_t=V\mathrm{diag}(a_1(t),\ldots,a_d(t))V^\top.\] In this basis, write \(a_j^\star(t)\) for the corresponding coefficient of \(A_t^\star\), and define \[\Delta_j(t)=a_j(t)-a_j^\star(t), \qquad M_j=\int_0^1\Delta_j(t)\,\mathrm{d}t.\] Under the learned ODE, the variance \(p_j(t)\) in coordinate \(j\) satisfies \[\frac{\mathrm{d}p_j(t)}{\mathrm{d}t}=2a_j(t)p_j(t), \qquad p_j(0)=1.\] The target flow has terminal variance \(s_j\), so \[p_j(1) =\exp\left\{2\int_0^1a_j(t)\,\mathrm{d}t\right\} =s_je^{2M_j}.\] Thus the learned terminal covariance is \[P_1=V\mathrm{diag}\{s_1e^{2M_1},\ldots,s_de^{2M_d}\}V^\top.\] Because \(P_1\) and \(\Sigma\) share the same eigenvectors, the Gaussian Wasserstein and KL formulas reduce to \[\begin{align} \mathcal{W}_2^2(\mathcal{N}(0,P_1),\mathcal{N}(0,\Sigma)) &=\sum_{j=1}^d s_j(e^{M_j}-1)^2,\tag{33}\\ {\sf KL}(\mathcal{N}(0,P_1)\,\|\,\mathcal{N}(0,\Sigma)) &=\frac{1}{2}\sum_{j=1}^d\{e^{2M_j}-1-2M_j\}. \tag{34} \end{align}\]
Finally, the reducible component of RF regression risk is the integrated squared velocity error along the true marginals: \[\int_0^1\mathbb{E}\!\left[\|v(t,X_t)-v^\star(t,X_t)\|^2\right]\mathrm{d}t =\sum_{j=1}^d\mathcal{E}_j, \qquad \mathcal{E}_j\mathrel{\vcenter{:}}=\int_0^1c_j(t)\Delta_j(t)^2\,\mathrm{d}t, \label{eq:gauss-risk}\tag{35}\] where \(c_j(t)=(1-t)^2+t^2s_j\). Weighted Cauchy–Schwarz yields \[M_j^2 \leq\left(\int_0^1\frac{\mathrm{d}t}{c_j(t)}\right)\mathcal{E}_j =\frac{\pi}{2\sqrt{s_j}}\,\mathcal{E}_j. \label{eq:gauss-cs}\tag{36}\] Equations 33 –36 therefore connect the exact terminal discrepancies to the same coordinatewise errors that comprise the RF excess risk. These identities underlie the Gaussian experiments in Section 5.
This appendix supplements Section 5 with detailed metrics and sensitivity checks for the image experiments, followed by additional studies of anisotropy, discretization, and regularization in Gaussian and Gaussian mixture models.
Figure 7: Results by dataset underlying Figure 1. The panels show additional generated samples from the experiments with controlled shrinkage.. a — Handwritten digits., b — Fashion-MNIST.
For the controlled Fashion-MNIST experiment, Feature-FD is the Fréchet distance between generated and held-out real samples in the feature space defined by the penultimate layer of a Fashion-MNIST classifier trained on real training images. Inception-FID is computed after converting grayscale images to RGB and resizing them to \(299\times299\). Conditioning accuracy is the fraction of generated samples classified as the requested class, and confidence is the average classifier probability assigned to that class. Table 2 uses \(10{,}000\) generated samples and \(10{,}000\) held-out real samples at NFE \(=32\).
The Fashion-MNIST result in Figure 1 deliberately uses a teacher with strong output shrinkage to isolate the correction from negative mixing. Table 3 compares this stress test with an unscaled, naturally trained teacher. The latter has much lower initial risk and generation error; accordingly, validation selects a correction closer to zero and yields only a modest gain, consistent with its interpretation as a local correction.
| Method | Feature-FD \(\downarrow\) | Inception-FID \(\downarrow\) | Acc. \(\uparrow\) | Conf. \(\uparrow\) |
|---|---|---|---|---|
| Teacher | 854.18 | 355.08 | 0.109 | 0.111 |
| Pure distillation | 983.10 | 388.11 | 0.102 | 0.103 |
| SD clipped \([0,1]\) | 857.26 | 355.10 | 0.113 | 0.114 |
| Optimal SD | 9.65 | 14.84 | 0.898 | 0.863 |
| Teacher regime | Scale | \(\widehat\xi\) | \(R_T\) | \(R_{\rm SD}\) | Feature-FD\(_T \to\) Feature-FD\(_{\rm SD}\) |
|---|---|---|---|---|---|
| Controlled shrinkage | \(0.35\) | \(-2.83\) | \(0.828\) | \(0.210\) | \(818.5 \to 8.8\) |
| Unscaled teacher | \(1.00\) | \(0.239\) | \(0.205\) | \(0.205\) | \(9.88 \to 9.38\) |
Using a fully neural RF model on CIFAR-10 at \(32\times32\) resolution, we test the correction mechanism beyond linear probing with fixed features. We first train a time-conditioned U-Net velocity field with the standard RF objective, then use its output scaled by \(a=0.75\) as a teacher with controlled miscalibration of the velocity scale. We train a pure-distilled U-Net on the scaled teacher velocities and select \(\xi\) by validation over \(v_\xi=v+\xi(v_{\mathrm{pd}}-v)\). Metrics are computed from \(10{,}000\) generated samples using Euler sampling at NFE \(=32\). FID and KID use Inception features with the preprocessing implemented by torch-fidelity. Table 4 also includes a fine-tuning baseline with matched compute, trained for the same number of adaptation steps as pure distillation.
| Method | RF risk \(\downarrow\) | FID \(\downarrow\) | KID \(\downarrow\) |
|---|---|---|---|
| Teacher | 0.2435 | 284.26 | 0.3215 |
| Pure distillation | 0.3885 | 324.77 | 0.3569 |
| SD clipped \([0,1]\) | 0.2435 | 283.56 | 0.3196 |
| Optimal SD | 0.1803 | 30.08 | 0.0186 |
| Compute-matched fine-tune | 0.1749 | 34.94 | 0.0268 |
We assess sensitivity to teacher output scaling in Table 5. Across the evaluated scales, the correction selected by validation remains negative, pure distillation worsens the teacher, and optimal SD substantially improves FID.
| Scale \(a\) | \(\widehat\xi\) | RF risk \(T\to\) SD | FID \(T\to\) SD | FID fine-tune |
|---|---|---|---|---|
| \(0.55\) | \(-1.78\) | \(0.396\to0.179\) | \(324.4\to35.0\) | \(47.0\) |
| \(0.65\) | \(-1.49\) | \(0.310\to0.181\) | \(321.1\to35.3\) | \(46.9\) |
| \(0.75\) | \(-1.23\) | \(0.245\to0.182\) | \(291.9\to41.6\) | \(46.9\) |
| \(0.85\) | \(-0.88\) | \(0.201\to0.182\) | \(212.4\to56.5\) | \(47.0\) |
We use an anisotropic Gaussian target with covariance eigenvalues logarithmically spaced between \(0.15\) and \(6\) and a model with a Legendre basis in time, \(v_W(t,x)=\sum_{\ell=0}^5b_\ell(t)W_\ell x\). Figure 12 shows that the sign rule and strict improvement persist beyond the scalar Gaussian model with exact features.
We use a four-component Gaussian mixture (MoG) in two dimensions. The population velocity \(\mathbb{E}[u\mid X_t=x]\) is nonlinear but computable from posterior responsibilities; we use an oracle feature map in which this velocity is exactly representable. At \(\lambda\approx7.5\), the teacher velocity risk is \(2.08\), pure distillation worsens it to \(2.64\), while oracle SD and one-shot SD reduce the risk to \(0.265\) and \(0.273\), respectively. Figure 13 shows that the same correction improves sliced Wasserstein distance, maximum mean discrepancy (MMD), mode mass error, and covariance error. Figure 5 visualizes the effect: the teacher produces a blurred mixture, while optimal SD with negative mixing recovers the four modes.
Because RF generation requires numerical ODE integration, we evaluate Euler sampling over a range of numbers of function evaluations (NFEs). This experiment examines whether the improvement in velocity risk translates into improved finite-step generation under the Euler discretization addressed by 24 . In the MoG experiment with oracle features, SD improves generation metrics across NFEs. At NFE \(=32\), sliced \(\mathcal{W}_2^2\) is approximately \(1.18\) for the teacher, \(1.52\) for pure distillation, \(0.096\) for oracle SD, and \(0.091\) for one-shot SD; mode mass error drops from about \(0.087\) for the teacher to \(0.025\) for one-shot SD. The corresponding curves appear in Figure 14.
Figure 15 varies the aspect ratio \(d/n\) and regularization level in Gaussian RF with exact features. The largest relative gains occur away from the optimal ridge region for the teacher, and the sign of \(\xi^\star\) changes across the boundary between under- and over-regularization.
Theorem 2 optimizes a single mixing coefficient shared across time. Allowing a measurable function \(t\mapsto\xi_t\) instead gives the time-dependent affine family \[v_{\mathrm{sd},\lambda,\xi(\cdot)}(t,x) =v_\lambda(t,x)+\xi_t\{v_{\mathrm{pd},\lambda}(t,x)-v_\lambda(t,x)\}. \label{eq:td-affine}\tag{37}\] Let \(R_t(\lambda)\), \(C_t(\lambda)\), and \(D_t(\lambda)\) denote the analogues of 11 –12 at a fixed time \(t\), and set \(A_t(\lambda)=R_t(\lambda)-C_t(\lambda)\). Then the time-slice risk is \[R_{\mathrm{sd},t}(\lambda,\xi_t) =R_t(\lambda)-2\xi_t A_t(\lambda)+\xi_t^2D_t(\lambda).\] Whenever \(D_t(\lambda)>0\), the pointwise optimum is \(\xi_t^\star(\lambda)=A_t(\lambda)/D_t(\lambda)\). If \(\int_0^1D_t(\lambda)\,\mathrm{d}t>0\), the best coefficient shared across time is \[\xi_{\mathrm U}^\star(\lambda)=\frac{\int_0^1A_t(\lambda)\,\mathrm{d}t}{\int_0^1D_t(\lambda)\,\mathrm{d}t}. \label{eq:td-xi}\tag{38}\] Assume now that \(D_t(\lambda)>0\) for almost every \(t\), and define \(w_\lambda(t)=D_t(\lambda)/\int_0^1D_s(\lambda)\,\mathrm{d}s\). The additional gain in integrated risk from pointwise rather than uniform mixing is \[\label{eq:td-extra-gain} \int_0^1\frac{A_t(\lambda)^2}{D_t(\lambda)}\,\mathrm{d}t -\frac{\left(\int_0^1A_t(\lambda)\,\mathrm{d}t\right)^2}{\int_0^1D_t(\lambda)\,\mathrm{d}t} =\left(\int_0^1D_t(\lambda)\,\mathrm{d}t\right) \mathrm{Var}_{w_\lambda}\{\xi_t^\star(\lambda)\}\ge0.\tag{39}\] Thus pointwise mixing performs at least as well as the best uniform mixture, with strict improvement precisely when \(\xi_t^\star(\lambda)\) is not \(w_\lambda\)-almost surely constant. The gridwise analogue in Appendix 8.2.2 yields the reduction in velocity error used in the Euler generation bound 24 .
We evaluate this time-dependent extension on four targets in two dimensions—a Gaussian mixture with four components (MoG-4), checkerboard, moons, and spiral—using RF with fixed random features and fixed interpolants. The affine identity in Proposition 1 applies directly to a coefficient shared across all training samples. With time-varying coefficients, however, refitting on targets mixed according to each sample’s time generally differs from applying the affine construction 37 directly to the fields. We therefore distinguish TD-affine, which evaluates 37 , from TD-refit, which retrains on time-dependent mixed targets.
Figure 16 compares their integrated RF risks with those of the teacher and the best uniform mixture across ridge penalties. TD-affine provides modest, consistent improvements over the best uniform mixture, whereas TD-refit is not covered by 39 and can perform worse. Figure 17 localizes the gains across time, and Figure 18 reports the corresponding finite-step generation metrics. The downstream improvements depend on the setting and metric, consistent with the scope of our guarantee on velocity risk.