Training Directly with FID as a Loss: From Gradient Computation to Streaming Training

直接以FID为Loss：从梯度计算到流式训练 苏剑林 · 2026-05-08 · Machine Learning matrix loss function generative model gradient

Anyone who follows visual generative models knows that FID has long been one of the key evaluation metrics — the smaller it is, the more realistic the generated results tend to be. A natural question then arises: why not just train the generative model directly using FID as the loss function? Is it because FID isn't differentiable? Not at all — FID is in fact differentiable, and there's no theoretical obstacle to using it as a loss. The real difficulty lies in the practical computation.

Recently, the paper Representation Fréchet Loss for Visual Generation made some attempts to overcome these difficulties, successfully applying FID to fine-tune generative models and noticeably improving the quality of one-step generation. This post gives a brief overview of the underlying mathematics and the implementation tricks involved.

Generation Metrics

FID stands for "Fréchet Inception Distance," and we can understand it in two parts: "Fréchet Distance (FD)" and "Inception (I)." more

Suppose we have two distributions $p$ and $q$, representing the real samples and the generated samples respectively. We pass their samples $\boldsymbol{x}$ through some pretrained encoder $\phi$ to get feature vectors $\boldsymbol{z}=\phi(\boldsymbol{x})\in\mathbb{R}^d$, and estimate the corresponding mean vectors $\boldsymbol{\mu}_p,\boldsymbol{\mu}_q$ and covariance matrices $\boldsymbol{\Sigma}_p,\boldsymbol{\Sigma}_q$. If we then assume that the encoded results follow a multivariate normal distribution, we can measure the discrepancy between the two distributions using a divergence function for normal distributions. The Fréchet Distance (FD) uses the W-distance for this purpose:

\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\begin{aligned} \mathcal{F}\triangleq\mathcal{W}_2^2[p,q]=&\,\Vert \boldsymbol{\mu}_p - \boldsymbol{\mu}_q\Vert^2 + \tr(\boldsymbol{\Sigma}_p + \boldsymbol{\Sigma}_q - 2(\boldsymbol{\Sigma}_p\boldsymbol{\Sigma}_q)^{1/2})\\[4pt] =&\,\Vert \boldsymbol{\mu}_p - \boldsymbol{\mu}_q\Vert^2 + \tr(\boldsymbol{\Sigma}_p + \boldsymbol{\Sigma}_q - 2(\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2})^{1/2}) \end{aligned}\label{eq:w-p-q}\end{equation}

Substituting the mean vectors and covariance matrices of the respective encoded results into the above formula gives what is called the "Fréchet Distance (FD)." For the derivation of the formula above, interested readers may refer to The KL Divergence, Bhattacharyya Distance, and W-Distance Between Two Multivariate Normal Distributions.

If we choose the encoder $\phi$ to be InceptionV3 (the "I" above), the resulting metric is called the "Fréchet Inception Distance," i.e. FID. This evaluation metric was first proposed in the 2017 paper GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium, which in some sense is already an artifact of the "ancient era."

Of course, whether for training or evaluation, there's no longer any strict requirement to use InceptionV3 — one can use other, more advanced feature models, such as SigLIP, or even compute the Fréchet Distance separately with multiple different encoders and sum the results, and so on. We'll refer to all such approaches collectively as "FD Loss."

Related Work

Although FID looks complicated, it doesn't actually involve any non-differentiable operations, so using it as a loss is a very natural idea. Attempts along these lines were already made several years ago, such as Image Generation Via Minimizing Fréchet Distance in Discriminator Feature Space and Backpropagating through Fréchet Inception Distance.

However, these earlier attempts never produced particularly impressive results, and the fundamental reason is batch size. An ordinary loss function computes a loss per individual sample and then averages over the whole batch, whereas FID first computes the mean and covariance over the whole batch and then plugs them into equation $\eqref{eq:w-p-q}$ for a nonlinear calculation. This means that the FID estimated from a small batch is biased, and this bias cannot be eliminated simply by training longer — it can only be mitigated by increasing the batch size, which makes training prohibitively expensive.

"Requires nonlinear computation across samples," "requires a large batch size" — does this sound familiar? Indeed, contrastive learning in vision typically shares these same two characteristics: because of the nonlinear computation across samples, we can't simply use gradient accumulation to simulate a larger batch size either. But this isn't entirely unsolvable, as discussed in Can Contrastive Learning Use Gradient Accumulation?. Later we'll see that the approach FID uses to address the batch-size problem is actually quite similar.

Additionally, from the angle of "using a pretrained model to extract features for constructing a loss function," there's a related line of work called Perceptual Loss. However, that is used as a reconstruction loss for individual samples, typically in training VAEs and the like, and doesn't involve any cross-sample statistical computation — so it doesn't run into the same computational difficulties.

Gradient Computation

Let's now work through the derivation step by step to see exactly what difficulties arise when using FD as a loss. First we need to address the computation of the gradient. $p$ represents the real distribution, whose $\boldsymbol{\mu}_p,\boldsymbol{\Sigma}_p$ is fixed, so we only need to take the gradient with respect to $\boldsymbol{\mu}_q,\boldsymbol{\Sigma}_q$. The gradient of $\boldsymbol{\mu}_q$ is fairly simple:

\begin{equation}\nabla_{\boldsymbol{\mu}_q}\mathcal{F} = \nabla_{\boldsymbol{\mu}_q}\Vert \boldsymbol{\mu}_p - \boldsymbol{\mu}_q\Vert^2 = 2(\boldsymbol{\mu}_q - \boldsymbol{\mu}_p) \end{equation}

while the gradient of $\boldsymbol{\Sigma}_q$ is

\begin{equation}\nabla_{\boldsymbol{\Sigma}_q}\mathcal{F} = \nabla_{\boldsymbol{\Sigma}_q}\tr(\boldsymbol{\Sigma}_p + \boldsymbol{\Sigma}_q - 2(\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2})^{1/2}) = \boldsymbol{I} - 2\nabla_{\boldsymbol{\Sigma}_q}\tr((\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2})^{1/2})\end{equation}

Here we use the second line of equation $\eqref{eq:w-p-q}$, which looks more complicated but has one advantage: the matrix $\boldsymbol{S} = \boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2}$ is positive definite and symmetric, a fact we can exploit to simplify the computation. Let the singular value decomposition (equivalently, the eigendecomposition) of $\boldsymbol{S}$ be $\boldsymbol{U}\boldsymbol{\Lambda}\boldsymbol{U}^{\top}$, so that $\boldsymbol{S}^{1/2}=\boldsymbol{U}\boldsymbol{\Lambda}^{1/2}\boldsymbol{U}^{\top}$. Then

\begin{align}\tr(\boldsymbol{S}^{1/2})=&\,\tr(\boldsymbol{\Lambda}^{1/2})=\sqrt{\lambda_1}+\sqrt{\lambda_2}+\cdots+\sqrt{\lambda_d} \\[4pt] \nabla_{\boldsymbol{S}}\tr(\boldsymbol{S}^{1/2}) =&\, \frac{1}{2}\sum_{i=1}^d\frac{\nabla_{\boldsymbol{S}} \lambda_i}{\sqrt{\lambda_i}} = \frac{1}{2}\sum_{i=1}^d\frac{\boldsymbol{u}_i\boldsymbol{u}_i^{\top}}{\sqrt{\lambda_i}} = \frac{1}{2}\boldsymbol{U}\boldsymbol{\Lambda}^{-1/2}\boldsymbol{U}^{\top} = \frac{1}{2}\boldsymbol{S}^{-1/2}\end{align}

For the derivative of the eigenvalues, see Derivatives of the SVD. The final result resembles the derivative of $\sqrt{x}$ being $\frac{1}{2\sqrt{x}}$, which seems intuitive, but it's not trivial — if $\boldsymbol{S}$ weren't positive definite and symmetric, this generally would not hold. Finally, by the chain rule, we have

\begin{equation}\nabla_{\boldsymbol{\Sigma}_q} \tr(\boldsymbol{S}^{1/2}) = \boldsymbol{\Sigma}_p^{1/2}[\nabla_{\boldsymbol{S}}\tr(\boldsymbol{S}^{1/2})] \boldsymbol{\Sigma}_p^{1/2} = \frac{1}{2}\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{S}^{-1/2}\boldsymbol{\Sigma}_p^{1/2} \end{equation}

Putting it all together, we get

\begin{equation}\nabla_{\boldsymbol{\Sigma}_q}\mathcal{W}_2^2[p,q] = \boldsymbol{I} - \boldsymbol{\Sigma}_p^{1/2}(\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2})^{-1/2}\boldsymbol{\Sigma}_p^{1/2}\label{eq:Sigma-grad}\end{equation}

This form looks complicated, but $\boldsymbol{\Sigma}_p^{1/2}$ can be precomputed in advance, and all we need is to take the square root and inverse square root of the positive definite symmetric matrix $\boldsymbol{S} = \boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2}$ in order to compute both FID and its gradient. This can be done with the function eigh, or using the Newton–Schulz iteration schemes described in Efficient Computation of Matrix Square Roots and Inverse Square Roots and Efficient Computation of Matrix r-th Roots and Inverse r-th Roots.

Extremely Large Batches

Let's introduce the notation

\begin{equation}\begin{gathered} \boldsymbol{\mu}_p = \mathbb{E}[\boldsymbol{z}_p], \qquad \boldsymbol{V}_p = \mathbb{E}[\boldsymbol{z}_p \boldsymbol{z}_p^{\top}], \qquad \boldsymbol{z}_p = \phi(\boldsymbol{x}_p),\qquad \boldsymbol{x}_p\sim p \\[4pt] \boldsymbol{\mu}_q = \mathbb{E}[\boldsymbol{z}_q], \qquad \boldsymbol{V}_q = \mathbb{E}[\boldsymbol{z}_q \boldsymbol{z}_q^{\top}], \qquad \boldsymbol{z}_q = \phi(\boldsymbol{x}_q),\qquad \boldsymbol{x}_q\sim q \end{gathered}\end{equation}

so that

\begin{equation}\boldsymbol{\Sigma}_p = \boldsymbol{V}_p - \boldsymbol{\mu}_p \boldsymbol{\mu}_p^{\top},\qquad\boldsymbol{\Sigma}_q = \boldsymbol{V}_q - \boldsymbol{\mu}_q \boldsymbol{\mu}_q^{\top}\end{equation}

Note that $\boldsymbol{z}=\phi(\boldsymbol{x})$ typically has thousands of dimensions (2048 for InceptionV3), so for the estimate to be accurate, we usually need tens of thousands of samples. The real distribution is fixed, so its $\boldsymbol{\mu}_p,\boldsymbol{\Sigma}_p$ can be computed ahead of time — no problem there. But the generated distribution changes in real time as training proceeds, and if we had to compute it from tens of thousands of samples at every step, that would mean a batch size in the tens of thousands, which is prohibitively expensive in most settings.

We can also see the necessity of a large batch from the gradient formula $\eqref{eq:Sigma-grad}$: if the batch size is too small, the estimated $\boldsymbol{V}_q$ will not have full rank, and consequently $\boldsymbol{\Sigma}_q$ won't have full rank either, at which point inverting $(\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2})^{-1/2}$ becomes ill-defined (we'd run into $0^{-1/2}$). So using FD as a loss imposes a requirement on the training batch size, and this is arguably the most fundamental practical difficulty.

Given limited compute, we need to find a way to simulate the effect of a large batch size using a small one — a need very similar to that of "contrastive learning + gradient accumulation."

Equivalent Loss

Suppose that a batch size of $B$ is required for $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ to be sufficiently accurate, but we can only run a small batch size of $b$ each time. Then we need to run $k=B/b$ separate iterations to simulate the effect of a large batch, producing results $\tilde{\boldsymbol{\mu}}_q^{(1)},\tilde{\boldsymbol{V}}_q^{(1)}$, $\tilde{\boldsymbol{\mu}}_q^{(2)},\tilde{\boldsymbol{V}}_q^{(2)}$, ..., $\tilde{\boldsymbol{\mu}}_q^{(k)},\tilde{\boldsymbol{V}}_q^{(k)}$ respectively, satisfying the relation

\begin{equation}\boldsymbol{\mu}_q = \frac{1}{k}\sum_{i=1}^k\tilde{\boldsymbol{\mu}}_q^{(i)},\qquad \boldsymbol{V}_q = \frac{1}{k}\sum_{i=1}^k\tilde{\boldsymbol{V}}_q^{(i)}\end{equation}

We want to find an ideal equivalent loss such that the total gradient is equal to the sum of the gradients from each small batch — this would give us an unbiased estimate. To this end, we take the differential of both sides of equation $\eqref{eq:w-p-q}$:

\begin{equation}\begin{aligned} d\mathcal{F}(\boldsymbol{\mu}_q,\boldsymbol{V}_q) =&\, \langle\nabla_{\boldsymbol{\mu}_q}\mathcal{F}, d\boldsymbol{\mu}_q \rangle + \langle\nabla_{\boldsymbol{V}_q}\mathcal{F}, d\boldsymbol{V}_q \rangle_F \\ =&\, \sum_{i=1}^k \left[\langle\nabla_{\boldsymbol{\mu}_q}\mathcal{F}, d\tilde{\boldsymbol{\mu}}_q^{(i)}/k \rangle + \langle\nabla_{\boldsymbol{V}_q}\mathcal{F}, d\tilde{\boldsymbol{V}}_q^{(i)}/k \rangle_F\right] \\ =&\, d\sum_{i=1}^k \mathcal{F}(\color{skyblue}{[}\boldsymbol{\mu}_q - \tilde{\boldsymbol{\mu}}_q^{(i)}/k\color{skyblue}{]_{sg}} + \tilde{\boldsymbol{\mu}}_q^{(i)}/k,\color{skyblue}{[}\boldsymbol{V}_q - \tilde{\boldsymbol{V}}_q^{(i)}/k\color{skyblue}{]_{sg}} + \tilde{\boldsymbol{V}}_q^{(i)}/k) \\ \end{aligned}\end{equation}

This equality means that we can compute $\tilde{\boldsymbol{\mu}}_q^{(i)},\tilde{\boldsymbol{V}}_q^{(i)}$ successively from small-batch forward passes, average them to get a sufficiently accurate $\boldsymbol{\mu}_q,\boldsymbol{V}_q$, and then, batch by batch, take the gradient of the loss as usual

\begin{equation}\mathcal{F}_i = \mathcal{F}(\color{skyblue}{[}\boldsymbol{\mu}_q - \tilde{\boldsymbol{\mu}}_q^{(i)}/k\color{skyblue}{]_{sg}} + \tilde{\boldsymbol{\mu}}_q^{(i)}/k,\color{skyblue}{[}\boldsymbol{V}_q - \tilde{\boldsymbol{V}}_q^{(i)}/k\color{skyblue}{]_{sg}} + \tilde{\boldsymbol{V}}_q^{(i)}/k)\label{eq:Fi}\end{equation}

Finally, summing up all these gradients gives a gradient equivalent to that of a batch of size $B$, where $\color{skyblue}{[\cdot]_{sg}}$ denotes the stop-gradient operator. Of course, we could also choose not to accumulate gradients but instead perform an update at every single step, correspondingly reducing the learning rate a bit — the effect would be similar.

Making Up for It with History

Although the scheme above is theoretically feasible, since it requires $k$ forward passes before we can compute an accurate $\boldsymbol{\mu}_q,\boldsymbol{V}_q$, and only then go back and compute the gradient at each step, the whole pipeline doesn't feel very "streamlined." The bottleneck here is that we must know the global $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ before we can find an unbiased local gradient.

A natural idea is to ask whether we can approximate $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ somehow. Given that the learning rate is small, the parameters change slowly, so $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ should also change slowly. When a new batch of data is incorporated, the new $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ should just be a slight adjustment on top of the old one. We can approximate this using an exponential moving average (EMA):

\begin{equation}\boldsymbol{\mu}_q^{(t)} = \beta \boldsymbol{\mu}_q^{(t-1)} + (1-\beta) \tilde{\boldsymbol{\mu}}_q^{(t)},\qquad \boldsymbol{V}_q^{(t)} = \beta \boldsymbol{V}_q^{(t-1)} + (1-\beta) \tilde{\boldsymbol{V}}_q^{(t)}\end{equation}

This roughly maintains an averaging window of size $\mathcal{O}(1/(1-\beta))$, which is approximately equivalent to expanding the effective statistical batch size of $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ by a factor of $\mathcal{O}(1/(1-\beta))$. With this, at each step we can compute the gradient of the following loss to perform the update:

\begin{equation}\mathcal{F}_t = \mathcal{F}(\underbrace{\beta \color{skyblue}{[}\boldsymbol{\mu}_q^{(t-1)}\color{skyblue}{]_{sg}} + (1-\beta) \tilde{\boldsymbol{\mu}}_q^{(t)}}_{\boldsymbol{\mu}_q^{(t)}},\underbrace{\beta \color{skyblue}{[}\boldsymbol{V}_q^{(t-1)}\color{skyblue}{]_{sg}} + (1-\beta) \tilde{\boldsymbol{V}}_q^{(t)}}_{\boldsymbol{V}_q^{(t)}})\end{equation}

The extra cost is having to cache $\boldsymbol{\mu}_q,\boldsymbol{V}_q$, which is quite cheap. This trick of "not enough batch size, so make up for it with history" is, in fact, another expression of the "streaming" idea behind streaming power iteration. The paper also discusses a queue-based approach, which maintains a queue of $k$ historical batches, incorporates them into the current batch according to equation $\eqref{eq:Fi}$ to compute the gradient, and discards the oldest batch. This approach is actually somewhat simpler, but it takes up far more memory than EMA, and in practice performs worse than EMA.

Experimental Highlights

The paper's experiments focus mainly on post-training of generative models, aiming to use FD Loss training to improve existing one-step generation models, or to fine-tune existing multi-step generation models into one-step ones. When mixing multiple different encoders to compute the FD Loss, the paper uses a loss normalization trick to balance losses of different magnitudes:

\begin{equation}\mathcal{L} = \sum_i \frac{\mathcal{F}[\phi_i]}{\color{skyblue}{[}\mathcal{F}[\phi_i]\color{skyblue}{]_{sg}} + \epsilon}\end{equation}

We discussed this same trick in Musings on Multi-Task Learning (I): In the Name of Loss.

The paper's core achievement is pushing one-step generation quality (measured by FID) to an entirely new level, surpassing all other one-step and multi-step generation models — a result that appears to have hit the ceiling. Some of the figures are shown below:

FD loss improves one-step generation quality
Comprehensive comparison of FD loss

Summary

This post has mainly analyzed, from a theoretical perspective, the difficulties involved in using FID as a loss function for generative models, and how the derivation process naturally leads to techniques for overcoming these difficulties.