% Source: https://kexue.fm/archives/11738
\documentclass[11pt]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{lmodern}
\usepackage{geometry}
\usepackage{hyperref}
\usepackage{graphicx}
\usepackage{amsmath,amssymb}
\usepackage{booktabs,longtable,array}
\usepackage{xcolor}
\usepackage{float}

\title{Direct FID as Loss: From Gradient Computation to Streaming Training}
\author{Su Jianlin}
\date{2026-05-08}

\begin{document}

\maketitle

\hypersetup{colorlinks=true,linkcolor=blue,urlcolor=blue}

\begin{abstract}
Readers familiar with visual generation models know that FID has long been one of the key evaluation metrics, with lower values typically indicating more realistic generation. A natural question arises: why not directly use FID as the loss function for training generative models? Is it because FID is non-differentiable? Not at all---FID is in fact differentiable, and in theory presents no issue as a loss function. However, practical difficulties in computation arise. Recently, the paper \emph{Representation Fr\'{e}chet Loss for Visual Generation} has made attempts to overcome these challenges, successfully applying FID to fine-tune generative models and significantly improving single-step generation performance. This article briefly explores the underlying mathematics and implementation techniques.
\end{abstract}

\newcommand{\tr}{\operatorname{tr}}

\section{Generation Metrics}

FID, short for ``Fr\'{e}chet Inception Distance'', can be understood in two parts: ``Fr\'{e}chet Distance (FD)'' and ``Inception (I)''. Suppose we have two distributions $p$ and $q$, representing real and generated samples respectively. We encode input samples $\boldsymbol{x}$ via a pretrained encoder $\phi$ into feature vectors $\boldsymbol{z} = \phi(\boldsymbol{x}) \in \mathbb{R}^d$, and estimate their mean vectors $\boldsymbol{\mu}_p, \boldsymbol{\mu}_q$ and covariance matrices $\boldsymbol{\Sigma}_p, \boldsymbol{\Sigma}_q$. Assuming the encoded features follow multivariate normal distributions, we can measure their discrepancy using the Wasserstein-2 distance:
\begin{equation}
\begin{aligned}
\mathcal{F} \triangleq \mathcal{W}_2^2[p,q] 
&= \|\boldsymbol{\mu}_p - \boldsymbol{\mu}_q\|^2 + \tr(\boldsymbol{\Sigma}_p + \boldsymbol{\Sigma}_q - 2(\boldsymbol{\Sigma}_p\boldsymbol{\Sigma}_q)^{1/2}) \\
&= \|\boldsymbol{\mu}_p - \boldsymbol{\mu}_q\|^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}
where we define $\tr$ as the trace of a matrix. Substituting the estimated mean and covariance from the encoded features into this equation gives the ``Fr\'{e}chet Distance (FD)''. For details on the derivation of this formula, interested readers may refer to \href{https://kexue.fm/archives/8512}{\emph{KL Divergence, Bhattacharyya Distance, and Wasserstein Distance between Two Multivariate Normal Distributions}}.

When the encoder $\phi$ is taken to be \href{https://papers.cool/arxiv/1512.00567}{InceptionV3} (I), the resulting metric is called ``Fr\'{e}chet Inception Distance'', or FID. This evaluation metric was first introduced in the 2017 paper \href{https://papers.cool/arxiv/1706.08500}{\emph{GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium}}, which in some sense belongs to the ``ancient era'' of generative modeling.

Nowadays, however, in both training and evaluation, InceptionV3 is not mandatory; more advanced feature extractors such as \href{https://papers.cool/arxiv/2303.15343}{SigLIP} can be used instead. Alternatively, one may compute the Fr\'{e}chet Distance using multiple distinct encoders and sum the results. We collectively refer to such generalizations as ``FD Loss''.

\section{Related Work}

Although FID appears complex, it involves no non-differentiable operations, making its use as a loss function a natural idea. Several years ago, such attempts already existed, for example, in \href{https://papers.cool/arxiv/2003.11774}{\emph{Image Generation Via Minimizing Fr\'{e}chet Distance in Discriminator Feature Space}} and \href{https://papers.cool/arxiv/2009.14075}{\emph{Backpropagating through Fr\'{e}chet Inception Distance}}.

However, early attempts did not achieve impressive results. The fundamental reason lies in batch size. Typical loss functions compute per-sample losses and then average them. In contrast, FID computes means and covariances over the entire batch before evaluating the nonlinear function in \eqref{eq:w-p-q}. Consequently, FID estimated with small batches is biased, and this bias cannot be eliminated through prolonged training, only mitigated by increasing the batch size---leading to prohibitively high training costs.

The phrases ``nonlinear operations involving cross-sample statistics'' and ``requirement for large batch size'' may sound familiar. Indeed, contrastive learning in vision typically shares these two characteristics. Due to nonlinear interactions across samples, gradient accumulation cannot easily increase effective batch size---but this is not impossible to solve. See for example \href{https://kexue.fm/archives/8471}{\emph{Can Gradient Accumulation Be Used in Contrastive Learning?}}. As we shall see, the solution to the batch size problem in FID is conceptually similar.

From another perspective---using pretrained models to extract features for loss construction---there is also the related work on \href{https://papers.cool/arxiv/1603.08155}{\emph{Perceptual Loss}}. However, this serves as a reconstruction loss for individual samples and is typically used in training models like VAEs. It does not involve cross-sample statistical computations, and thus presents no computational difficulties.

\section{Gradient Computation}

We now proceed step by step to derive the gradients and understand the challenges in using FD as a loss. First, we need to compute the gradients. Since $p$ represents the real data distribution, the statistics $\boldsymbol{\mu}_p, \boldsymbol{\Sigma}_p$ are fixed, and we only need to compute gradients with respect to $\boldsymbol{\mu}_q, \boldsymbol{\Sigma}_q$.

The gradient with respect to $\boldsymbol{\mu}_q$ is straightforward:
\begin{equation}
\nabla_{\boldsymbol{\mu}_q}\mathcal{F} = \nabla_{\boldsymbol{\mu}_q}\|\boldsymbol{\mu}_p - \boldsymbol{\mu}_q\|^2 = 2(\boldsymbol{\mu}_q - \boldsymbol{\mu}_p)
\end{equation}

The gradient with respect to $\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}

We use the second line of \eqref{eq:w-p-q}, which is more complex but has an advantage: the matrix $\boldsymbol{S} = \boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2}$ is symmetric positive definite, a property that enables simplification. Let the eigen-decomposition of $\boldsymbol{S}$ be $\boldsymbol{U}\boldsymbol{\Lambda}\boldsymbol{U}^\top$, so $\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} \\
\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 derivations of eigenvalue gradients, refer to \href{https://kexue.fm/archives/10878}{\emph{Derivatives of SVD}}. The result resembles the derivative of $\sqrt{x}$, namely $\frac{1}{2\sqrt{x}}$, which feels intuitive. However, this is not trivial and generally fails if $\boldsymbol{S}$ is not symmetric positive definite. Finally, by the chain rule:
\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}
Combining everything, we obtain:
\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 may appear complex, but $\boldsymbol{\Sigma}_p^{1/2}$ can be precomputed. It suffices to compute the square root and inverse square root of the symmetric positive definite matrix $\boldsymbol{S} = \boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2}$, which can be done using the \texttt{eigh} function, or via the Newton-Schulz iteration method described in \href{https://kexue.fm/archives/11158}{\emph{Efficient Computation of Matrix Square Root and Inverse Square Root}} and \href{https://kexue.fm/archives/11175}{\emph{Efficient Computation of Matrix $r$-th Root and Inverse $r$-th Root}}.

\section{Large Batches}

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 \\
\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}
Then:
\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 (e.g., 2048 for InceptionV3), so accurate estimation often requires tens of thousands of samples. The real distribution is fixed and $\boldsymbol{\mu}_p,\boldsymbol{\Sigma}_p$ can be precomputed. However, the generated distribution changes dynamically, and computing accurate estimates every step with tens of thousands of samples implies an enormous batch size, which is computationally prohibitive in many settings.

On the other hand, the gradient formula \eqref{eq:Sigma-grad} also reveals the necessity of large batches. With small batches, the estimated $\boldsymbol{V}_q$ may not be full rank, so $\boldsymbol{\Sigma}_q$ may not be invertible, making $(\boldsymbol{\Sigma}_p^{1/2}\boldsymbol{\Sigma}_q\boldsymbol{\Sigma}_p^{1/2})^{-1/2}$ undefined (due to division by zero). Therefore, using FD as a loss imposes a requirement on the training batch size---this is likely the most critical practical difficulty.

Given computational constraints, we aim to simulate large batch effects using small batches, similar to the ``contrastive learning + gradient accumulation'' scenario.

\section{Equivalent Loss}

Suppose a batch size $B$ is required for accurate estimation of $\boldsymbol{\mu}_q,\boldsymbol{V}_q$. Instead, we process $k = B/b$ small batches of size $b$ with intermediate 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)}$, such that:
\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 seek an equivalent total gradient equal to the sum of gradients over small batches, achieving unbiased estimation. Differentiating both sides of \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}\left( \bigl[ \boldsymbol{\mu}_q - \tilde{\boldsymbol{\mu}}_q^{(i)}/k \bigr]_{\text{sg}} + \tilde{\boldsymbol{\mu}}_q^{(i)}/k, \bigl[ \boldsymbol{V}_q - \tilde{\boldsymbol{V}}_q^{(i)}/k \bigr]_{\text{sg}} + \tilde{\boldsymbol{V}}_q^{(i)}/k \right)
\end{aligned}
\end{equation}
This means we compute $\tilde{\boldsymbol{\mu}}_q^{(i)}, \tilde{\boldsymbol{V}}_q^{(i)}$ in small batches, average them to obtain an accurate $\boldsymbol{\mu}_q,\boldsymbol{V}_q$, then compute gradients using the loss:
\begin{equation}
\mathcal{F}_i = \mathcal{F}\left( \bigl[ \boldsymbol{\mu}_q - \tilde{\boldsymbol{\mu}}_q^{(i)}/k \bigr]_{\text{sg}} + \tilde{\boldsymbol{\mu}}_q^{(i)}/k, \bigl[ \boldsymbol{V}_q - \tilde{\boldsymbol{V}}_q^{(i)}/k \bigr]_{\text{sg}} + \tilde{\boldsymbol{V}}_q^{(i)}/k \right)
\label{eq:Fi}
\end{equation}
and sum the gradients. This yields gradients equivalent to using a batch size $B$. Here, $[\cdot]_{\text{sg}}$ denotes the stop-gradient operator. Alternatively, we can update parameters each step with a smaller learning rate, achieving a similar effect.

\section{Historical Averaging}

The above approach is theoretically sound, but not smooth: $k$ forward passes are needed before computing accurate $\boldsymbol{\mu}_q,\boldsymbol{V}_q$, and then gradients are computed retroactively. The bottleneck is the need for global statistics to compute unbiased local gradients.

A natural idea: can we approximate $\boldsymbol{\mu}_q,\boldsymbol{V}_q$? Given small learning rates, model parameters change slowly, so $\boldsymbol{\mu}_q,\boldsymbol{V}_q$ should also change slowly. After including a new batch, the updated statistic is only a small adjustment to the previous one. We use exponential moving averages (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 maintains an effective averaging window of size $\mathcal{O}(1/(1-\beta))$, effectively increasing the statistical batch size by a factor of $\mathcal{O}(1/(1-\beta))$. Thus, at each step we can compute the gradient using the loss:
\begin{equation}
\mathcal{F}_t = \mathcal{F}\left( \underbrace{\beta \bigl[ \boldsymbol{\mu}_q^{(t-1)} \bigr]_{\text{sg}} + (1-\beta) \tilde{\boldsymbol{\mu}}_q^{(t)}}_{\boldsymbol{\mu}_q^{(t)}}, \underbrace{\beta \bigl[ \boldsymbol{V}_q^{(t-1)} \bigr]_{\text{sg}} + (1-\beta) \tilde{\boldsymbol{V}}_q^{(t)}}_{\boldsymbol{V}_q^{(t)}} \right)
\end{equation}
Extra cost: storing $\boldsymbol{\mu}_q,\boldsymbol{V}_q$, which is minimal. This practice of ``making up for small batches with historical data'' embodies the ``streaming'' idea in \href{https://kexue.fm/search/streaming+power+iteration}{Streaming Power Iteration}. Additionally, the paper discusses a queue-based method that retains $k$ historical batches, computes gradients using \eqref{eq:Fi} and removes the oldest batch. This method is simpler but more memory-intensive than EMA and performs worse in practice.

\section{Experimental Overview}

The paper's experiments focus on post-training of generative models, aiming to improve one-step generation via FD Loss, or fine-tune multi-step models to one-step. When combining multiple encoders to compute FD Loss, a loss normalization technique balances losses of different scales:
\begin{equation}
\mathcal{L} = \sum_i \frac{\mathcal{F}[\phi_i]}{\bigl[ \mathcal{F}[\phi_i] \bigr]_{\text{sg}} + \epsilon}
\end{equation}
This technique was also discussed in \href{https://kexue.fm/archives/8870}{\emph{On Multi-Task Learning (I): The Name of the Loss}}.

The paper's key achievement is pushing one-step generation performance (FID) to unprecedented levels, surpassing both one-step and multi-step baselines. This result appears near the performance ceiling. Some figures are shown below:

\begin{figure}[htbp]
\centering
\includegraphics[width=0.95\linewidth]{assets/11738/3014248186.png}
\caption{FD loss improves one-step generation performance}
\end{figure}

\begin{figure}[htbp]
\centering
\includegraphics[width=0.95\linewidth]{assets/11738/3462357209.png}
\caption{Comprehensive comparison of FD loss}
\end{figure}

\section{Summary}

This article theoretically analyzes the challenges of using FID as a loss function for generative models and derives corresponding techniques to overcome them.

\end{document}
