Efficiently Access Diffusion Fisher: Within the Outer Product Span Space
Extracted source text · arXiv 2505.23264
%%%%%%%% ICML 2025 EXAMPLE LATEX SUBMISSION FILE %%%%%%%%%%%%%%%%%
\documentclass{article}
% Use the following line for the initial blind version submitted for review:
% \usepackage{icml2025}
% If accepted, instead use the following line for the camera-ready submission:
\usepackage[accepted]{icml2025}
% For theorems and such
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{tocloft}
\usepackage{etoc}
\usepackage{hyperref}
% Recommended, but optional, packages for figures and better typesetting:
\usepackage{microtype}
\usepackage{graphicx}
%\usepackage{subfigure}
\usepackage{subcaption}
\usepackage{booktabs} % for professional tables
\usepackage{diagbox}
\usepackage{cancel}
\usepackage{bbm}
\input{math_commands.tex}
\usepackage{amsthm}
\usepackage{hyperref}
\usepackage{url}
\usepackage{physics}
\usepackage{bbding}
\usepackage{pifont}
\usepackage{multirow}
\usepackage{graphicx}
\usepackage{multirow}
\usepackage{amssymb}
\usepackage{makecell}
\usepackage{tcolorbox}
\usepackage{xcolor}
\usepackage{caption}
\definecolor{mygreen}{RGB}{0, 150, 0}
\newcommand{\cmark}{\textcolor{mygreen}{\ding{51}}}%
\newcommand{\xmark}{\textcolor{red}{\ding{55}}}%
% if you use cleveref..
\usepackage[capitalize,noabbrev]{cleveref}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% THEOREMS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\theoremstyle{plain}
\theoremstyle{definition}
\theoremstyle{remark}
% Attempt to make hyperref and algorithmic work together better:
\newcommand{\theHalgorithm}{\arabic{algorithm}}
% \usepackage{algorithm}
% \usepackage{algpseudocode}
% \usepackage{amsfonts}
\usepackage{algorithm}
% \etocdepthtag.toc{mtchapter}
% \etocsettagdepth{mtchapter}{subsection}
% \etocsettagdepth{mtappendix}{none}
\newtheorem{example}{Example}
\newtheorem{proposition}{Proposition}
\newtheorem{definition}{Definition}
\newtheorem{theorem}{Theorem}
\newtheorem{remark}{Remark}
\newtheorem{assumption}{Assumption}
\newtheorem{corollary}{Corollary}
\newtheorem{note}{Note}
\newtheorem{lemma}{Lemma}
%\newtheorem{remark}[theorem]{Remark}
% Todonotes is useful during development; simply uncomment the next line
% and comment out the line below the next line to turn off comments
%\usepackage[disable,textsize=tiny]{todonotes}
\usepackage[textsize=tiny]{todonotes}
% The \icmltitle you define below is probably too long as a header.
% Therefore, a short form for the running title is supplied here:
\icmltitlerunning{Efficiently Access Diffusion Fisher: Within the Outer Product Span Space}
\begin{document}
\twocolumn[
\icmltitle{Efficiently Access Diffusion Fisher: Within the Outer Product Span Space}
% It is OKAY to include author information, even for blind
% submissions: the style file will automatically remove it for you
% unless you've provided the [accepted] option to the icml2025
% package.
% List of affiliations: The first argument should be a (short)
% identifier you will use later to specify author affiliations
% Academic affiliations should list Department, University, City, Region, Country
% Industry affiliations should list Company, City, Region, Country
% You can specify symbols, otherwise they are numbered in order.
% Ideally, you should not use this facility. Affiliations will be numbered
% in order of appearance and this is the preferred way.
\icmlsetsymbol{equal}{*}
\icmlsetsymbol{corres}{$\dagger$}
\icmlsetsymbol{intern}{$\ddagger$}
\begin{icmlauthorlist}
\icmlauthor{Fangyikang Wang}{zju,equal,intern}
\icmlauthor{Hubery Yin}{wxg,equal}
\icmlauthor{Shaobin Zhuang}{sjtu,intern}
\icmlauthor{Huminhao Zhu}{zju}
\icmlauthor{Yinan Li}{zju}
\icmlauthor{Lei Qian}{zju}
\icmlauthor{Chao Zhang}{zju}
%\icmlauthor{}{sch}
\icmlauthor{Hanbin Zhao}{zju}
\icmlauthor{Hui Qian}{zju}
\icmlauthor{Chen Li}{wxg}
%\icmlauthor{}{sch}
%\icmlauthor{}{sch}
\end{icmlauthorlist}
\icmlaffiliation{zju}{Zhejiang University.}
\icmlaffiliation{wxg}{WeChat Vision, Tencent Inc.}
\icmlaffiliation{sjtu}{Shanghai Jiao Tong University.
$^*$Equal contribution.
$^\ddagger$Work done as interns at WeChat Vision, Tencent Inc}
\icmlcorrespondingauthor{Chao Zhang}{zczju@zju.edu.cn}
% \icmlcorrespondingauthor{Firstname2 Lastname2}{first2.last2@www.uk}
% You may provide any keywords that you
% find helpful for describing your paper; these are used to populate
% the "keywords" metadata in the PDF but will not be shown in the document
\icmlkeywords{Machine Learning, Generative Models, ICML} %关键词待填
\vskip 0.3in]
% this must go after the closing bracket ] following \twocolumn[ ...
% This command actually creates the footnote in the first column
% listing the affiliations and the copyright notice.
% The command takes one argument, which is text to display at the start of the footnote.
% The \icmlEqualContribution command is standard text for equal contribution.
% Remove it (just {}) if you do not need this facility.
\printAffiliationsAndNotice{} % leave blank if no need to mention equal contribution
%\printAffiliationsAndNotice{\icmlEqualContribution} % otherwise use the standard text.
\begin{abstract}
% Diffusion models (DMs) have demonstrated powerful distributional modeling capabilities by matching the first-order diffusion score.
Recent Diffusion models (DMs) advancements have explored incorporating the second-order diffusion Fisher information (DF), defined as the negative Hessian of log density, into various downstream tasks and theoretical analysis.
However, current practices typically approximate the diffusion Fisher by applying auto-differentiation to the learned score network. This black-box method, though straightforward, lacks any accuracy guarantee and is time-consuming.
In this paper, we show that the diffusion Fisher actually resides within a space spanned by the outer products of score and initial data.
Based on the outer-product structure, we develop two efficient approximation algorithms to access the trace and matrix-vector multiplication of DF, respectively.
These algorithms bypass the auto-differentiation operations with time-efficient vector-product calculations.
Furthermore, we establish the approximation error bounds for the proposed algorithms.
Experiments in likelihood evaluation and adjoint optimization demonstrate the superior accuracy and reduced computational cost of our proposed algorithms.
Additionally, based on the novel outer-product formulation of DF, we design the first numerical verification experiment for the optimal transport property of the general PF-ODE deduced map.
\end{abstract}
\input{content/introduction}
\input{content/preliminary}
\input{content/analytical}
\input{content/trace}
\input{content/endpoint}
\input{content/ot}
\input{content/conclusion}
\section*{Acknowledgments}
This work was supported in part by the National Natural Science Foundation of China under Grants 62206248 and 62402430, and the Zhejiang Provincial Natural Science Foundation of China under Grant LQN25F020008.
We would like to thank all the reviewers for their constructive comments.
Fangyikang Wang wishes to express gratitude to Pengze Zhang from ByteDance, Binxin Yang, and Xinhang Leng from WeChat Vision for their insightful discussions on the experiments.
\section*{Impact Statement}
The development of accurate diffusion Fisher, as discussed in this paper, holds significant potential for several domains, including machine learning, healthcare, environmental modeling, and economics.
However, while this research holds great potential for positive impacts, it is also important to consider potential negative societal impacts.
The enhanced ability of generative models given by DF could potentially be misused.
For instance, it could be exploited to create aesthetically improved deepfakes, leading to misinformation.
In healthcare, if not properly regulated, the use of synthetic patient data could lead to ethical issues.
Therefore, it is crucial to ensure that the findings of this research are applied ethically and responsibly, with necessary safeguards in place to prevent misuse and protect privacy.
\bibliographystyle{icml2025}
\bibliography{example_paper}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% APPENDIX
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\onecolumn
{\huge \bfseries Appendix}
% \etocdepthtag.toc{mtappendix}
% \etocsettagdepth{mtchapter}{none}
% \etocsettagdepth{mtappendix}{subsection}
% \tableofcontents
\appendix
\input{content/appendix/proofs}
\input{content/appendix/extra}
\input{content/appendix/discussion}
\end{document}
%%%%% NEW MATH DEFINITIONS %%%%%
\usepackage{amsmath,amsfonts,bm}
% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}
% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
\def\vlambda{{\bm{\lambda}}}
% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\newcommand{\vepsilon}{\boldsymbol{\varepsilon}}
\section{Introduction}
The emerging diffusion models (DMs) \cite{sohl2015deep, ho2020denoising,song2019generative, song2020score}, generating samples of data distribution from initial noise by learning a reverse diffusion process, have been proven to be an effective technique for modeling data distribution, especially in generating high-quality images \cite{pmlr-v162-nichol22a, NEURIPS2021_49ad23d1, NEURIPS2022_ec795aea, ramesh2022hierarchical, rombach2022high,JMLR:v23:21-0635}. The training process of DMs can be seen as employing a neural network to match the first-order diffusion score $\nabla_\vx \log q_t(\vx)$ at varying noise levels.
There has been a growing trend to recognize the importance of the diffusion Fisher (DF) in DMs, defined as the negative Hessian of the diffused distributions' log density, $-\nabla^2_\vx \log q_t(\vx)$.
The DF provides valuable second-order information on DMs and plays a crucial role in downstream tasks, e.g., likelihood evaluation \citep{pmlr-v162-lu22f, pmlr-v202-zheng23c}, adjoint optimization \citep{pan2023adjointdpm,pan2023towards,blasingame2024adjointdeis}, and the optimal transport analysis of DMs \cite{zhang2024formulating}.
%The DF provides valuable second-order information on DMs and plays a crucial role in both practice tasks, e.g. likelihood evaluation \citep{pmlr-v162-lu22f, pmlr-v202-zheng23c} and adjoint optimization \citep{pan2023adjointdpm,pan2023towards,blasingame2024adjointdeis}, and theoretical analysis, e.g. the optimal transport property \cite{zhang2024formulating} and convergence property \cite{benton2020nearly} of DMs.
It is straightforward to access DFs by applying auto-differentiation to the learned first-order diffusion score.
Nevertheless, the repeated auto-differentiation operations would be time-consuming for large-scale tasks, even when incorporating the Vector-Jacobian-product (VJP) and Hutchinson's estimation technique \cite{song2024fisher}.
Recently, certain research efforts investigated the time-evolving property of DM and demonstrated that the DF can be explicitly expressed in terms of the score and its covariance matrix \cite{benton2020nearly,pmlr-v162-lu22f}.
However, these results predominantly concentrate on the theoretical aspects and fail to enable efficient access to the DF.
In this paper, we derive a novel form for the diffusion Fisher, which is composed of weighted outer-product sums.
This innovative formulation uncovers that the diffusion Fisher invariably lies within the span of a set of outer-product bases.
Notably, these bases solely depend on the initial data and the noise schedule.
Leveraging this outer-product structure, we devise two efficient approximation algorithms to access the trace and matrix-vector multiplication of DF, respectively.
Moreover, we design the first numerical verification experiment to investigate the optimal transport property of the diffusion-ODE deduced map under diverse types of initial data and noise schedules.
The main contributions of our paper are listed as follows:
\begin{itemize}
\item
We derive an analytical formulation of the diffusion Fisher, which is the weighted summation of the outer product, under the setting when the initial distribution is a sum of Dirac.
The derivation process involves applying the consecutive partial differential chain rule to the data-dependent form of marginal distributions.
Subsequently, we extend this formulation to a more general setting. In this general case, we only assume that the initial distribution has a finite second moment.
These novel formulations reveal that the diffusion Fisher invariably lies within the span of a set of outer-product bases.
\item
Based on the outer-product structure of DF, we propose two efficient approximation algorithms, each tailored to access the trace and matrix-vector multiplication of the DF, respectively.
When it comes to evaluating the trace of the diffusion Fisher, we introduce a parameterized network to learn the trace, significantly reducing the time complexity of trace evaluation from quadratic to linear w.r.t. the dimension.
In situations where we need to access the matrix-vector multiplication of the diffusion Fisher, we present a training-free method that simplifies the complex linear transformation into several simple vector-product calculations.
Furthermore, we rigorously establish the approximation error bounds for these two algorithms.
% Leveraging the novel insights into the Diffusion Fisher, we derive a corollary and design the first numerical experiment to examine the optimal transport (OT) property of the map deduced from the general diffusion ordinary differential equation (ODE).
% Our numerical experiments indicate that for all commonly employed noise schedules (VE, VP, sub-VP, and EDM), the diffusion-ODE map exhibits the OT property when the initial data follows a single Gaussian distribution or is affine, but does not exhibit this property when the initial data is non-affine.
% These numerical experiments validate the theoretical OT conclusions for the special cases presented in \cite{khrulkov2023understanding} and \cite{zhang2024formulating}, and also enable the exploration of general cases.
\item
Utilizing the novel outer-product formulation of DF, we derive a corollary for the optimal transport (OT) property of the diffusion-ODE deduced map.
Subsequently, we design the first numerical experiment to examine the OT property of any general diffusion ODE deduced map based on the corollary.
Our numerical experiments indicate that for all commonly employed noise schedules (VE \cite{ho2020denoising}, VP \cite{song2019generative}, sub-VP \cite{song2020score}, and EDM \cite{karras2022elucidating}), the diffusion ODE map exhibits the OT property when the initial data follows a single Gaussian distribution or has an affine form.
Conversely, it does not exhibit this property when the initial data is non-affine.
These numerical results not only validate the conclusions presented in \cite{khrulkov2023understanding, zhang2024formulating,lavenant2022flow}, but also open up new avenues for the exploration of more general cases.
\end{itemize}
We evaluate our DF access algorithms on likelihood evaluation and adjoint optimization tasks. The empirical results demonstrate our DF methods' enhanced accuracy and reduced time cost.
\section{Preliminaries}
\textbf{Notation}. The Euclidean norm over $\sR^{d}$ is denoted by $\norm{\cdot}$, and the Euclidean inner product is denoted by $\ip{\cdot}{\cdot}$.
% The Lebesgue measure on $\sR^d$ is denoted as $\mathcal{L}(\sR^d)$, and $\ll$ means absolutely continuous.
Throughout, we simply write $\int g$ to denote the integral with respect to the Lebesgue measure: $\int g(x)\rd x$.
When the integral is with respect to a different measure $\mu$, we explicitly write $\int g \rd \mu$.
% The expectation and variance of $g(X)$when $X \sim p$ are respectively denoted $\E_\mu g=\int g\rd\mu$ and $\text{var}_\mu g:=\int(g-\E_\mu g)^2\rd \mu$.
When clear from context, we sometimes abuse notation by identifying a measure $\mu$ with its Lebesgue density.
We also use $\delta(\cdot)$ to denote the Dirac Delta function.
For a vector $\vv\in\sR^d$, we denote the $d\times d$ matrix $\vv\vv^\top$ as the outer-product of $\vv$.
\subsection{Diffusion Models and Diffusion SDEs}
Suppose that we have a d-dimensional random variable $\vx_0 \in \sR^d$ following an unknown target distribution $q_0(\vx_0)$. Diffusion Models (DMs) define a forward process $\{\vx_t\}_{t\in[0,T]}$ with $T>0$ starting with $\vx_0$, such that the distribution of $\vx_t$ conditioned on $\vx_0$ satisfies
\begin{equation} \label{diffusion_kernel}
% \small
\begin{aligned}
% (\text{Diffusion Transition Kernel})~~~~~~~
q_{t\vert 0}(\vx_t \vert \vx_0) = \mathcal{N}(\vx_t; \alpha(t) \vx_0, \sigma^2(t)\mathbf{I}),
\end{aligned}
\end{equation}
where $\alpha(\cdot),\sigma(\cdot)\in \mathcal{C}(\left[ 0,T\right],\mathbb{R}^{+})$ have bounded derivatives, and we denote them as $\alpha_t$ and $\sigma_t$ for simplicity. The choice for $\alpha_t$ and $\sigma_t$ is referred to as the noise schedule of a DM.
According to \cite{NEURIPS2021_b578f2a5,karras2022elucidating}, with some assumption on $\alpha(\cdot)$ and $\sigma(\cdot)$, the forward process can be modeled as a linear SDE which is also called the Ornstein–Uhlenbeck process:
\begin{equation}\label{diffusion_sde}
% \small
\begin{aligned}
\rd \vx_t = f(t) \vx_t \rd t + g(t) \rd B_t,
\end{aligned}
\end{equation}
where $B_t$ is the standard d-dimensional Brownian Motion (BM),
% \begin{equation}\label{f_g_detail}
$f(t) = \frac{\rd \log \alpha_t}{\rd t}$ and $g^2(t)=\frac{\rd \sigma_t^2}{\rd t}-2 \frac{\rd \log \alpha_t}{\rd t}\sigma_t^2$.
% \end{equation}
Under some regularity conditions, the above forward SDE \eqref{diffusion_sde} have a reverse SDE from time $T$ to $0$, which starts from $\vx_t$ \cite{anderson1982reverse}:
\begin{equation} \label{diffusion_sde_rev}
% \small
\begin{aligned}
\rd \vx_t = \left[f(t) \vx_t - g^2(t) \nabla_{\vx_t} \log q(\vx_t, t)\right]\rd t + g(t) \rd \tilde{B}_t,
\end{aligned}
\end{equation}
where $\tilde{B}_t$ is the reverse-time Brownian motion and $q(\vx_t, t)$ is the single-time marginal distribution of the forward process. In practice, DMs \cite{ho2020denoising,song2020score} use $\vepsilon_\theta(\vx_t,t)$ to estimate $-\sigma(t)\nabla_{\vx_t} \log q(\vx_t, t)$ and the parameter $\theta$ is optimized by the following loss:
\begin{equation}\label{dpm_objective}
\small
\begin{aligned}
% \theta^* = \mathop{\arg\min}\limits_{\theta}
\mathcal{L}=\mathbb{E}_t \left\{ \lambda_t \mathbb{E}_{x_0, x_t} \left[ \lVert s_\theta(x_t, t) - \nabla_{x_t} \log p(x_t, t | x_0, 0) \rVert^2 \right]\right\},
\end{aligned}
\end{equation}
where $s_\theta$ represents the parameterized score, i.e., $s_\theta(\vx_t,t) = -\frac{\vepsilon_\theta(\vx_t,t)}{\sigma_t}$.
This familiar parameterization is called $\epsilon$-prediction. There are also $\vy$-prediction and $\vv$-prediction \cite{salimans2022progressive}. The corresponding loss is equal to replace the term $|\epsilon - \vepsilon_\theta(\vx_t, t)|$ with $\frac{\alpha_t}{\sigma_t} |\vx_0 - \bar{\vy}_\theta(\vx_t, t)|$ and $ |\alpha_t \epsilon - \sigma_t \vx_0 - \vv_\theta(\vx_t, t)|$. The learned $\vepsilon_\theta(\vx_t, t)$ can be also transformed to a $\vy$-prediction form by $\bar{\vy}_\theta(\vx_t, t) = \frac{\vx_t - \sigma_t \vepsilon_\theta(\vx_t, t)}{\alpha_t}$.
\subsection{Diffusion Models Inference as Neural ODE}
% The diffusion process along with its sampling processes in DMs can be delineated as the forward and corresponding backward stochastic differential equations (SDE) \cite{song2020score, anderson1982reverse}.
It is noted that the reverse diffusion SDE in \eqref{diffusion_sde_rev} has an associated probability flow ODE (PF-ODE, also called diffusion ODE), which is a deterministic process that shares the same single-time marginal distribution \cite{song2020score}:
% . The diffusion ODE is governed by the following:
% The SDE has an associated probability flow ODE, which has the same marginal distribution at each time $t$:
\begin{equation}\label{PFODE}
% \small
\begin{aligned}
% (\text{PF-ODE})~~~~~~~
\rd \vx_t = \left[ f(t)\vx_t - \frac{1}{2}g^2(t) \nabla_{\vx_t} \log q_t(\vx_t, t) \right]\rd t.
\end{aligned}
\end{equation}
By replacing the score function in \eqref{PFODE} with the noise predictor $\vepsilon_\theta$, the inference process of DMs can be constructed by the following neural ODE:
\begin{equation}\label{dm_neural_ode}
% \small
\begin{aligned}
\frac{\mathrm{d} \boldsymbol{x}_t}{\mathrm{~d} t}=\boldsymbol{h}_\theta\left(\boldsymbol{x}_t, t\right):=f(t) \boldsymbol{x}_t+\frac{g^2(t)}{2 \sigma_t} \boldsymbol{\epsilon}_\theta\left(\boldsymbol{x}_t, t\right), \quad
\end{aligned}
\end{equation}
where initial $\boldsymbol{x}_T \sim \mathcal{N}\left(\mathbf{0}, \sigma_T^2 \boldsymbol{I}\right)$.
\subsection{Fisher Information in Diffusion Models}
The diffusion Fisher information matrix in DMs is defined as the negative Hessian of the marginal diffused log-density function, which takes the following matrix-valued form \cite{song2021maximum,song2024fisher}:
\begin{equation} \label{fisher_def}
% \small
\begin{aligned}
\mF_t(\vx_t,t) &:= - \frac{\partial^2}{\partial \vx_t^2} \log q_t\left(\vx_t, t\right)\\
\end{aligned}
\end{equation}
The current technique typically approximately accesses the diffusion Fisher by accessing the scaled Jacobian matrix of the learned score estimator network $\vepsilon_\theta$:
\begin{equation} \label{}
% \small
\begin{aligned}
\mF_t(\vx_t,t) &= - \frac{\partial}{\partial \vx_t}\left(\frac{\partial}{\partial \vx_t} \log p\left(x_t, t\right) \right)\\
&\approx - \frac{\partial}{\partial \vx_t}\left(-\frac{\vepsilon_\theta(\vx_t,t)}{\sigma_t} \right) = \frac{1}{\sigma_t}\frac{\partial \vepsilon_\theta(\vx_t,t)}{\partial\vx_t}
\end{aligned}
\end{equation}
The full diffusion Fisher matrix within DMs cannot be obtained due to dimensional constraints.
For instance, the Stable Diffusion-1.5 model \citep{rombach2022high} features a latent dimension of $d = 4\times64\times64 = 16384$, resulting in a diffusion Fisher matrix of $16384 \times 16384$.
Fortunately, for applications that only need to access the trace or multiplication of diffusion Fisher, it is feasible to use the Vector-Jacobian-product (VJP) to access diffusion Fisher.
For any $d$-dimensional vector $\vv$, the approximation of $\vv$ left multiplied by $\mF_t(\vx_t,t)$ using VJP is as follows:
\begin{equation} \label{VJP}
% \small
\begin{aligned}
(\text{VJP})~~~~~~~~~~
\vv^\top\mF_t(\vx_t,t) &\approx \frac{1}{\sigma_t} \vv^\top\frac{\partial \vepsilon_\theta(\vx_t,t)}{\partial\vx_t} \\
&= \frac{1}{\sigma_t}\frac{\partial \left[\ip{\vepsilon_\theta(\vx_t,t)}{\vv} \right]}{\partial\vx_t}
\end{aligned}
\end{equation}
The VJP is a time-consuming process due to its requirement for gradient calculations within the neural network. In addition, empirical evidence from synthetic distributions, as demonstrated in \citet{pmlr-v162-lu22f}, shows that the approximation results from the VJP significantly deviate from the true underlying diffusion Fisher. To our knowledge, there is no theoretical guarantee that the diffusion Fisher can be accurately accessed through the VJP. Moreover, the VJP fails to provide any theoretical insight into the diffusion Fisher.
% \subsubsection{Adjoint Optimization}
\begin{table}[t]
\centering
\resizebox{0.5\textwidth}{!}{%
\begin{tabular}{c||c|ccc}
\Xhline{1pt}
\multirow{2}{*}{Access} & \multirow{2}{*}{\makecell{Methods}} & \makecell{Theoretical\\ Time-cost} & \makecell{Practical\\ Time-cost (s)} & \makecell{Theoretical\\ Error Bound} \\
\Xhline{1pt}
\multirow{2}{*}{Trace} &VJP (eq. \ref{nll_VJP}) & $c_1d+c_2d^2$ & 2195.48& \xmark \\
&DF-TM \textbf{(Ours)} & $2c_1d$ & \textbf{0.072}\textcolor{red}{\textsubscript{\text{$99\% \downarrow$}}} & \cmark (Prop. \ref{prop:error_bound_trace}) \\
\Xhline{0.5pt}
\multirow{2}{*}{Operator} &VJP (eq. \ref{Adjoint_VJP}) & $c_1d+c_2d$ & 0.155& \xmark \\
&DF-EA \textbf{(Ours)}& $2c_1d$ & \textbf{0.063}\textcolor{red}{\textsubscript{\text{$59\% \downarrow$}}}& \cmark (Prop. \ref{prop:error_bound_ea})\\
\Xhline{1pt}
\end{tabular}
}
\caption{\small Comparison of Vector-Jacobian-product (VJP) and our proposed efficient access methods in terms of per-iteration theoretical time-cost, practical time-cost, and approximation error bound. The theoretical time cost is based on assumptions of a network access cost of $c_1d$ time and backpropagation on network cost of $c_2d$ time ($c_2\approx4c_1$). The practical time-cost is tested using the SD-v1.5 model over 10k COCO prompts. Here, our VJP baseline calculates every element of the trace, and discussion on its approximation is deferred to Appendix \ref{app:hutchinson}.
}
\vspace{-3mm}
\label{tab:my_label}
\end{table}
\section{Diffusion Fisher in Outer-Products Span}\label{sec:analytical}
Accessing diffusion Fisher via the VJP as shown in \eqref{VJP} is straightforward, but it does not take advantage of any inherent structure of the diffusion Fisher.
Though certain recent research efforts \cite{pmlr-v162-lu22f,benton2020nearly} demonstrated that the DF can be explicitly expressed in terms of the score and its covariance matrix. Nevertheless, their formulations cannot directly facilitate efficient access to the diffusion Fisher.
In this section, we derive a novel form for the diffusion Fisher, which is composed of weighted outer-product sums.
We first conduct this under a simplified scenario, assuming that the initial distribution $q_0$ is a sum of Dirac distributions. Subsequently, we generalize this outer-product span form to a more general setting.
Significantly, the outer-product span form of the diffusion Fisher obtained in both settings does not require any gradient calculations and depends solely on the initial data and the noise schedule.
Notably, there exist fast methods for accessing the trace and matrix-vector multiplication of an outer-product matrix.
This advantageous characteristic allows for the development of novel algorithms for efficiently accessing the diffusion Fisher.
% A detailed discussion on this topic is provided in Appendix \ref{app:covariance}.
\begin{figure}[t!]
\small
\centering
\begin{minipage}{0.24\textwidth}
\centering
\includegraphics[width=1\textwidth]{figures/tnet_loss_sd2b.pdf}
\subcaption{\scriptsize }\label{fig:train_loss}
\end{minipage}\hfill
\begin{minipage}{0.24\textwidth}
\centering
\includegraphics[width=1\textwidth]{figures/guidance.pdf}
\subcaption{\scriptsize }\label{fig:tradeoff}
\end{minipage}
\vspace{-3mm}
\caption{(a) {\small The training loss of DF-TM for SD-1.5 and SD-2base. It demonstrates commendable convergence behavior.} (b) {\small The trade-off curve of NLL and Clip score of SD-1.5 and SD-2base across various guidance scales in [1.5, 2.5, ..., 12.5, 13.5]}}
\label{fig:xx}
\end{figure}
% \subsection{The Dirac Setting}
We start with a simple setting where we assume that the initial distribution is characterized as a sum of Dirac distributions composed of the set of samples in the dataset. If we suppose the dataset is denoted as $\{ \vy_i \}_{i=0}^N$, then the initial distribution follows
\begin{equation} \label{dirac_initial_dist}
% \small
\begin{aligned}
(\text{Dirac Setting})~~~~~~~~~~q(\vx, t)|_{t=0} = \frac{1}{N} \sum_{i=0}^{N} \delta(\vx - \vy_i),
\end{aligned}
\end{equation}
where exists a $0<\mathcal{D}_y<\infty$ such that $\norm{\vy_i}\leq\mathcal{D}_y$ holds true for every $i$.
In this Dirac setting, we derive the following formulation of diffusion Fisher, which is a weighted outer-product sum devoid of gradients and composed solely of the initial distribution and the noise schedule.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{proposition}\label{prop:analytic_dirac}
Defines $v_i(\vx_t, t)$ as $\exp(-\frac{|\vx_t - \alpha_t \vy_i|^2}{2\sigma_t^2})\in \sR$ and $w_i(\vx_t, t)$ as $\frac{v_i(\vx_t, t)}{\sum_j v_j(\vx_t, t)} \in \sR$. If $q_0$ takes the form as in equation \eqref{dirac_initial_dist}, the diffusion Fisher matrix of the diffused distribution $q_t$ for $t\in(0,1]$ can be analytically formulated as follows:
\begin{equation} \label{information_dirac}
\small
\begin{aligned}
% (\text{Dirac diffusion Fisher formulation})~~~~~~~~~~
\mF_t(\vx_t,t) =& \frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left[ \sum_i w_i \vy_i\vy_i^{\top} \right. \\ & \left.- \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right]
\end{aligned}
\end{equation}
where we have simplified $w_i(\vx_t, t)$ to $w_i$, as it does not lead to any confusion.
\end{proposition}
\end{tcolorbox}
We also find that the $\sum_i w_i \vy_i$ component in \eqref{information_dirac} can be effectively approximated by the trained score network in the form of $y$-prediction, as demonstrated in the following proposition.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{proposition}\label{prop:optimal_y}
Given the diffusion training loss in \eqref{dpm_objective}, and if $q_0$ conforms to the form presented in \eqref{dirac_initial_dist}, then the optimal $\bar{\vy}_\theta(\vx_t, t)$ can accurately estimate $\sum_i w_i \vy_i$.
\end{proposition}
\end{tcolorbox}
% Given that $\vepsilon_\theta$ is designed to align with $-\sigma(t)\nabla_{\vx(t)} \log p(\vx(t), t)$, and in conjunction with Proposition \ref{prop:analytic_dirac} and \ref{prop:optimal_y}, we can deduce the following:
% % \begin{corollary}\label{corollary:theta_dirac}
% \begin{equation} \label{theta_dirac}
% \begin{aligned}
% \frac{\partial \vepsilon_\theta(\vx_t,t)}{\partial\vx_t} &\approx \frac{1}{\sigma_t}\mI - \frac{\alpha_t^2}{\sigma_t^3}\left( \sum_i w_i \vy_i\vy_i^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)
% \end{aligned}
% \end{equation}
% \end{corollary}
\paragraph{The General Setting}
We further extend the diffusion Fisher formulation in \eqref{information_dirac} to a more general setting, where we only assume that the initial distribution $q_0$ is a probability measure on $\sR^d$ with finite second moments.
% and is absolutely continuous w.r.t. the Lebesgue measure.
\begin{equation} \label{general_initial_dist}
% \small
\begin{aligned}
(\text{General Setting})~~~~~~~~~~ q_0\in \mathcal{P}_2(\sR^d),
% ,\quad q_0 \ll \mathcal{L}(\sR^d)
\end{aligned}
\end{equation}
% In this general setting, we derive the analogous propositions to Proposition \ref{prop:analytic_dirac} and \ref{prop:optimal_y} in Appendix \ref{app:general_df}.
which means that $q_0\in \mathcal{P}(\mathbb{R}^d),\quad \int_{\mathbb{R}^d} ||x||^{2} q_0(x)\rd x<\infty$.
In this general setting, we derive the following form of diffusion Fisher, which is a weighted outer-product integral, which resides in a space spanned by an infinite set of outer products.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt] \small
\begin{proposition}\label{prop:analytic_general}
Let us define $v(\vx_t, t, \vy)$ as $\exp(-\frac{|\vx_t - \alpha_t \vy|^2}{2\sigma_t^2})\in \sR$ and $w(\vx_t, t, \vy)$ as $\frac{v(\vx_t, t, \vy)}{\int_{\sR^d} v(\vx_t, t, \vy)\rd q_0(\vy)} \in \sR$. If $q_0$ takes the form as in \eqref{general_initial_dist}, the diffusion Fisher matrix of the diffused distribution $q_t$ for $t\in(0,1]$ can be analytically formulated as follows:
\begin{equation} \label{information_general}
\small
\begin{aligned}
% (\text{General diffusion Fisher formulation})~~~~
\mF_t(\vx_t,t)\! =& \frac{1}{\sigma_t^2}\mI -
\frac{\alpha_t^2}{\sigma_t^4}\left[ \int w(\vy) \vy\vy^{\top}\rd q_0 \right. \\ & \left. \!- \!\left(\int w(\vy) \vy \rd q_0\right)\left(\int w(\vy) \vy \rd q_0\right)^{\top} \right]
\end{aligned}
\end{equation}
where we simply write $w(\vx_t, t, \vy)$ as $w(\vy)$, as long as it does not lead to any confusion.
\end{proposition}
\end{tcolorbox}
We further ascertain that the $\int w(\vy) \vy \rd q_0(\vy)$ component in \eqref{information_general} can be effectively approximated by the score network in the form of $y$-prediction, as demonstrated in the following proposition.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt] \small
\begin{proposition}\label{prop:optimal_y_general}
Given the diffusion loss in \eqref{dpm_objective}, and if $q_0$ conforms to the form in \eqref{general_initial_dist}, then the optimal $\bar{\vy}_\theta(\vx_t, t)$ can accurately estimate $\int w(\vy) \vy \rd q_0(\vy)$.
\end{proposition}
\end{tcolorbox}
The derivation of the outer-product form of diffusion Fisher under the general setting is akin to the Dirac case, but in an integral form.
For the remainder of the paper, we will focus on developing our method based on the Dirac setting diffusion Fisher. However, the same results can be naturally extended to the general setting.
% Similar to \eqref{theta_dirac}, and in conjunction with Proposition \ref{prop:analytic_general} and \ref{prop:optimal_y_general}, we can deduce the following:
% \begin{equation} \label{theta_general}
% \begin{aligned}
% \frac{\partial \vepsilon_\theta(\vx_t,t)}{\partial\vx_t} &\approx \frac{1}{\sigma_t}\mI - \frac{\alpha_t^2}{\sigma_t^3}\left( \int w(\vy) \vy\vy^{\top}\rd q_0(\vy) - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)
% \end{aligned}
% \end{equation}
% \section{Methodology}
% \subsection{Gauss-Newton Transformation}
% \begin{figure}[t!]
% \begin{minipage}{0.23\textwidth}
% \centering
% \includegraphics[width=0.95\textwidth]{figures/tnet_loss_sd2b.pdf}
% \captionof{figure}{The training loss of DF-TM for SD-1.5 and SD-2base. It demonstrates commendable convergence behavior.}
% \label{fig:train_loss}
% \end{minipage}
% \hfill
% \begin{minipage}{0.23\textwidth}
% \centering
% \includegraphics[width=0.95\textwidth]{figures/guidance.pdf}
% \captionof{figure}{The trades-off curve of NLL and Clip score of SD-1.5 and SD-2base across various guidance scales in [1.5, 2.5, ..., 12.5, 13.5]}
% \label{fig:tradeoff}
% \end{minipage}
% \end{figure}
\begin{figure*}[!t]
\centering
\includegraphics[width=0.90\textwidth]{figures/nll.drawio.png} % Reduce the figure size so that it is slightly narrower than the column.
\caption{\small Our DF-TM method facilitates the effective evaluation of the NLL of generated samples with varying seeds. It can be demonstrated that a lower NLL signifies a region of higher possibility, thereby consistently indicating superior image quality.
}
\label{fig:nll_compare}
\end{figure*}
\section{Diffusion Fisher Trace Matching}\label{sec:tm}
% \vspace{-2mm}
The likelihood evaluation of DMs would require access to diffusion Fisher's trace. In this section, we introduce a network to learn the trace, thus facilitating effective likelihood evaluation in DMs.
\paragraph{Log-Likelihood in DMs}
Log-likelihood is a classic and significant metric for probabilistic generative models, extensively utilized for comparison between samples or models \cite{bengio2013bounding,theis2015note}.
According to \citet{chen2018neural,song2021maximum}, the log-likelihood of samples generated by PF-ODE in \eqref{dm_neural_ode} from DMs can be computed through a connection to continuous normalizing flows as follows:
\begin{equation} \label{nll_ode}
\small
\begin{aligned}
\frac{\partial \log q_t (\vx_t,t)}{\partial t} &\!=\! -\text{tr}\left( \frac{\partial }{\partial \vx_t} \left(f(t)\vx_t - \frac{1}{2}g^2(t) \partial_{\vx_t} \log q_t(\vx_t, t)\right)\right) \\
&\!=\! -\text{tr}\left( \left(f(t)\mI - \frac{1}{2}g^2(t) \frac{\partial^2}{\partial {\vx_t}^2} \log q_t(\vx_t, t)\right)\right)\\
&\!=\! -f(t)d-\frac{g^2(t)}{2 }\text{tr}\left(\mF_t(\vx_t, t)\right)
\end{aligned}
\end{equation}
where $\text{tr}(\cdot)$ denotes the trace of a matrix, which is defined to be the sum of elements on the diagonal.
\begin{figure*}[t!]
\small
\centering
\begin{minipage}{0.2\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figures/aes_sac.pdf}\vspace{-2mm}
\subcaption{\scriptsize SAC Aesthetic Score ($\big\uparrow$)}
\end{minipage}\hfill
\begin{minipage}{0.2\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figures/aes_ava.pdf}\vspace{-2mm}
\subcaption{\scriptsize AVA Aesthetic Score ($\big\uparrow$)}
\end{minipage}\hfill
\begin{minipage}{0.2\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figures/pick_score.pdf}\vspace{-2mm}
\subcaption{\scriptsize Pick-Score ($\big\uparrow$)}
\end{minipage}\hfill
\begin{minipage}{0.2\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figures/clip_score.pdf}\vspace{-2mm}
\subcaption{\scriptsize Clip Loss ($\big\downarrow$)}
\end{minipage}\hfill
\begin{minipage}{0.2\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figures/face_score.pdf}\vspace{-2mm}
\subcaption{\scriptsize Face ID Loss ($\big\downarrow$)}
\end{minipage}
\caption{\small Comparison between our DF method and the VJP method on adjoint guidance sampling across five objective scores: SAC/AVA aesthetic score, Pick-Score, clip loss, and Face ID loss. Notably, our DF consistently achieves superior scores with less time expenditure.}
\vspace{-3mm}
\label{fig:adjoint_score}
\end{figure*}
\paragraph{Log-Likelihood Evaluation via VJP}
The current technique is only capable of conducting backpropagation of scalar value to the neural network. Therefore, the VJP in \eqref{VJP} cannot directly calculate the trace of the diffusion Fisher. The VJP must iterate through each dimension to compute the individual elements on the diagonal, and then sum them up as follows
\begin{equation} \label{nll_VJP}
% \small
\begin{aligned}
% (\text{VJP for trace})~~~~~~~~~~
\text{tr}\left(\mF_t(\vx_t, t)\right)\approx \frac{1}{\sigma_t} \sum_{i=1}^d \frac{\partial \left[\ip{\vepsilon_\theta(\vx_t,t)}{\ve^{(i)}} \right]}{\partial\vx_t}. \\
\end{aligned}
\end{equation}
Evaluating the trace using the VJP method would be extremely time-consuming due to the curse of dimensionality.
If the time-complexity of a single backpropagation is $\mathcal{O}(d)$, then the calculation in \eqref{nll_VJP} would have a time-complexity of $\mathcal{O}(d^2)$. In practice, as demonstrated in Table \ref{tab:my_label}, evaluating the trace of diffusion Fisher on the SD-v1.5 model would require half an hour, rendering it nearly infeasible.
Some Monte-Carlo type approximations can be used to accelerate the NLL evaluation of models, but they still fall short in per-sample NLL evaluation, see a detailed discussion in \ref{app:hutchinson}.
\paragraph{Log-Likelihood Evaluation via DF trace matching}
To overcome the limitations of the VJP method in evaluating the trace of the diffusion Fisher, we propose to directly obtain its analytical form and train a network to match it.
Note that the trace of an outer-product of the vector results in the vector's norm.
Thus, given the diffusion Fisher in Proposition \ref{prop:analytic_dirac}, we can also derive its trace in the form of a weighted norm sum, as highlighted in the following proposition:
\begin{table}[t]
\centering
\resizebox{0.5\textwidth}{!}{%
\begin{tabular}{c|cccccc}
\Xhline{1pt}
{\multirow{2}{*}{Methods}} & \multicolumn{6}{c}{The relative error of NLL evaluation} \\
\Xcline{2-7}{0.5pt}
& \makecell{t = 1.0} & \makecell{t = 0.8} & \makecell{t = 0.6} & \makecell{t = 0.4} & \makecell{t = 0.2} & \makecell{t = 0.0} \\
\Xhline{1pt}
VJP (eq. \ref{nll_VJP}) & 6.68\% & 5.79\% & 10.46\% & 20.13\% & 51.14\% & 70.95\% \\
DF-TM \textbf{(Ours)} & \textbf{3.41\%} & \textbf{4.56\%} & \textbf{4.13\%} & \textbf{4.28\%} & \textbf{5.33\%} & \textbf{5.81\%} \\
\Xhline{1pt}
\end{tabular}
}
\caption{\small Comparison of the VJP method and our DF-TM in terms of the
diffusion Fisher trace evaluation error across different timesteps. The error is evaluated on the 2-D chessboard data with the VE schedule.
}
\label{tab:trace_toy}
\vspace{-3mm}
\end{table}
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{proposition}\label{prop:trace_matching}
In the same context as Proposition \ref{prop:analytic_dirac}, the trace of the diffusion Fisher matrix for the diffused distribution $q_t$, where $t\in(0,1]$, is given by:
\begin{equation} \label{trace_matching}
\small
\begin{aligned}
\mathrm{tr}\left(\mF_t(\vx_t,t)\right) = \frac{d}{\sigma_t^2} - \frac{\alpha_t^2}{\sigma_t^4}\left[ \sum_i w_i \norm{\vy_i}^2 - \norm{\sum_i w_i \vy_i}^2 \right]
\end{aligned}
\end{equation}
\end{proposition}
\end{tcolorbox}
As demonstrated in Proposition \ref{prop:optimal_y}, the $\norm{\sum_i w_i \vy_i}^2$ can be directly estimated by $\norm{\bar{\vy}_\theta(\vx_t, t)}^2$. Therefore, the only unknown element in \eqref{trace_matching} is $\sum_i w_i \norm{\vy_i}^2$.
Consequently, we suggest learning this term using a scalar-valued network, as per the following training scheme:
\begin{algorithm}[h]
\caption{Training of DF-TM Network}\label{algo:tnet}
\small
% \textbf{Input}: number of timesteps $N$, noise scheduler $\{\alpha_t\}$ and $\{\sigma_t\}$.
\begin{algorithmic}[1] %[1] enables line numbers
\STATE \textbf{Input}: data space dimension $d$, initial network $\vt_\theta(\cdot,\cdot):\sR^d \times \sR\mapsto \sR$, noise schedule $\{\alpha_t\}$ and $\{\sigma_t\}$.
\REPEAT
\STATE $\vx_0 \sim q_0\left(\vx_0\right)$
\STATE $t \sim \operatorname{Uniform}(\{1, \ldots, T\})$
\STATE $\vepsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
\STATE $\vx_t = \alpha_t \vx_0 + \sigma_t \vepsilon$
\STATE Take gradient descent on $\nabla_\theta \abs{\vt_\theta(\vx_t,t) - \frac{\norm{\vx_0}^2}{d} }^2$
\UNTIL{converged}
\STATE \textbf{Output}: $\vt_\theta(\cdot,\cdot) $
% \approx \frac{1}{d}\mathrm{tr}\left(\sum_i w_i \norm{\vy_i}^2\right)$
\end{algorithmic}
\end{algorithm}
The training scheme detailed in Algorithm \ref{algo:tnet} can indeed enable $\vt_\theta(\vx_t,t)$ to estimate the weighted norm term $\frac{1}{d}\sum_i w_i(\vx_t,t) \norm{\vy_i}^2$. This is substantiated by the convergence analysis Proposition \ref{prop:optimal_t}, as presented below.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{proposition}\label{prop:optimal_t}
$\forall (x_t, t) \in \mathbb{R}^d \times \mathbb{R}_{\ge 0}$, the optimal $t_\theta(\vx_t, t)$s trained by the objective in Algorithm \ref{algo:tnet} are equal to $\frac{1}{d}\sum_i w_i(\vx_t,t) \norm{\vy_i}^2$.
\end{proposition}
\end{tcolorbox}
Once we have obtained $\vt_\theta$, we can evaluate the trace of diffusion Fisher efficiently, as illustrated below. This approach is a straightforward result of \eqref{trace_matching} and Propositions \ref{prop:optimal_y} and \ref{prop:optimal_t}, which we refer to as DF trace matching (DF-TM).
\begin{equation}\label{trace_appro}
% \small
\begin{aligned}
% (\text{DF-TM})~~~~~~~~~~
\mathrm{tr}\left(\mF_t(\vx_t,t)\right) \approx \frac{d}{\sigma_t^2} - \frac{\alpha_t^2}{\sigma_t^4}\left( d*t_\theta(\vx_t,t) - \norm{\vy_\theta(\vx_t,t)}^2 \right)
\end{aligned}
\end{equation}
To estimate the trace using DF-TM in \eqref{trace_appro}, we simply need one access to $\vt_\theta$ and $\vepsilon_\theta$.
DF-TM enables us to effectively evaluate the log-likelihood in a gradient-free manner with linear time complexity.
We can further substantiate the theoretical approximation error bound when using DF-TM to calculate the trace of the diffusion Fisher, as illustrated in the Proposition \ref{prop:error_bound_trace}.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{proposition}\label{prop:error_bound_trace}
Assume the approximation error on $t_\theta(\vx_t,t)$ is $\delta_1$ and on $\vepsilon_\theta(\vx_t,t)$ is $\delta_2$, then the approximation error of the approximated Fisher trace in \eqref{trace_appro} is at most $\frac{\alpha_t^2}{\sigma_t^4}\delta_1 +\frac{1}{\sigma_t^2}\delta_2^2$.
\end{proposition}
\end{tcolorbox}
\paragraph{Experiments}
We conduct toy experiments to show the accuracy and commercial-level experiments to show the efficiency of our DF-TM method.
In toy experiments, as shown in Table \ref{tab:trace_toy}, we trained a simple DF-TM network on 2-D chessboard data and then evaluated the relative error of the trace estimated by VJP and DF-TM.
It is shown that our DF-TM method consistently achieves more accurate trace estimation of DF compared to the VJP method.
In commercial-level experiments, we trained two DF-TM networks for the SD-1.5 and SD-2base pipeline on the Laion2B-en dataset \citep{schuhmann2022laionb}, which contains 2.32 billion text-image pairs.
Our $t_\theta$ follows the U-net structure, similar to stable diffusion models, but with an added MLP head to produce a scalar-valued output.
We utilize the AdamW optimizer \citep{loshchilov2018decoupled} with a learning rate of 1e-4.
The training is executed across 8 Ascend 910B chips with a batch size of 384 and completed after 150K steps.
In Figure \ref{fig:train_loss}, we demonstrate that the training loss of DF-TM nets converges smoothly, indicating the robustness of the DF-TM training scheme in Algorithm \ref{algo:tnet}.
\begin{figure*}[!t]
\centering
\includegraphics[width=0.94\textwidth]{figures/figure_icml_adjoint.pdf}
\caption{Visual comparison of DF-EA \textbf{(Ours)} and VJP in the adjoint improvement task on (left) SAC aesthetic score and (right) Pick-Score. DF-EA consistently generates images with better visual effects and reduced time expenditure.
% The smoother visual effect is a desirable outcome of the target score, but not a result of our DF-EA.
}
\vspace{-4mm}
\label{fig:aes_shown}
\end{figure*}
In Figure \ref{fig:tradeoff}, we evaluate the average NLL and Clip score of samples generated by various SD models, using 10k randomly selected prompts from the COCO dataset \cite{lin2014microsoft}. A lower NLL suggests more realistic data generation, while a higher Clip score indicates a better match between the generated images and the input prompts. The results imply that the NLL and the Clip score form a trade-off curve across different guidance scales. This phenomenon, previously hypothesized in theory \citep{wu2024theoretical}, is now confirmed in the SD models, thanks to the effective NLL evaluation via the DF-TM method.
In Figure \ref{fig:nll_compare}, we display images with varying NLL under the same prompt with 10 steps on SD-1.5 using DDIM. It's clear that images with lower NLL exhibit greater visual realism, while those with higher NLL often contain deformed elements (emphasized by the yellow rectangle). Our proposed NLL evaluation method proves to be an effective tool for automatic sample selection.
\section{Diffusion Fisher Endpoint Approximation}\label{sec:ea}
% \vspace{-2mm}
The adjoint optimization of DMs would require access to the matrix-vector multiplication of the diffusion Fisher, which can also be seen as applying diffusion Fisher as a linear operator.
In this section, we present a training-free method that simplifies the complex linear transformation calculations, thus enabling faster and more accurate adjoint optimization based on the outer-product form diffusion Fisher.
\paragraph{Adjoint optimization sampling.}
Guided sampling techniques are extensively utilized in diffusion models to facilitate controllable generation. Recently, to address the inflexibility of commonly used classifier-based guidance \citep{dhariwal2021diffusion} and classifier-free guidance \citep{ho2022classifier}, a series of adjoint guidance methods have been explored \citep{pan2023adjointdpm,pan2023towards}.
Consider optimizing a scalar-valued loss function $\gL(\cdot):\sR^d\mapsto\sR$, which takes $\vx_0$ in the data space as input.
Adjoint guidance is implemented by applying gradient descent on $\vx_t$ in the direction of $\frac{\partial \gL(\vx_0(\vx_t))}{\partial \vx_t}$. The essence of adjoint guidance is to use the gradient at $t=0$ and follow the adjoint ODE \citep{pollini2018adjoint,chen2018neural} to compute $\vlambda_t := \frac{\partial \gL(\vx_0(\vx_t))}{\partial \vx_t}$ for $t>0$.
\begin{equation} \label{Adjoint_ODE}
% \small
\begin{aligned}
% (\text{Adjoint ODE})~~~~~~~~~~~~
\frac{\rd\vlambda_t}{\rd t} = - \vlambda_t^{\top}\frac{\partial \boldsymbol{h}_\theta\left(\vx_t, t\right)}{\partial \vx_t},\quad \vlambda_0 = \frac{\partial\mathcal{L}(\vx_0)}{\partial \vx_0}
\end{aligned}
\end{equation}
\paragraph{Adjoint ODE via VJP.}
Regardless of the ODE solver being used, it is necessary to compute the right-hand side of \eqref{Adjoint_ODE}, or equivalently, $\mF(\vx_t,t)^{\top}\vlambda_t$. This computation can be interpreted as applying the diffusion Fisher matrix as a linear operator to the adjoint state $\vlambda_t$, from a functional analysis perspective \citep{yosida2012functional}.
Current practices utilize the VJP technique to approximate this linear transformation operation as follows:
\begin{equation} \label{Adjoint_VJP}
% \small
\begin{aligned}
% (\text{VJP for Adjoint})~~~~~~~~~~~~
\mF(\vx_t,t)^{\top}\vlambda_t&\approx\frac{1}{\sigma_t}\frac{\partial \vepsilon_\theta\left(\vx_t, t\right)}{\partial \vx_t}^{\top}\vlambda_t \\
&\approx\frac{1}{\sigma_t} \frac{\partial \left[\ip{\vepsilon_\theta\left(\vx_t, t\right)}{\vlambda_t}\right]}{\partial \vx_t}
\end{aligned}
\end{equation}
% In practical code implementation, the VJP as per Equation \eqref{Adjoint_VJP} is computed by first calculating the inner product of $\vepsilon_\theta\left(\vx_t, t\right)$ and $\vlambda_t$, followed by computing the gradient of this inner product with respect to $\vx_t$.
This process involves computationally expensive neural network auto-differentiations, and the approximation errors introduced by the VJP technique have no theoretical bound.
\paragraph{Adjoint ODE via DF-EA.}
As previously discussed in Section \ref{sec:analytical}, the DF inherently doesn't require gradients, suggesting that we could potentially apply the diffusion Fisher as a linear operator in a gradient-free manner.
The challenging part in \eqref{information_dirac} is $\sum_i w_i \vy_i\vy_i^{\top}$, which represents a weighted form of outer-products of data.
Based on the definition of $w_i$, the closest $\vy_i$ to $\vx_0$ will dominate as $t\to0$. This makes it intuitive to replace this sum with a single final sample outer-product $\vx_0\vx_0^{\top}$.
It's also important to note that the adjoint guidance itself needs to compute $\vx_0$ at each guidance step, eliminating the need for additional computation to obtain $\vx_0$.
Given that we utilize the endpoint sample $\vx_0$, we refer to this approximation technique as DF Endpoint Approximation (EA). The formulation for DF-EA in the adjoint ODE is as follows:
\begin{equation} \label{EA_approx}
\small
\begin{aligned}
% (\text{DF-EA})~~~~~~
&\mF(\vx_t,t)^{\top}\vlambda_t \\
\approx& \left(\frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left( \sum_i w_i \vy_i\vy_i^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)\right)^{\top}\vlambda_t\\
\approx& \left(\frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left( \vx_0\vx_0^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)\right)^{\top}\vlambda_t\\
=& \frac{1}{\sigma_t^2}\vlambda_t - \frac{\alpha_t^2}{\sigma_t^4} \left< \vx_0,\vlambda_t\right>\vx_0 + \frac{\alpha_t^2}{\sigma_t^4} \left< \bar{\vy}_\theta(\vx_t, t),\vlambda_t\right>\bar{\vy}_\theta(\vx_t, t)
\end{aligned}
\end{equation}
The DF-EA approximation leads to a scalar-weighted combination of $\vlambda_t$, $\vx_0$, and $\bar{\vy}_\theta(\vx_t, t)$, which importantly, does not involve any gradients.
Additionally, we derive the theoretical approximation error bound of the DF-EA in Proposition \ref{prop:error_bound_ea}. To measure the accuracy of DF-EA as a linear operator, we opt to use the Hilbert–Schmidt norm \citep{gohberg1990hilbert} for measurement, as follows:
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{proposition}\label{prop:error_bound_ea}
Assume that the approximation error on $\vepsilon_\theta(\vx_t,t)$ is $\delta_2$, the approximation error of the DF-EA linear operator, as referenced in \ref{EA_approx}, is at most $\frac{\alpha_t^2}{\sigma_t^3}\left(2\mathcal{D}_y^2+ \sqrt{d}\delta_2\right)$ when measured in terms of the Hilbert–Schmidt norm.
\end{proposition}
\end{tcolorbox}
\paragraph{Experiments on DF-EA.}
As depicted in Figure \ref{fig:adjoint_score}, we conducted experiments comparing our DF-EA and VJP methods in adjoint guidance sampling, using five different scores and two different base models.
% We examined varying numbers of adjoint guidance, ranging from 0 to 20, under a full inference number of 50. The adjoint guidance scale was grid searched by the VJP method.
DF-EA consistently achieves better scores due to its bounded approximation error. Furthermore, DF-EA requires less processing time as it eliminates the need for time-consuming gradient operations. DF-EA and VJP are compared under the same guidance scales and schemes across various numbers of steps. Details regarding the score function can be found in Appendix \ref{app:extra_adjoint}.
As depicted in Figure \ref{fig:aes_shown}, our DF-EA consistently generates samples with higher scores with a reduced time cost compared to VJP.
% It's worth noting that this enhancement in aesthetics results in final images that are more vibrant and smoother.
All samples are generated within 50 steps, with adjoint applied from the 15\textsuperscript{th} to the 35\textsuperscript{th} step. Details on hyperparameters can be found in Appendix \ref{app:extra_adjoint}.
% \begin{figure}[t!]
% \begin{minipage}{0.49\textwidth}
% \centering
% \includegraphics[width=1.15\textwidth]{ICLR 2025 Template/figures/aes_sac.pdf}
% \captionof{figure}{xx}
% \label{fig:}
% \end{minipage}
% \hfill
% \begin{minipage}{0.49\textwidth}
% \centering
% \includegraphics[width=1.15\textwidth]{ICLR 2025 Template/figures/aes_ava.pdf}
% \captionof{figure}{xx}
% \label{fig:}
% \end{minipage}
% \end{figure}
% \vspace{-4mm}
\section{Numerical OT Verification of DMs}\label{sec:ot}
% \vspace{-4mm}
There is an increasing trend towards analyzing the probability modeling capabilities of DMs by interpreting them from an optimal transport (OT) perspective \citep{albergo2023stochastic,chen2024rethinking}.
The foundational concepts of optimal transport can be found in Appendix \ref{appendix:preliminary_ot}.
One of the central questions is whether the map deduced by the PF-ODE could represent an optimal transport.
\citet{khrulkov2023understanding} have proven that, given single-Gaussian initial data and a VE noise schedule, the PF-ODE deduced map is optimal transport.
\citet{zhang2024formulating} has demonstrated that affine initial data suffices for the PF-ODE map to be optimal transport.
However, the OT property of the general PF-ODE map remains an open question.
% For a specific noise schedule where $f(t)\equiv0$, \cite{zhang2024formulating} demonstrates that having all data points lying on a single line is a sufficient condition for the map to represent an optimal transport.
% In this section, we find out that the AFI, as derived in section \ref{sec:analytical}, can contribute to the first equivalence condition for the OT property of PF-ODE deduced mapping under a general noise schedule. We refer to this theorem as the AFI Optimal Transport (AFI-OT) theorem. The AFI-OT condition is presented in the following:
In this section, we propose the first numerical OT verification experiment for a general PF-ODE deduced map based on the outer-product form diffusion Fisher we obtained in section \ref{sec:analytical}.
We first derive the following corollary for the OT property of the PF-ODE deduced map.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt] \small
\begin{corollary}\label{prop:ot}
Denote the diffeomorphism deduced by the PF-ODE in \eqref{PFODE} as follows
\begin{equation}
\small
\begin{aligned}
T_{s,t}: \mathbb{R}^n \longrightarrow \mathbb{R}^n; \vx_s \longmapsto \vx_t, \quad \forall t \geq s > 0.
\end{aligned}
\end{equation}
% This diffeomorphism is well-posed guaranteed by the global version of Picard-Lindelöf theorem \cite{amann2011ordinary,zhang2024formulating}.
The diffeomorphism $T_{s,T}$ is a Monge optimal transport map \textbf{if and only if} the normalized fundamental matrix for $\mB(t)\equiv \mB(t,\vx_t)$ at $s$ is s.p.d. for every PF-ODE chain that starts from a $\vx_T\in\mathbb{R}^d$.
where
\begin{equation}
\small
\begin{aligned}\label{eq:ot_b}
\mB(t,\vx_t) =& \left[f(t)-\frac{g^2(t)}{2\sigma_t^2}\right]\mI+\frac{\alpha_t^2 g^2(t)}{2\sigma_t^4}\left[ \sum_i w_i \vy_i\vy_i^{\top} \right.
\\ &\left.- \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right].
\end{aligned}
\end{equation}
The definition of the normalized fundamental matrix is deferred to Appendix \ref{appendix:proof_ot}.
\end{corollary}
\end{tcolorbox}
Note that $T_{s,t}$ is well-posed, guaranteed by the global version of the Picard-Lindelöf theorem \cite{amann2011ordinary,zhang2024formulating}.
% The outline of the proof is as follows: We first apply Brenier's theorem \citep{brenier1991polar,santambrogio2015optimal} to convert the problem of whether the PF-ODE mapping is an optimal transport into the task of finding a convex potential, where the existence is guaranteed by the Poincar\'e's Theorem \citep{lang2012fundamentals}. We then use adjoint methods to express the second derivatives of the convex potential function in the form of an integral of information \cite{rockafellar2015convex,pollini2018adjoint}. Finally, we apply matrix exponential integration theory \cite{masuyama2016limit} to reformulate the condition into the AFI-OT theorem.
Detailed proofs can be found in Appendix \ref{appendix:proof_ot}.
We then design the numerical OT verification experiment for a given noise schedule and initial data as shown in the Algorithm \ref{algo:ot}. The detailed version of Algorithm \ref{algo:ot} can be found in Appendix \ref{app:extra_ot}.
\begin{algorithm}[h]
\caption{Numerical OT test for PF-ODE map}\label{algo:ot}
\small
\begin{algorithmic}[1] %[1] enables line numbers
\STATE \textbf{Input}: initial data $\{\vy_i\}_{i=1}^{N}$, noise schedule $\{\alpha_t\}$ and $\{\sigma_t\}$, discretization steps $M$.
\STATE Initialize $\mA_M = \mI$, $\vx_M \sim \mathcal{N}(0,\sigma_T\mI)$.
\FOR{$i=M,M-1,\cdots,1$}
\STATE $\rd t = t_{i-1} - t_{i}$.
\STATE Calculate $\mB_i$ by \eqref{eq:ot_b}.
\STATE $\mA_{i-1} = \mA_{i} + \rd t * \mA_{i}^\top\mB_i $\hfill\COMMENT{solve fundamental matrix.}
\STATE $\vx_{i-1} = \text{PF-ODE Solver}(\vx_i,i)$
\ENDFOR
% \STATE $\phi = \frac{\norm{\mA-\mA^T}_F}{\sqrt{2}\norm{\mA}_F}$
\STATE \textbf{Output}:$\mA_0$. \hfill\COMMENT{The result fundamental matrix.}
\end{algorithmic}
\end{algorithm}
In Table \ref{tab:ot}, we numerically tested the OT property of the PF-ODE map under different nose schedules and initial data.
We tested four commonly used noise schedules, VE \cite{ho2020denoising}, VP \cite{song2019generative}, sub-VP \cite{song2020score}, and EDM \cite{karras2022elucidating} on 2-D data.
We first tested when the initial data is single-Gaussian and affine, we found that the result normalized fundamental matrix is p.s.d., which means the PF-ODE map is optimal transport. These results coincide with the previous proposed theoretical results in \cite{khrulkov2023understanding} and \cite{zhang2024formulating}.
We further tested an extremely simple non-affine case, that is, there are only three initial data points, (0,0), (0,0.5), and (0.5,0). We showed that all noise schedules result in an obvious asymmetric fundamental matrix, which means that the PF-ODE map fails to be OT in these cases.
Here, we use the Frobenius norm of the difference of $\mA$ and its transpose divided by the Frobenius norm of $\mA$ to measure its asymmetric rate.
Based on this finding, we hypothesize that the PF-ODE map can only be OT in the affine initial data (Single-Gaussian is a special case of affine data), but fails to be OT in general non-affine data.
Our hypothesis coincides with the statement that the flow map of a Fokker–Planck equation is not necessarily OT as proposed in \cite{lavenant2022flow}. (Note that PF-ODE is a special class of Fokker–Planck equations)
\begin{table}[t]
\centering
\small
\resizebox{0.49\textwidth}{!}{%
\begin{tabular}{c||cc|cc|cc}
\Xhline{1pt}
Initial Data & \multicolumn{2}{c|}{Single-Gaussian} & \multicolumn{2}{c|}{Affine}& \multicolumn{2}{c}{Non-affine}\\
\Xcline{1 - 7}{0.5pt}
Noise Schedule & Asym. & OT & Asym. & OT& Asym. & OT\\
\Xhline{1pt}
\makecell{VE\\ \citep{song2019generative}} & 0.00\% & \cmark & 0.00\% & \cmark & 25.28\% & \xmark\\
\Xhline{0.5pt}
\makecell{VP\\ \cite{ho2020denoising}} & 0.00\% & \cmark & 0.00\% & \cmark & 23.36\% & \xmark\\
\Xhline{0.5pt}
\makecell{sub-VP\\ \cite{song2020score}} & 0.00\% & \cmark & 0.00\% & \cmark & 13.84\% & \xmark\\
\Xhline{0.5pt}
\makecell{EDM\\ \cite{karras2022elucidating}} & 0.00\% & \cmark & 0.00\% & \cmark & 27.09\% & \xmark\\
\Xhline{1pt}
\end{tabular}
}
\caption{\small Comparison of numerical OT verification results of four commonly used noise schedulers with different initial data.}
\label{tab:ot}
\vspace{-5mm}
\end{table}
\section{Conclusions}
\vspace{-2mm}
% We tackle the inexact inversion issue of DMs in a training-free manner.
% We introduce the generic Bidirectional Explicit Linear Multi-step (BELM) framework based on a linear multi-step observation, which encompasses existing heuristic exact inversion samplers as special cases.
% Furthermore, we devise a Local Truncation Error (LTE) minimization approach to construct the Optimal-BELM (O-BELM) within the BELM framework, which achieves a higher order of local error.
% We provide a theoretical guarantee of global stability and convergence for O-BELM and conduct various experiments to demonstrate that O-BELM not only accomplishes exact inversion but also maintains a high-quality sampling capability.
% Please refer to further discussion and limitations in appendix \ref{app:discussion}.
This paper derives an outer-product span formulation of the diffusion Fisher and introduces two approximation algorithms for different types of access to the diffusion Fisher: diffusion Fisher trace matching (DF-TM) and diffusion Fisher endpoint approximation (DF-EA).
Both methods are theoretically guaranteed in terms of approximation error bounds, and offer improved accuracy and reduced time-cost compared to the traditional VJP method.
% Theoretically, we have established the first general theorem for the PF-ODE map to be optimal transport.
We also designed the numerical OT property verification experiment for the PF-ODE map based on the outer-product form of DF.
This work not only improves the efficiency of diffusion Fisher access but also widens our understanding of the diffusion models.
Please refer to further discussions in appendix \ref{app:discussion}.
The code is available at \url{https://github.com/zituitui/DiffusionFisher}.
\section{Proofs and Formulations}
% \subsection{The diffusion Fisher outer-product formulation in general setting}\label{app:general_df}
\subsection{Proof of Proposition \ref{prop:analytic_dirac}}
Notice that, in the subsection, we can do an interchange of sum and gradient; this is due to Leibniz's rule \citep{osler1970leibniz} and the boundness condition we set in \eqref{dirac_initial_dist}.
Before we give the proof of Proposition \ref{prop:analytic_dirac}, we would like to establish two technical lemmas. The first lemma is about the first partial derivative of $v_i(\vx_t,t)$ with respect to $\vx_t$.
\begin{lemma}
\begin{equation} \label{lemma:v}
\begin{aligned}
\frac{\partial v_i(\vx_t,t)}{\partial \vx_t} &=\frac{\partial\exp(-\frac{|\vx_t - \alpha_t \vy_i|^2}{2\sigma_t^2})}{\partial \vx_t}\\
& = -\frac{1}{\sigma_t^2}(\vx_t - \alpha_t \vy_i)\exp(-\frac{|\vx_t - \alpha_t \vy_i|^2}{2\sigma_t^2})\\
& = -\frac{1}{\sigma_t^2}(\vx_t - \alpha_t \vy_i)v_i(\vx_t,t)
\end{aligned}
\end{equation}
\end{lemma}
The second lemma is about the Jacobian of $\sum_i w_i\left(\vx_t, t\right) \vy_i$ w.r.t. $\vx_t$.
\begin{lemma}\label{lemma:chain}
\begin{equation}
\begin{aligned}
&\frac{\partial \sum_i w_i\left(\vx_t, t\right) \vy_i}{\partial \vx_t}\\
&=\sum_i \vy_i\left(\frac{\partial w_i\left(\vx_t, t\right)}{\partial \vx_t}\right)^{\top} \\
& =\sum_i \vy_i\left(\frac{\partial}{\partial \vx_t}\left[\frac{v_i\left(\vx_t, t\right)}{\sum_j v_j\left(\vx_t, t\right)}\right]\right)^{\top} \\
& =\sum_i \vy_i\left\{\frac{\frac{\partial v_i\left(\vx_t, t\right)}{\partial \vx_t} \left[\sum_k v_k\left(\vx_t, t\right)\right]-v_i\left(\vx_t, t\right) \frac{\partial}{\partial \vx_t}\left[\sum_j v_j\left(\vx_t, t\right)\right]}{\left[\sum_k v_k\left(\vx_t, t\right)\right]^2}\right\}^{\top} \\
& =\frac{1}{\left(\sum_k v_k\right)^2} \sum_i \vy_i\left\{-\frac{1}{\sigma_t^2}\left(\vx_t-\alpha_t \vy_i\right) v_i \sum_k v_k-v_i\left(\vx_t,t\right)\left(-\frac{1}{\sigma_t^2}\right) \sum_j\left(\vx_t-\alpha_t \vy_j\right) v_j\right\}^{\top} \\
& =\frac{1}{\left(\sum_k v_k\right)^2}\left(-\frac{1}{\sigma_t^2}\right) \sum_i v_i \vy_i\left\{\left(\vx_t-\alpha_t \vy_i\right) \sum_k v_k-\sum_j\left(\vx_t-\alpha_t \vy_i\right) v_j\right\}^{\top} \\
& =\frac{1}{\left(\sum_k v_k\right)^2}\left(-\frac{1}{\sigma_t^2}\right) \sum_i v_i \vy_i\left\{-\alpha_t \vy_i \sum_k v_k+\alpha_t \sum_j \vy_j v_j\right\}^{\top} \\
& =\frac{\alpha_t}{\sigma_t^2}\left[\frac{\sum_i v_i\vy_i\vy_i^{\top}\bcancel{\left(\sum_k v_k\right)}}{\left(\sum_k v_k\right)^{\bcancel{2}}} - \left(\frac{\sum_i v_i\vy_i}{\sum_k v_k}\right)\left(\frac{\sum_i v_i\vy_i}{\sum_k v_k}\right)^{\top}\right]\\
& =\frac{\alpha_t}{\sigma_t^2}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right]
\end{aligned}
\end{equation}
\end{lemma}
Now we are ready to give the Proof of Proposition \ref{prop:analytic_dirac}
% $$
% \begin{aligned}
% \varepsilon_\theta & =\frac{1}{\sigma_t} x_t-\frac{\alpha_t}{\sigma_t} \sum_j w_j\left(\vx_t, t\right) y_j \\
% \frac{\partial \varepsilon_\theta}{\partial x_t} & =\frac{1}{\sigma_t}-\frac{\alpha_t}{\sigma_t}\left[\frac{\alpha_t}{\sigma_t^2}\left(\sum_i w_i \vy_i \vy_i^{\top}-\bar{y} \bar{y}^{\top}\right)\right] \\
% & =\frac{1}{\sigma_t} \frac{1}{\sigma_t}-\frac{\alpha_t^2}{\sigma^3}\left(\sum_i w_i y_T \vy_i^{\top}-\bar{y} \bar{y}^{\top}\right) \\
% \mathcal{L}\left(\vx_t, t\right) & =-\nabla_{x t}\left(\nabla_{x t} \log p_t\right)=-\nabla_{x t}\left(-\frac{\varepsilon_\theta}{\sigma_t}\right)=\frac{1}{\sigma_t} \frac{\partial \varepsilon_\theta}{\partial x_t}
% \end{aligned}
% $$
\begin{proof}
According to the initial distribution (\eqref{dirac_initial_dist}) and the diffusion kernel (\eqref{diffusion_kernel}), the marginal distribution at some time $t>0$ would be
\begin{equation}
\begin{aligned}
p\left(\vx_t, t\right) & =\frac{1}{N} \sum_i\left(2 \pi \sigma_t^2\right)^{-\frac{d}{2}} \exp \left(-\frac{\left|x_t-\alpha_t \vy_i\right|^2}{2 \sigma_{t^2}}\right)
\end{aligned}
\end{equation}
Thus, the log-density has the following analytical formulation
\begin{equation}
\begin{aligned}
\log p(\vx_t, t) & =\log \left[\frac{1}{N}\left(2 \pi \sigma_t^2\right)^{-\frac{d}{2}} \sum_i \exp \left(-\frac{\left|x_t-\alpha_t \vy_i\right|^2}{2 \sigma_{t^2}}\right)\right] \\
& =\log \left[\sum_i \exp \left(-\frac{\left|x_t-\alpha_t \vy_i\right|^2}{2 \sigma_t^2}\right)\right]+C \\
& =\log \left[\sum_i v_i\left(\vx_t, t\right)\right]+C
\end{aligned}
\end{equation}
The score can be expressed as follows
\begin{equation}
\begin{aligned}
\frac{\partial}{\partial \vx_t} \log p\left(\vx_t, t\right) & =\frac{\partial}{\partial \vx_t} \log \left[\sum_i v_i\left(\vx_t, t\right)\right] \\
& =\frac{\frac{\partial}{\partial \vx_t}\left[\sum_i v_i\left(\vx_t, t\right)\right]}{\sum_i v_i\left(\vx_t, t\right)} \\
& =\frac{-\frac{1}{\sigma_t^2} \sum_j\left(\vx_t-\alpha_t y_j\right) v_j}{\sum_i v_i\left(\vx_t, t\right)} \\
& =-\frac{1}{\sigma_t^2}\left[x_t-\alpha_t \sum_j w_j\left(\vx_t, t\right) y_j\right]
\end{aligned}
\end{equation}
The Fisher information we want can then be calculated by further applying a gradient on the score.
\begin{equation}
\begin{aligned}
\mF_t(\vx_t,t) &= - \frac{\partial}{\partial \vx_t}\left(\frac{\partial}{\partial \vx_t} \log p\left(\vx_t, t\right) \right)\\
& = - \frac{\partial}{\partial \vx_t}\left\{-\frac{1}{\sigma_t^2}\left[x_t-\alpha_t \sum_j w_j\left(\vx_t, t\right) y_j\right]\right\}\\
& = \frac{1}{\sigma_t^2}\mI - \frac{\alpha_t}{\sigma_t^2}\frac{\partial \sum_i w_i\left(\vx_t, t\right) \vy_i}{\partial \vx_t}\\
&= \frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right]~~~(\text{by Lemma \ref{lemma:chain}})
\end{aligned}
\end{equation}
\end{proof}
This proof is inherently the calculation of the Hessian of a log-convolution of a density, we provide the detailed derivation here for completeness.
\subsection{Proof of Proposition \ref{prop:optimal_y}}
\begin{proof}
Given fixed $(x_{t}, t)$, $\mathcal{L}$ is a quadratic form of $y_\theta$. To obtain the optimal $\bar{\vy}_\theta$, we differentiate $\mathcal{L}$ and set this derivative equal to zero, resulting in the following
\begin{align}
0 = \frac{\partial \mathcal{L}}{\partial \bar{\vy}_\theta(\vx_t, t)} &= \frac{\partial}{\partial \bar{\vy}_\theta(\vx_t, t)} \sum_j \underbrace{\frac{1}{N} (2\pi \sigma_{t}^2)^{-\frac{d}{2}}}_{A_{t}} v_j(\vx_{t}, t) \lambda_{t} \frac{\alpha_{t}^2}{\sigma_{t}^2} \norm{\bar{\vy}_\theta(\vx_t, t) - \vy_j}^2 \notag \\
&=2 A_{t} \lambda_{t} \frac{\alpha_{t}^2}{\sigma_{t}^2} \sum_j \vv_j(x_{t}, t) (\bar{\vy}_\theta(\vx_t, t) - \vy_j),
\end{align}
which yields
\begin{equation}
\bar{\vy}^*_\theta(\vx_t, t) = \sum_k \frac{v_k(\vx_{t}, t)}{\sum_j v_j(\vx_{t}, t)} \vy_k = \sum_i w_i \vy_i.
\end{equation}
\end{proof}
% \subsection{Proof of Corollary \ref{corollary:theta_dirac}}
% \begin{equation} \label{}
% \begin{aligned}
% \frac{\partial \vepsilon_\theta(\vx_t,t)}{\partial\vx_t} &\approx \frac{\partial \left(-\sigma(t)\nabla_{\vx(t)} \log p(\vx(t), t)\right)}{\partial\vx_t}\\
% &=\sigma(t)\mF_t(\vx_t,t)\\
% &=\frac{1}{\sigma_t}\mI - \frac{\alpha_t^2}{\sigma_t^3}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right]\\
% &\approx \frac{1}{\sigma_t}\mI - \frac{\alpha_t^2}{\sigma_t^3}\left( \sum_i w_i \vy_i\vy_i^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)
% \end{aligned}
% \end{equation}
\subsection{Proof of Proposition \ref{prop:analytic_general}}
Notice that, in the subsection, we can do an interchange of integral and gradient; this is due to Leibniz's rule \citep{osler1970leibniz} and the bounded moments condition we set in \eqref{general_initial_dist}.
Before we give the proof of Proposition \ref{prop:analytic_general}, we would like to establish two technical lemmas. The first lemma is about the first partial derivative of $v(\vx_t,t,\vy)$ w.r.t. $\vx_t$.
\begin{lemma}
\begin{equation} \label{lemma:v_general}
\begin{aligned}
\frac{\partial v(\vx_t,t, \vy)}{\partial \vx_t} &=\frac{\partial\exp(-\frac{|\vx_t - \alpha_t \vy|^2}{2\sigma_t^2})}{\partial \vx_t}\\
& = -\frac{1}{\sigma_t^2}(\vx_t - \alpha_t \vy)\exp(-\frac{|\vx_t - \alpha_t \vy|^2}{2\sigma_t^2})\\
& = -\frac{1}{\sigma_t^2}(\vx_t - \alpha_t \vy)v(\vx_t,t,\vy)
\end{aligned}
\end{equation}
\end{lemma}
\begin{lemma}\label{lemma:chain_general}
\begin{equation}
\small
\begin{aligned}
&\frac{\partial \int_{\sR^d} w(\vx_t, t, \vy) \vy \rd q_0(\vy)}{\partial \vx_t}\\
&=\int_{\sR^d} \vy\left(\frac{\partial w\left(\vx_t, t, \vy\right)}{\partial \vx_t}\right)^{\top} \rd q_0(\vy) \\
& =\int_{\sR^d} \vy\left(\frac{\partial}{\partial \vx_t}\left[\frac{v(\vx_t, t, \vy)}{\int_{\sR^d} v(\vx_t, t, \vy')\rd q_0(\vy')}\right]\right)^{\top}\rd q_0(\vy) \\
& =\int_{\sR^d} \vy\left\{\frac{\frac{\partial v_i\left(\vx_t, t\right)}{\partial \vx_t} \left[\int_{\sR^d} v(\vx_t, t, \vy')\rd q_0(\vy')\right]-v\left(\vx_t, t, \vy \right) \frac{\partial}{\partial \vx_t}\left[\int_{\sR^d} v(\vx_t, t, \vy')\rd q_0(\vy')\right]}{\left[\int_{\sR^d} v(\vx_t, t, \vy'')\rd q_0(\vy'')\right]^2}\right\}^{\top} \rd q_0(\vy)\\
& =\frac{1}{\left[\int v(\vy'')\rd q_0(\vy'') \right]^2} \int \vy\left\{-\frac{\left(\vx_t-\alpha_t \vy\right) v(\vy)}{\sigma_t^2} \int v(\vy')\rd q_0(\vy') -v(\vy)\left(-\frac{1}{\sigma_t^2}\right) \int\left(\vx_t-\alpha_t \vy'\right) v(\vy')\rd q_0(\vy')\right\}^{\top}\rd q_0(\vy) \\
& =\frac{1}{\left[\int v(\vy'')\rd q_0(\vy'') \right]^2}\left(-\frac{1}{\sigma_t^2}\right) \int v(\vy) \vy\left\{\left(\vx_t-\alpha_t \vy\right) \int v(\vy')\rd q_0(\vy') -\int \left(\vx_t-\alpha_t \vy'\right) v(\vy')\rd\vy'\right\}^{\top} \rd q_0(\vy)\\
& =\frac{1}{\left[\int v(\vy'')\rd q_0(\vy'') \right]^2}\left(-\frac{1}{\sigma_t^2}\right) \int v(\vy) \vy\left\{-\alpha_t \vy \int v(\vy')\rd q_0(\vy') +\alpha_t \int v(\vy')\vy'\rd\vy' \right\}^{\top} \rd q_0(\vy)\\
& =\frac{\alpha_t}{\sigma_t^2}\left[\frac{\int v(\vy)\vy\vy^{\top}\rd q_0(\vy')\bcancel{\left[\int v(\vy')\rd q_0(\vy') \right]}}{\left[\int v(\vy'')\rd q_0(\vy'') \right]^{\bcancel{2}}} - \left(\frac{\int v(\vy)\vy\rd q_0(\vy)}{\int v(\vy'')\rd q_0(\vy'') }\right)\left(\frac{\int v(\vy')\vy'\rd q_0(\vy')}{\int v(\vy'')\rd q_0(\vy'') }\right)^{\top}\right]\\
& =\frac{\alpha_t}{\sigma_t^2}\left[ \int w(\vy) \vy\vy^{\top}\rd q_0(\vy) - \left(\int w(\vy) \vy \rd q_0(\vy)\right)\left(\int w(\vy) \vy \rd q_0(\vy) \right)^{\top} \right]
\end{aligned}
\end{equation}
\end{lemma}
\begin{proof}
According to the initial distribution (\eqref{general_initial_dist}) and the diffusion kernel (\eqref{diffusion_kernel}), the marginal distribution at some time $t>0$ would be
\begin{equation}
\begin{aligned}
p\left(\vx_t, t\right) & =\int_{\sR^d} \left(2 \pi \sigma_t^2\right)^{-\frac{d}{2}} \exp \left(-\frac{\left|x_t-\alpha_t \vy\right|^2}{2 \sigma_{t^2}}\right) \rd q_0(\vy)
\end{aligned}
\end{equation}
Thus the log-density has the following analytical formulation
\begin{equation}
\begin{aligned}
\log q_t(\vx_t, t) & =\log \left[\int_{\sR^d} \left(2 \pi \sigma_t^2\right)^{-\frac{d}{2}} \exp \left(-\frac{\left|x_t-\alpha_t \vy\right|^2}{2 \sigma_{t^2}}\right) \rd q_0(\vy)\right] \\
& =\log \left[\int_{\sR^d} \exp \left(-\frac{\left|x_t-\alpha_t \vy\right|^2}{2 \sigma_{t^2}}\right) \rd q_0(\vy)\right]+C \\
& =\log \left[\int_{\sR^d} v\left(\vx_t, t, \vy\right) \rd q_0(\vy)\right]+C
\end{aligned}
\end{equation}
The score can be expressed as follows
\begin{equation}
\begin{aligned}
\frac{\partial}{\partial \vx_t} \log p\left(\vx_t, t\right) & =\frac{\partial}{\partial \vx_t} \log\left[\int_{\sR^d} v\left(\vx_t, t, \vy\right) \rd q_0(\vy)\right] \\
& =\frac{\frac{\partial}{\partial \vx_t}\left[\int_{\sR^d} v\left(\vx_t, t, \vy\right) \rd q_0(\vy)\right]}{\int_{\sR^d} v\left(\vx_t, t, \vy\right) \rd q_0(\vy)} \\
& =\frac{\int_{\sR^d}\frac{\partial}{\partial \vx_t}\left[ v\left(\vx_t, t, \vy\right) \right] \rd q_0(\vy)}{\int_{\sR^d} v\left(\vx_t, t, \vy\right) \rd q_0(\vy)} ~~~(\text{by xx and Leibniz integral rule})\\
& =\frac{-\frac{1}{\sigma_t^2} \int_{\sR^d} (\vx_t - \alpha_t \vy)v\left(\vx_t, t, \vy\right) \rd q_0(\vy)}{\int_{\sR^d} v\left(\vx_t, t, \vy\right) \rd q_0(\vy)} \\
& =-\frac{1}{\sigma_t^2}\left[x_t-\alpha_t \int_{\sR^d} w(\vx_t, t, \vy) \vy \rd q_0(\vy)\right]
\end{aligned}
\end{equation}
The Fisher information we want can then be calculated by further applying a gradient on the score.
\begin{equation}
\begin{aligned}
\mF_t(\vx_t,t) &= - \frac{\partial}{\partial \vx_t}\left(\frac{\partial}{\partial \vx_t} \log p\left(\vx_t, t\right) \right)\\
& = - \frac{\partial}{\partial \vx_t}\left\{-\frac{1}{\sigma_t^2}\left[x_t-\alpha_t \int_{\sR^d} w(\vx_t, t, \vy) \vy \rd q_0(\vy)\right]\right\}\\
& = \frac{1}{\sigma_t^2}\mI - \frac{\alpha_t}{\sigma_t^2}\frac{\partial \int_{\sR^d} w(\vx_t, t, \vy) \vy \rd q_0(\vy)}{\partial \vx_t}\\
&= \frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left[ \int w_i \vy\vy^{\top}\rd q_0(\vy) - \left(\int w_i \vy \rd q_0(\vy)\right)\left(\int w_i \vy \rd q_0(\vy)\right)^{\top} \right]~~~(\text{by Lemma \ref{lemma:chain_general}})
\end{aligned}
\end{equation}
\end{proof}
This proof is also inherently the calculation of the Hessian of a log-convolution of a density, we provide the detailed derivation here for completeness.
\subsection{Proof of Proposition \ref{prop:optimal_y_general}}
\begin{proof}
Given fixed $(x_{t}, t)$, $\mathcal{L}$ is a quadratic form of $y_\theta$. To obtain the optimal $\bar{\vy}_\theta$, we differentiate $\mathcal{L}$ and set this derivative equal to zero, resulting in the following
\begin{align}
0 = \frac{\partial \mathcal{L}}{\partial \bar{\vy}_\theta(\vx_t, t)} &= \frac{\partial}{\partial \bar{\vy}_\theta(\vx_t, t)} \int_{\sR^d} \underbrace{\frac{1}{N} (2\pi \sigma_{t}^2)^{-\frac{d}{2}}}_{A_{t}} v(\vx_{t}, t, \vy) \lambda_{t} \frac{\alpha_{t}^2}{\sigma_{t}^2} \norm{\bar{\vy}_\theta(\vx_t, t) - \vy}^2\rd q_0(\vy) \notag \\
&=2 A_{t} \lambda_{t} \frac{\alpha_{t}^2}{\sigma_{t}^2} \int_{\sR^d} \vv(x_{t}, t,\vy) (\bar{\vy}_\theta(\vx_t, t) - \vy)\rd q_0(\vy),
\end{align}
which yields
\begin{equation}
\bar{\vy}^*_\theta(\vx_t, t) = \int_{\sR^d} \frac{v(\vx_t, t, \vy')}{\int_{\sR^d} v(\vx_t, t, \vy'')\rd q_0(\vy'')} \vy'\rd q_0(\vy') = \int_{\sR^d} w(\vx_t, t, \vy') \vy'\rd q_0(\vy').
\end{equation}
\end{proof}
\subsection{Proof of Proposition \ref{prop:trace_matching}}
\begin{lemma}
Given a vector $\vv\in\sR^d$, the trace of the outer-product matrix of this vector is precisely equal to the square of its 2-norm. This can be shown as follows:
\begin{equation}
\begin{aligned}
\mathrm{tr}\left(\vv\vv^T\right)&=\sum_{i=1}^d\left(\vv\vv^T\right)_{i,i}=\sum_{i=1}^d v_i * v_i = \norm{\vv}^2
\end{aligned}
\end{equation}
\end{lemma}
\begin{proof}
Then we can start to give the derivation of Proposition \ref{prop:trace_matching}
\begin{equation}
\begin{aligned}
\mathrm{tr}\left(\mF_t(\vx_t,t)\right) &=\mathrm{tr}\left(\frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right]\right)\\
=&\frac{1}{\sigma_t^2}\mathrm{tr}\left(\mI\right) - \frac{\alpha_t^2}{\sigma_t^4}\left[\sum_i w_i\mathrm{tr}\left(\vy_i\vy_i^{\top}\right)-\mathrm{tr}\left(\left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right) \right] \\
=&\frac{d}{\sigma_t^2} - \frac{\alpha_t^2}{\sigma_t^4}\left[ \sum_i w_i \norm{\vy_i}^2 - \norm{\sum_i w_i \vy_i}^2 \right]
\end{aligned}
\end{equation}
\end{proof}
\subsection{Proof of Proposition \ref{prop:optimal_t}}
\begin{proof}
The objective in Algorithm \ref{algo:tnet} obviously equals to:
\begin{equation}\label{optimal_p_eq1}
\begin{aligned}
\argmin_{t_\theta}\E_{\vx_0\sim q_0\left(\mathbf{x}_0\right),\vx_t \sim \mathcal{N}( \alpha(t) \vx_0, \sigma^2(t)\mI} \abs{\vt_\theta(\vx_t,t) - \frac{\norm{\vx_0}^2}{d} }^2.
\end{aligned}
\end{equation}
By expressing the expectation of Equation \eqref{optimal_p_eq1} in the form of a marginal distribution, we can transform the objective as follows:
\begin{equation}
\begin{aligned}
\argmin_{t_\theta}\sum_i {\frac{1}{N} (2 \pi \sigma_t^2)^{-\frac{d}{2}}}\abs{\vt_\theta(\vx_t,t) - \frac{\norm{\vy_i}^2}{d} }^2
\end{aligned}
\end{equation}
The optimal $t^*_\theta$ must satisfy the condition that the gradient of the loss equals $0$. Therefore, we have:
\begin{equation}
\begin{aligned}
0&= \nabla_{t^*_\theta(\vx_t, t)} \left[\sum_i \underbrace{\frac{1}{N} (2 \pi \sigma_t^2)^{-\frac{d}{2}}}_{A_t} v_i(\vx_t, t) \abs{\frac{\norm{\vy_i}^2}{d} - t^*_\theta(\vx_t, t)}^2 \right] \notag \\
& = \sum_i A_t v_i(\vx_t, t) (t^*_\theta(\vx_t, t) - \frac{\norm{\vy_i}^2}{d}) \notag \\
&= A_t\sum_j v_j(\vx_t, t) t^*_\theta(\vx_t, t) - A_t\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} ,
\end{aligned}
\end{equation}
Thus
\begin{equation}
\begin{aligned}
t^*_\theta(\vx_t, t) &= \frac{ A_t\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} }{A_t\sum_j v_j(\vx_t, t)}\\
&=\sum_i\frac{ v_i(\vx_t, t) }{\sum_j v_j(\vx_t, t)}\frac{\norm{\vy_i}^2}{d}\\
&=\frac{1}{d}\sum_i w_i(\vx_t,t) \norm{\vy_i}^2
\end{aligned}
\end{equation}
We have successfully completed the proof that the optimal $t_\theta(\vx_t, t)$, as trained by Algorithm \ref{algo:tnet}, is equivalent to $\frac{1}{d}\sum_i w_i(\vx_t,t) \norm{\vy_i}^2$.
\end{proof}
% \begin{proof}
% % Here, we adopt a proof routine similar to that which is presented in \cite{zhang2024formulating}.
% We first turn the objective in Algorithm \ref{algo:tnet} to be a constrained optimization problem:
% \begin{equation}\label{constrain_opt}
% \begin{cases}
% \mathop{\arg\min}\limits_{t_\theta} ~~~\mathcal{L}_2, \\
% s.t. ~~ t_\theta = \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d},~~ t_\theta \ge 0.
% \end{cases}
% \end{equation}
% Note that the above constraints only represent a constructive trick, which we will later prove can be omitted. By applying the Karush–Kuhn–Tucker condition \cite{nocedal1999numerical,boyd2004convex} to equation \eqref{constrain_opt}, we can deduce the following:
% \begin{align}
% 0&=\nabla_{t_\theta(\vx_t, t)}\left[ \mathcal{L}_2 + \nu ( t_\theta(\vx_t, t) - \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d}) - \mu^T t_\theta(\vx_t,t)\right] \notag \\
% &= \nabla_{t_\theta(\vx_t, t)} \left[\sum_i \underbrace{\frac{1}{N} (2 \pi \sigma_t^2)^{-\frac{d}{2}}}_{A_t} v_i(\vx_t, t) \abs{\frac{\norm{\vy_i}^2}{d} - t_\theta(\vx_t, t)}^2 + \nu ( t_\theta(\vx_t, t) - \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d}) -\mu^T t_\theta(\vx_t,t)\right] \notag \\
% & = \sum_i A_t v_i(\vx_t, t) (t_\theta(\vx_t, t) - \frac{\norm{\vy_i}^2}{d}) + \nu -\mu \notag \\
% &= A_t\sum_i v_i(\vx_t, t) t_\theta(\vx_t, t) - A_t\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} + \nu -\mu,
% \end{align}
% which leads to
% \begin{equation}
% t^*_\theta(\vx_t, t) = \frac{\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} - \nu / A_t +\mu /A_t}{\sum_j v_j(\vx_t, t)},
% \end{equation}
% where $\mu \ge 0$.
% \begin{equation}
% t^*_\theta(\vx_t, t) = \frac{\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} - \nu / A_t +\mu /A_t}{\sum_j v_j(\vx_t, t)} = \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} - \nu / A_t + \mu / A_t,
% \end{equation}
% Because $ t^*_\theta(\vx_t, t) = \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d}$, we have, we can indicate that $\nu = \mu \ge 0$.
% Since $t^*_\theta(\vx_t, t) \mu = 0$, $\forall c$, we have
% \begin{equation}
% (\frac{\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} - \nu / A_t + \mu / A_t}{\sum_j v_j(\vx_t, t)}) \mu = 0.
% \end{equation}
% If let $\mu \neq 0$, then
% \begin{equation}
% \sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} - \nu / A_t + \mu /A_t = 0,
% \end{equation}
% then
% \begin{equation}
% \nu > \mu \ge 0,
% \end{equation}
% which lead to the contradiction of $\nu > \mu$ and previous $\nu = \mu \ge 0$.
% As as result, it must be that $\mu=0$ and $\nu=0$, and
% \begin{equation}
% t^*_\theta(\vx_t, t) = \frac{\sum_i v_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d} - \nu / A_t + \mu /A_t}{\sum_j v_j(\vx_t, t)} = \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d}.
% \end{equation}
% The Lagrange multiplier $\nu$ and $\mu$ are zero, which means we can omit the constrain $ t_\theta(\vx_t, t) = \sum_i w_i(\vx_t, t) \frac{\norm{\vy_i}^2}{d}$ and $t_\theta \ge 0$ and still get to the conclusion.
% \end{proof}
\subsection{Proof of Proposition \ref{prop:error_bound_trace}}
The approximation error of the estimated trace \eqref{trace_appro} will be its difference from the true Fisher information trace \eqref{trace_matching}. We use consecutive Cauchy–Schwarz and triangle inequality to get the bound of the approximation error:
\begin{equation}
\begin{aligned}
&\abs{\frac{d}{\sigma_t^2} - \frac{\alpha_t^2}{\sigma_t^4}\left[ \sum_i w_i \norm{\vy_i}^2 - \norm{\sum_i w_i \vy_i}^2 \right] - \left\{d \left[\frac{1}{\sigma_t^2} - \frac{\alpha_t^2}{\sigma_t^4}\left( t_\theta(\vx_t,t) - \norm{\frac{\vx_t-\sigma_t\vepsilon_\theta(\vx_t,t)}{\alpha_t}}^2 \right)\right]\right\}}\\
=&\frac{\alpha_t^2}{\sigma_t^4}\abs{\sum_i w_i \norm{\vy_i}^2 - \norm{\sum_i w_i \vy_i}^2 - \left( t_\theta(\vx_t,t) - \norm{\frac{\vx_t-\sigma_t\vepsilon_\theta(\vx_t,t)}{\alpha_t}}^2 \right)}\\
\leq&\frac{\alpha_t^2}{\sigma_t^4}\left[\abs{\sum_i w_i \norm{\vy_i}^2 -t_\theta(\vx_t,t) }+\abs{ \norm{\sum_i w_i \vy_i - \frac{\vx_t-\sigma_t\vepsilon_\theta(\vx_t,t)}{\alpha_t}}}\right]\\
% \leq&\frac{\alpha_t^2}{\sigma_t^4}\left[\delta_1+ \norm{\sum_i w_i \vy_i}^2 + \norm{\frac{\vx_t-\sigma_t\vepsilon_\theta(\vx_t,t)}{\alpha_t}}^2\right]\\
\leq&\frac{\alpha_t^2}{\sigma_t^4}\left[\delta_1+\frac{\sigma_t^2}{\alpha_t^2}\delta_2^2\right]\\
=&\frac{\alpha_t^2}{\sigma_t^4}\delta_1 +\frac{1}{\sigma_t^2}\delta_2^2
\end{aligned}
\end{equation}
\subsection{Proof of Proposition \ref{prop:error_bound_ea}}
\begin{equation}
\begin{aligned}
&\norm{\frac{1}{\sigma_t}\mI - \frac{\alpha_t^2}{\sigma_t^3}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right] - \left(\frac{1}{\sigma_t}\mI - \frac{\alpha_t^2}{\sigma_t^3}\left( \vx_0\vx_0^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)\right)}_{HS}\\
=&\norm{- \frac{\alpha_t^2}{\sigma_t^3}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right] - \left( - \frac{\alpha_t^2}{\sigma_t^3}\left( \vx_0\vx_0^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top} \right)\right)}_{HS}\\
\leq& \frac{\alpha_t^2}{\sigma_t^3}\norm{\sum_i w_i \vy_i\vy_i^{\top} - \vx_0\vx_0^{\top}}_{HS}+\frac{\alpha_t^2}{\sigma_t^3}\norm{\left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top}}_{HS}\\
=&\frac{\alpha_t^2}{\sigma_t^3}\sum_i w_i\norm{ \vy_i\vy_i^{\top} - \vx_0\vx_0^{\top}}_{HS}+\frac{\alpha_t^2}{\sigma_t^3}\norm{\left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} - \bar{\vy}_\theta(\vx_t, t)\bar{\vy}_\theta(\vx_t, t)^{\top}}_{HS}\\
\leq&\frac{\alpha_t^2}{\sigma_t^3}\sum_i w_i \max_i \norm{ \vy_i\vy_i^{\top} - \vx_0\vx_0^{\top}}_{HS}+\frac{\alpha_t^2}{\sigma_t^3}\sum_j \abs{\sum_i w_i \vy_i[j] - \bar{\vy}_\theta( \vx_t, t)[j]}\norm{\sum_i w_i \vy_i -\bar{\vy}_\theta( \vx_t, t)}\\
\leq& \frac{\alpha_t^2}{\sigma_t^3}\left(2\mathcal{D}_y^2+ \sqrt{d}\delta_2\right)
\end{aligned}
\end{equation}
\subsection{Preliminaries on Optimal Transport}\label{appendix:preliminary_ot}
The optimal transport is the general problem of moving one distribution of mass to another as efficiently as possible.
The \emph{optimal transport problem} can be formulated in two primary ways, namely the Monge formulation \citep{monge1781memoire} and the Kantorovich formulation \citep{kantorovich1960mathematical}. Suppose there are two probability measures $\mu$ and $\nu$ on $(\mathbb{R}^n, \mathcal{B})$, and a cost function $c : \mathbb{R}^n \times \mathbb{R}^n \rightarrow \left[ 0, + \infty \right] $. The \emph{Monge problem} is
\begin{equation}
\text{(MP) } ~~~~~~~~~~~~~~\inf_{\text{T}} \left\{ \int c(x, \text{T}(x)) \,{\rm d}\mu(x) : \text{T}_{\texttt{\#}}\mu = \nu \right\}.
\end{equation}
The measure $\text{T}_{\texttt{\#}}\mu$ is defined through $\text{T}_{\texttt{\#}}\mu(A) = \mu(\text{T}^{-1} (A))$ for every $A \in \mathcal{B}$ and is called the \emph{pushforward} of $\mu$ through T.
It is evident that the Monge Problem (MP) transports the entire mass from a particular point, denoted as $x$, to a single point $\text{T}(x)$. In contrast, Kantorovich provided a more general formulation, referred to as the \emph{Kantorovich problem}:
\begin{equation}
\text{(KP) } ~~~~~~~~~~~~~~\inf_{\gamma} \left\{ \int_{\mathbb{R}^n \times \mathbb{R}^n} c\,{\rm d}\gamma : \gamma \in \mit\Pi(\mu, \nu) \right\},
\end{equation}
where $\mit\Pi(\mu, \nu)$ is the set of \emph{transport plans}, i.e.,
\begin{equation}
\mit\Pi(\mu, \nu) = \left\{ \gamma \in \mathcal{P} (\mathbb{R}^n \times \mathbb{R}^n) : (\pi_x)_{\texttt{\#}}\gamma = \mu, (\pi_y)_{\texttt{\#}}\gamma = \nu \right\},
\end{equation}
where $\pi_x$ and $\pi_y$ are the two projections of $\mathbb{R}^n \times \mathbb{R}^n$ onto $\mathbb{R}^n$.
For measures absolutely continuous with respect to the Lebesgue measure, these two problems are equivalent \cite{villani2009optimal}. However, when the measures are discrete, they are entirely distinct as the constraint of the Monge Problem may never be fulfilled.
\subsection{Proof of Corollary \ref{prop:ot}}\label{appendix:proof_ot}
To prove the Corollary \ref{prop:ot}, we first introduce two theorems to transform the problem of whether the PF-ODE mapping is a Monge map into the task of deciding the convexity of the potential function of $T_{s,T}$.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{theorem}\label{theorem:convex}\citep[Theorem 1.48]{santambrogio2015optimal}
Suppose that $\mu$ is a probability measure on $(\mathbb{R}^n, \mathcal{B})$ such that $\int |x|^2\rd\mu(x) < \infty$ and that $u: \mathbb{R}^n \rightarrow \mathbb{R} \cup \{+\infty\}$ is convex and differentiable $\mu {\text -}a.e$. Set $\text{\normalfont T} = \nabla u$ and suppose $\int |\text{\normalfont T}(x)|^2\rd\mu(x) < \infty$. Then {\normalfont T} is optimal for the transport cost $c(x, y) = \frac{1}{2} |x - y| ^ 2$ between the measures $\mu$ and $\nu = \text{\normalfont T}_{\texttt{\#}}\mu$.
\end{theorem}
\end{tcolorbox}
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{theorem}\label{theorem:brenier}\textbf{The Brenier's Theorem.} \citep[Theorem 1.22]{santambrogio2015optimal} \citep{brenier1987decomposition,brenier1991polar}
Let $\mu, v$ be probabilities over $\mathbb{R}^d$ and $c(x, y)=\frac{1}{2}|x-y|^2$. Suppose $\int|x|^2 \mathrm{~d} x, \int|y|^2 \mathrm{~d} y<+\infty$, which implies $\min (\mathrm{KP})<+\infty$ and suppose that $\mu$ gives no mass to $(d-1)$ surfaces of class $C^2$. Then there exists, uniquely, an optimal transport map T from $\mu$ to $v$, and it is of the form $\mathrm{T}=\nabla u$ for a convex function $u$.
\end{theorem}
\end{tcolorbox}
To ensure the existence of the potential function, we need to leverage the following
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{theorem}\label{theorem:poincare}\textbf{The Poincar\'e's Theorem.} \citep[Theorem 4.1 of Chapter V, \S 4]{lang2012fundamentals}
Let $U$ be an open ball in $\mathbb{R}^n$ and let $\omega$ be a differential form of degree $\geq 1$ on $U$ such that ${\rm d} \omega = 0$. Then there exists a differential form $\phi$ on $U$ such that ${\rm d} \phi = \omega$.
\end{theorem}
\end{tcolorbox}
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{remark}
The conclusion remains valid when the open ball $U$ is substituted with the entirety of $\mathbb{R}^n$
\end{remark}
\end{tcolorbox}
For $s>0$, it is clear that $\frac{\rd q_T(x_T)}{\rd q_s(x_s)}$ is non-singular, and therefore, it satisfies the requirements of Brenier's Theorem \ref{theorem:brenier}, leading to the existence of a unique optimal transport map. According to Theorem \ref{theorem:convex}, if we can establish that the potential function of $T_{s,T}$ is convex, then the PF-ODE mapping will indeed be a Monge map.
Notice that the existence of the potential map is guaranteed by Poincaré's Theorem \ref{theorem:poincare}.
We can now convert the condition of the potential function of $T_{s,T}$ being convex into the condition that its Jacobian, $\frac{\partial T_{s,T}(x_s)}{\partial x_s}$, is positive semi-definite, as per the following theorem in convex analysis.
\begin{tcolorbox}[colback=gray!10, colframe=white, sharp corners, boxrule=0pt]
\begin{theorem}\label{theorem:hessian_convex}
\citep[Theorem 4.5]{rockafellar2015convex}
Let $f$ be a twice continuously differentiable real-valued function on an open convex set $C$ in $R^n$. Then $f$ is convex on $C$ if and only if its Hessian matrix
\begin{equation}
Q_x=\left(q_{i j}(x)\right), \quad q_{i j}(x)=\frac{\partial^2 f}{\partial \xi_i \partial \xi_j}\left(\xi_1, \ldots, \xi_n\right)
\end{equation}
is positive semi-definite for every $x \in C$.
\end{theorem}
\end{tcolorbox}
If we denote that
\begin{equation}
\begin{aligned}
\mA(t) = \frac{\partial T_{t,T}(x_t)}{\partial x_t}
\end{aligned}
\end{equation}
obviously, $\mA(T) = I$ is p.s.d., our goal is to answer when $\mA(t)$ is p.s.d..
We try to answer this to set up a connection between $\mA(t)$ and $\mA(T) = I$.
We can derive that:
\begin{equation}
\begin{aligned}
\frac{\rd \mA(t)}{\rd t} = & \lim_{\epsilon \rightarrow 0^+} \frac{\mA(t + \epsilon) - \mA(t)}{\epsilon} \\
= & \lim_{\epsilon \rightarrow 0^+} \frac{\mA(t + \epsilon) - \mA(t + \epsilon) \nabla_{\vx_t} \left[\vx_t + \epsilon \left(f(t)\vx_t -\frac{g^2(t)}{2}\nabla_{\vx_t}\log_t q_t(\vx_t) \right)+ \mathcal{O}(\epsilon^2)\right]}{\epsilon} \\
= & \lim_{\epsilon \rightarrow 0^+} \frac{\epsilon f(t) \mA(t + \epsilon) + \frac{g^2(t)}{2} \epsilon \mA(t + \epsilon) \nabla_{\vx_t}\log_t q_t(\vx_t) + \mathcal{O}(\epsilon^2)}{\epsilon}\\
= & f(t) \mA(t) + \frac{g^2(t)}{2} \mA(t) \nabla_{\vx_t}\log_t q_t(\vx_t)\\
= & f(t) \mA(t) - \frac{g^2(t)}{2} \mA(t) \left\{\frac{1}{\sigma_t^2}\mI - \frac{\alpha_t^2}{\sigma_t^4}\left[ \int w(\vy) \vy\vy^{\top}\rd q_0(\vy) \!- \!\left(\int w(\vy) \vy \rd q_0(\vy)\right)\left(\int w(\vy) \vy \rd q_0(\vy)\right)^{\top} \right]\right\} \\
=&\mA(t)\underbrace{\left\{\left[f(t)-\frac{g^2(t)}{2\sigma_t^2}\right]\mI+\frac{\alpha_t^2 g^2(t)}{2\sigma_t^4}\left[ \sum_i w_i \vy_i\vy_i^{\top} - \left(\sum_i w_i \vy_i\right)\left(\sum_i w_i \vy_i\right)^{\top} \right]\right\}}_{\mB(t)}
\end{aligned}
\end{equation}
Notice that the above ODE starts from $T$.
According to the Solution Matrices theory \citep{masuyama2016limit}\footnote{The Definition 6.2 in \url{https://math.mit.edu/~jorloff/suppnotes/suppnotes03/ls6.pdf} suffice the result here.}, let us denote the $\mC(t)$ is the normalized fundamental matrix at $T$ for $\mB(T)$, which implies $C(t)$ is the solution to the following ODE:
\begin{equation}
\begin{aligned}
\quad \mC'(t) = \mC(t)\mB(t),\mC(T)=I,\quad(\text{flow from T to t})
\end{aligned}
\end{equation}
Then we can deduce that
\begin{equation}
\begin{aligned}
\frac{\partial T_{t,T}(x_t)}{\partial x_t}&=\mA(t)\\
&=\mC(t)\mA(T)\\
&=\mC(t)\mI\\
&=\mC(t)\\
\end{aligned}
\end{equation}
Thus, the diffeomorphism $T_{s,T}$ is a Monge optimal transport map if and only if $\mC(s)$ is semi-positive definite. Note that the above requirement needs to be satisfied for every PF-ODE chain $x_t,t\in[T,t]$.
\section{Experiments Details}
\subsection{Evaluation of NLL}\label{app:extra_nll}
For toy experiments, we follow the experiment design in \citet{pmlr-v162-lu22f} and adopt the same network architecture for the score network and our DF-TM network.
We calculate the true diffusion Fisher with 5,000 data samples from the cheeseboard distribution.
For commercial-level experiments, we employ the explicit Euler method to compute the NLL, excluding the final step near $t=0$, for reasons discussed in Appendix \ref{app:singularity}. We evaluate the NLL across 10 steps throughout the timeline of the PF-ODE. The DF-TM network we trained uses the float-point16 data type.
\paragraph{Network architectures}
In terms of network architecture, we employ an SD Unet structure with an additional MLP head. However, we believe that a lighter network could potentially be sufficient for DF-TM.
\paragraph{Training cost of DF-TM}
For training the DF-TM network, we approximately spend 24 hours using 8 Ascend 910B chips. Given the size of the Laion2B-en dataset (which contains 2.32 billion images), this is quite an efficient speed. Additionally, the convergence behavior of the training loss is robust, as illustrated in Figure \ref{fig:train_loss}. We also hypothesize that our network design has redundancy, suggesting that we could further reduce costs by opting for lighter networks. We will provide more details about training costs in our revised manuscript.
\subsection{Adjoint Guidance Sampling}\label{app:extra_adjoint}
For Figure \ref{fig:adjoint_score}, we examined varying numbers of adjoint guidance, ranging from 0 to 20, under a full inference number of 50. The adjoint guidance scale was grid-searched by the VJP method.
For experiments on Pick-Score, we use NVIDIA V100 chips, and the rest experiments use Ascend 910B chips.
\paragraph{Base method for DF-EA}
Several variants of adjoint-optimization algorithms exist, such as AdjointDPM \citep{pan2023adjointdpm}, AdjointDES \citep{blasingame2024adjointdeis}, and SAG \citep{pan2023towards}. However, while these algorithms differ in their design of solvers for the adjoint ODE, they all utilize VJP when accessing the diffusion Fisher. We selected SAG as our base method due to its state-of-the-art performance. We believe that replacing VJP with DF-EA could also enhance the performance of algorithms like AdjointDPM and AdjointDES.
\paragraph{The Design of Score functions}
\begin{itemize}
\item Aesthetic Score (SAC/AVA/Pick-Score)
For the aesthetic score predictor $f_{aes}:\sR^d\mapsto\sR$, the adjoint optimization target is simply $f_{aes}$ itself.
\begin{equation}
\begin{aligned}
\mathcal{L}(\vx_0) := f_{aes}(\vx_0)
\end{aligned}
\end{equation}
\item Clip Loss
Following the implementation of \citep{pan2023adjointdpm}, we use the features from the CLIP image encoder as our feature vector. The loss function is $L_2$-norm between the Gram matrix of the style image and the Gram matrix of the estimated clean image.
\begin{equation}
\begin{aligned}
\mathcal{L}(\vx_0) := \norm{\text{clip}(\vx_0)\text{clip}(\vx_0)^\top-\text{clip}(\vx_{ref})\text{clip}(\vx_{ref})^\top}
\end{aligned}
\end{equation}
\item FaceID Loss
Following the implementation of \citep{pan2023towards}, we use ArcFace to extract the target features of reference faces to represent face IDs and compute the $l_2$ Euclidean distance between the extracted ID features of the estimated clean image and the reference face image as the loss function.
\begin{equation}
\begin{aligned}
\mathcal{L}(\vx_0) := \norm{\text{ArcFace}(\vx_0)-\text{ArcFace}(\vx_{ref})}
\end{aligned}
\end{equation}
\end{itemize}
\paragraph{Hyperparameters}
For the hyperparameters in adjoint-guided sampling, we ensure a fair comparison between the VJP and DF-EA methods. For most hyperparameters, we directly adopt the settings from previous works \cite{pan2023towards} for both the baseline VJP method and our DF-EA method. For the guidance scale, we tune the value for the VJP method and use the same value for our DF-EA method. The DF-EA method does not introduce additional hyperparameters, and for mutual hyperparameters, our DF-EA method uses the exact same values as the VJP method. We adopt this strategy because our DF-EA method solely improves the approximation of the diffusion Fisher linear operator, without altering the adjoint sampling mechanism. Therefore, the suitable hyperparameters should remain unchanged, and we simply use the same parameters from the VJP method for our DF-EA method.
For all experiments, we set the number of sampling steps to $T=50$. Adjoint guidance is applied starting from steps ranging from 15 to 35 and ending at step 35, with one guidance per step.
Thus the only parameter we tune is the guidance strength.
We determine this value for the VJP method via a grid search from 0.1 to 0.5 with a step size of 0.1 and find that the optimal guidance strength for VJP is 0.2. We then use this value for our DF-EA method. Notice that, we apply a normalization to the guidance gradient for both VJP and DF-EA methods, making our optimal guidance strength consistent across different scores.
The tuning is conducted on 1k COCO prompts, and the computational budget for tuning is 4 * 5 * 3 GPU hours (4 tasks * 5 grids * 3 hours per single test) on Ascend 910B chips.
\subsection{Numerical OT Experiments}\label{app:extra_ot}
\paragraph{The detailed algorithm }
The detailed version of the Algorithm \ref{algo:ot} is presented in Algorithm \ref{algo:to_detail}. We use all the data points in the initial dataset to calculate every quantum we need in the sampling trajectory to validate the condition discussed in Corollary \ref{prop:ot}.
\begin{algorithm}[h]
\caption{Detailed numerical OT test for PF-ODE map}\label{algo:to_detail}
% \small
\begin{algorithmic}[1] %[1] enables line numbers
\STATE \textbf{Input}: initial data $\{\vy_j\}_{j=1}^{N}$, noise schedule $\{\alpha(t)\}$ and $\{\sigma(t)\}$, discretization steps $M$.
\STATE Initialize $\mA_M = \mI$, $\vx_M \sim \mathcal{N}(0,\sigma_T\mI)$.
\FOR{$i=M,M-1,\cdots,1$}
\STATE $\alpha_i = \alpha(t_i),\quad \sigma_i = \sigma(t_i)$
\STATE $\rd t = t_{i-1} - t_{i}$.
\STATE $f_i = \frac{\rd \log \alpha_i}{\rd t}$.
\STATE $g_i = \sqrt{\frac{\rd \sigma_i^2}{\rd t}-2 \frac{\rd \log \alpha_i}{\rd t}\sigma_i^2} $.
\FOR{$j=1,\cdots,N$}
\STATE calculate $v_j$ via $v_j = \exp(-\frac{|\vx_i - \alpha_i \vy_j|^2}{2\sigma_i^2})$
\ENDFOR
\FOR{$j=1,\cdots,N$}
\STATE calculate $w_j$ via $w_j = \frac{v_j}{\sum_k v_k}$.
\ENDFOR
\STATE $\mB_i = \left[f_i-\frac{g^2_i}{2\sigma_i^2}\right]\mI+\frac{\alpha_i^2 g^2_i}{2\sigma_i^4}\left[ \sum_j w_j \vy_j\vy_j^{\top} - \left(\sum_j w_j \vy_j\right)\left(\sum_j w_j \vy_j\right)^{\top} \right].$
\STATE $\mA_{i-1} = \mA_{i} + \rd t * \mA_{i}^\top\mB_i $\hfill\COMMENT{solve fundamental matrix}
\STATE $s_i = \frac{\vx_i - \sum_jw_jy_j}{\sigma_i^2}$
\STATE $\vx_{i-1} =\vx_i + \left[ f_i\vx_i - \frac{1}{2}g^2_i s_i \right] * \rd t$
\ENDFOR
% \STATE $\phi = \frac{\norm{\mA-\mA^T}_F}{\sqrt{2}\norm{\mA}_F}$
\STATE \textbf{Output}:$\mA_0$.
\end{algorithmic}
\end{algorithm}
\paragraph{Initial data}
In Table \ref{tab:ot}, we adopt numerical verification of OT on 2-D synthesized data.
For the single-Gaussian case, we fix the one mode central in (0.5,0.5) and test at $s=0.1$;
For the affine data case, we set three data points as (0.2,-0.4), (0.2,0.0), and (0.2,0.9), and test at $s=0.0$;
For the non-affine data case, we set three data points as (0.0,0.5), (0.0,0.0), and (0.5,0.0), and test at $s=0.0$.
\paragraph{The definition of Asymmetric rate}
In Table \ref{tab:ot}, we need the asymmetry rate $I$ to point out that the PF-ODE map in very simple non-affine data cannot be OT.
Let $A$ be an $n \times n$ matrix. The asymmetry rate $I$ is defined as:
$$
I=\frac{\left\|A-A^T\right\|_F}{\sqrt{2}\|A\|_F}
$$
where $\|M\|_F=\sqrt{\sum_{i=1}^n \sum_{j=1}^n m_{i j}^2}$ is the Frobenius norm of the matrix $M$. The denominator $\sqrt{2}\|A\|_F$ is used to normalize the index to the range of $[0,1]$.
The asymmetry rate of the non-affine data in Table \ref{tab:ot} has significantly deviated from zero.
\subsection{Pretrained Models}\label{app:exp_models}
All of the pretrained models used in our research are open-sourced and available online as follows:
\begin{itemize}
% \item CIFAR10 generation : ddpm\_ema\_cifar10
% \href{https://github.com/VainF/Diff-Pruning/releases/download/v0.0.1/ddpm_ema_cifar10.zip}{https://github.com/VainF/Diff-Pruning/releases/download/v0.0.1/ddpm\_ema\_cifar10.zip}
% \item CelebA-HQ generation and interpolation : ddpm-ema-celebahq-256
% \href{https://huggingface.co/google/ddpm-ema-celebahq-256}
% {https://huggingface.co/google/ddpm-ema-celebahq-256}
\item stable-diffusion-v1-5
% \href{https://huggingface.co/runwayml/stable-diffusion-v1-5}
% {https://huggingface.co/runwayml/stable-diffusion-v1-5}
\url{https://huggingface.co/runwayml/stable-diffusion-v1-5}
\item stable-diffusion-2-base
% \href{https://huggingface.co/stabilityai/stable-diffusion-2-base}
% {https://huggingface.co/stabilityai/stable-diffusion-2-base}
\url{https://huggingface.co/stabilityai/stable-diffusion-2-base}
\item SAC-aesthetic score predictor
% \href{https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/sac%2Blogos%2Bava1-l14-linearMSE.pth}
% {https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/sac\%2Blogos\%2Bava1-l14-linearMSE.pth}
\url{https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/sac\%2Blogos\%2Bava1-l14-linearMSE.pth}
\item AVA-aesthetic score predictor
% \href{https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/ava%2Blogos-l14-linearMSE.pth}
% {https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/ava\%2Blogos-l14-linearMSE.pth}
\url{https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/ava\%2Blogos-l14-linearMSE.pth}
\item ArcFace ID loss
% \href{https://github.com/TreB1eN/InsightFace_Pytorch}
% {https://github.com/TreB1eN/InsightFace\_Pytorch}
\url{https://github.com/TreB1eN/InsightFace\_Pytorch}
\item Clip loss
% \href{https://huggingface.co/openai/clip-vit-large-patch14}
% {https://huggingface.co/openai/clip-vit-large-patch14}
\url{https://huggingface.co/openai/clip-vit-large-patch14}
\item Pick-Score
\url{https://github.com/yuvalkirstain/PickScore}
\end{itemize}
\section{Discussions}\label{app:discussion}
\subsection{Singularity of Fisher information at $t=0$}\label{app:singularity}
Previous studies \citep{yang2023lipschitz,zhang2024tackling} have shown that the diffusion model, particularly when learned in $\epsilon$-prediction, can encounter a singularity issue at $t=0$. Our DF in \eqref{information_dirac} reaffirms this issue, as this formulation becomes ill-formed at $t=0$ due to division by zero ($\sigma_0$). Consequently, our formulation does not describe the behavior at $t=0$. The deep theoretical exploration of the singularity problem remains an open question in the diffusion model field. However, as it is not the primary focus of this paper, we will not discuss the DF at $t=0$.
\subsection{Statistical Caliber of Negative Log-Likelihood}
When dealing with high-dimensional data such as images, direct likelihood comparisons may encounter scaling issues due to the dimensionality. In this study, unless explicitly indicated otherwise, we adopt the approach of \cite{pmlr-v202-zheng23c} and typically use Negative Log-Likelihood (NLL) to refer to Bits Per Dimension (BPD).
\begin{equation}
\begin{aligned}
\text{BPD} =\mathbb{E}_{\boldsymbol{x}_0 \sim q_0}\left[\frac{-\log P_0\left(\boldsymbol{x}_0\right)}{d \log 2}\right]
\end{aligned}
\end{equation}
\subsection{The ODE Solvers}
To compute the numerical solutions for the PF-ODE in \eqref{PFODE}, the likelihood ODE in \eqref{nll_ode}, and the adjoint ODE in \eqref{Adjoint_ODE}, we require ODE solvers. In our paper's experiments, we consistently use the explicit Euler method (referred to as DDIM when applied to PF-ODE). However, it's important to note that our approach is not dependent on a specific ODE solver. We can also utilize alternatives like fast ODE solvers \citep{lu2022dpm,lu2022dpm++,liu2022pseudo} or exact inversion ODE solvers \citep{wallace2023edict,zhang2023exact}.
\subsection{The relation of our outer-product form DF to Covariance in DMs}\label{app:covariance}
We note that there is a series of studies aiming to learn the covariance of the reverse diffusion Stochastic Differential Equation (SDE) \citep{bao2022analytic,bao2022estimating}. In Bayesian statistics, Fisher information is defined as the covariance of the score. These studies derive their formulation by analyzing the covariance of the score and obtaining the Fisher information in terms of the score. However, our DF is derived directly from the marginal distribution and is composed solely of the initial distribution and noise schedule. Furthermore, our application is unique; we are the first to replace the use of VJP with DF, while the focus of these studies is to enhance the performance of Diffusion Models (DMs) with analytical covariance.
A very similar form of the Fisher information is proposed in Lemma 5 of \cite{benton2020nearly}.
The difference is that our Proposition 3 presents the specific form of Fisher information in terms of data distribution, which is not included in Lemma 5 of \cite{benton2020nearly}. This distinction is crucial as it facilitates the derivation of our new algorithms.
Also, we adopt the currently more commonly used $\alpha_t,\sigma_t$ notation to represent the noise schedule. Instead, \cite{benton2020nearly} $\alpha_t\equiv e^{-2t}$.
We gave out this detailed formulation to avoid unnecessary misunderstandings in the development of the subsequent training scheme.
\subsection{Comparison of DF-TM to Hutchinson trace estimator}\label{app:hutchinson}
We notice that the naive trace calculation of the VJP method can be accelerated by the Hutchinson trace estimator.
We have not conducted experiments using the Hutchinson estimator, as it is a Monte-Carlo type estimation, unlike the direct evaluation methods like the full-VJP baseline and our DF-TM. Furthermore, the VJP method, with the help of the Hutchinson method and its variants, is still considerably more expensive than our method, and its applicability is also more restricted.
\begin{itemize}
\item \textbf{The Hutchinson method is much more costly:}
To attain a relative error less than $\epsilon$ with a probability of $1-\delta$, the Hutchinson method requires $\frac{2(1-\frac{8}{3}\epsilon)\log\frac{1}{\delta}}{\epsilon^2}$ samples \cite{skorski2021modern}.
Assuming that the goal is to obtain an estimation with a relative error of less than 10\% with a probability exceeding 90\%, a minimum of 351 NFEs is required. This equates to 1 or 2 minutes on SD-v1.5 for a single trace estimation. For a complete NLL estimation of an image with 20 steps, this would take 20 minutes, which is entirely impractical in any business context. In contrast, our DF-TM only requires 1 NFE for a single trace estimation, needing merely 10 seconds for the full NLL estimation of an image.
\item \textbf{The application of the Hutchinson method is more restricted:}
Due to its Monte-Carlo characteristics, the Hutchinson method is more appropriate for contributing to the computation of certain training objectives, as in \cite{pmlr-v162-lu22f}, where the unbiased property is sufficient, and large variance may be absorbed into network training. However, the Hutchinson method may encounter difficulties with accurate per-sample trace estimation. Our method can accommodate both scenarios, including per-sample computation.
\end{itemize}
There are already attempts to use the Hutchinson method to expedite the VJP of trace estimation in diffusion models \cite{pmlr-v162-lu22f}. Nonetheless, due to Hutchinson's limitations, these practices are restricted to relatively small DMs (CIFAR-10 at most). We will add discussions on the Hutchinson method in our revision.
\subsection{Discussions on the theoretical bound of DF-EA}
We notice that the error bound in Proposition \ref{prop:error_bound_ea} does not vanish as the training error decreases.
From the rigorous theoretical perspective, considering the error bounds in Propositions \ref{prop:error_bound_trace} and \ref{prop:error_bound_ea}, the DF-EA method is less valid than the DF-TM method, as we currently cannot establish vanishing bounds for it.
However, from an empirical perspective, our DF-EA outperforms the naive VJP method in terms of score improvement across various tasks and pretrained models, as demonstrated in Figure 4. Thus, the DF-EA approximation proves to be valid in a practical sense.
The replacement $\sum_i\vy_i\vy_i^{\top}\approx \vx_0\vx_0^{\top}$ originates from the observation that $w_i$ is a weighting of the summation equal to 1, and as $t$ approaches 0, the $w_i$ closest to $\vx_0$ will dominate due to the diffusion kernel. Therefore, this approximation is intuitively reasonable near $t=0$, which is precisely where we apply adjoint guidance.
\subsection{Limitations}
This paper does not explore the integration of DF into accelerated ODE solvers \cite{lu2022dpm}, or exact inversion ODE samplers \cite{wang2024belm}.
This paper is constrained in the scope of DMs, but similar second-order information may also exist in flow-matching generative models \cite{lipman2022flow, zhu2024analyzing} or variational inference models \cite{zhu2024neural,wang2024gad}, language models \cite{fu2025iap}.
Our technique may also have downstream applications in automated driving \cite{tu2025driveditfit}, touch-generation \cite{tu2025texttoucher}, domain-transfer \cite{feng2024lw2g,feng2024pectp}, and missing data imputation \cite{chen2024rethinking}.
The DF may also contribute to a more effective inference method, like L-BFGS, which we did not explore.
\begin{thebibliography}{70}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Albergo et~al.(2023)Albergo, Boffi, and Vanden-Eijnden]{albergo2023stochastic}
Albergo, M.~S., Boffi, N.~M., and Vanden-Eijnden, E.
\newblock Stochastic interpolants: A unifying framework for flows and diffusions.
\newblock \emph{arXiv preprint arXiv:2303.08797}, 2023.
\bibitem[Amann(2011)]{amann2011ordinary}
Amann, H.
\newblock \emph{Ordinary differential equations: an introduction to nonlinear analysis}, volume~13.
\newblock Walter de gruyter, 2011.
\bibitem[Anderson(1982)]{anderson1982reverse}
Anderson, B.~D.
\newblock Reverse-time diffusion equation models.
\newblock \emph{Stochastic Processes and their Applications}, 12\penalty0 (3):\penalty0 313--326, 1982.
\bibitem[Bao et~al.(2022{\natexlab{a}})Bao, Li, Sun, Zhu, and Zhang]{bao2022estimating}
Bao, F., Li, C., Sun, J., Zhu, J., and Zhang, B.
\newblock Estimating the optimal covariance with imperfect mean in diffusion probabilistic models.
\newblock \emph{arXiv preprint arXiv:2206.07309}, 2022{\natexlab{a}}.
\bibitem[Bao et~al.(2022{\natexlab{b}})Bao, Li, Zhu, and Zhang]{bao2022analytic}
Bao, F., Li, C., Zhu, J., and Zhang, B.
\newblock Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models.
\newblock \emph{arXiv preprint arXiv:2201.06503}, 2022{\natexlab{b}}.
\bibitem[Bengio et~al.(2013)Bengio, Yao, and Cho]{bengio2013bounding}
Bengio, Y., Yao, L., and Cho, K.
\newblock Bounding the test log-likelihood of generative models.
\newblock \emph{arXiv preprint arXiv:1311.6184}, 2013.
\bibitem[Benton et~al.(2024)]{benton2020nearly}
Benton, J. et~al.
\newblock Nearly $ d $-linear convergence bounds for diffusion models via stochastic localization.
\newblock In \emph{The Twelfth International Conference on Learning Representations}, 2024.
\bibitem[Blasingame \& Liu(2024)Blasingame and Liu]{blasingame2024adjointdeis}
Blasingame, Z.~W. and Liu, C.
\newblock Adjointdeis: Efficient gradients for diffusion models.
\newblock \emph{arXiv preprint arXiv:2405.15020}, 2024.
\bibitem[Brenier(1987)]{brenier1987decomposition}
Brenier, Y.
\newblock D{\'e}composition polaire et r{\'e}arrangement monotone des champs de vecteurs.
\newblock \emph{CR Acad. Sci. Paris S{\'e}r. I Math.}, 305:\penalty0 805--808, 1987.
\bibitem[Brenier(1991)]{brenier1991polar}
Brenier, Y.
\newblock Polar factorization and monotone rearrangement of vector-valued functions.
\newblock \emph{Communications on pure and applied mathematics}, 44\penalty0 (4):\penalty0 375--417, 1991.
\bibitem[Chen et~al.(2018)Chen, Rubanova, Bettencourt, and Duvenaud]{chen2018neural}
Chen, R.~T., Rubanova, Y., Bettencourt, J., and Duvenaud, D.~K.
\newblock Neural ordinary differential equations.
\newblock \emph{Advances in neural information processing systems}, 31, 2018.
\bibitem[Chen et~al.(2024)Chen, Li, Wang, Zhang, Xu, Jiang, Song, and Wang]{chen2024rethinking}
Chen, Z., Li, H., Wang, F., Zhang, O., Xu, H., Jiang, X., Song, Z., and Wang, E.~H.
\newblock Rethinking the diffusion models for numerical tabular data imputation from the perspective of wasserstein gradient flow.
\newblock \emph{arXiv preprint arXiv:2406.15762}, 2024.
\bibitem[Dhariwal \& Nichol(2021{\natexlab{a}})Dhariwal and Nichol]{NEURIPS2021_49ad23d1}
Dhariwal, P. and Nichol, A.
\newblock Diffusion models beat gans on image synthesis.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~34, pp.\ 8780--8794, 2021{\natexlab{a}}.
\bibitem[Dhariwal \& Nichol(2021{\natexlab{b}})Dhariwal and Nichol]{dhariwal2021diffusion}
Dhariwal, P. and Nichol, A.
\newblock Diffusion models beat gans on image synthesis.
\newblock \emph{Advances in Neural Information Processing Systems}, 34:\penalty0 8780--8794, 2021{\natexlab{b}}.
\bibitem[Feng et~al.(2024{\natexlab{a}})Feng, Zhao, Zhang, Dong, Ding, Jiang, and Qian]{feng2024pectp}
Feng, Q., Zhao, H., Zhang, C., Dong, J., Ding, H., Jiang, Y.-G., and Qian, H.
\newblock Pectp: Parameter-efficient cross-task prompts for incremental vision transformer.
\newblock \emph{arXiv preprint arXiv:2407.03813}, 2024{\natexlab{a}}.
\bibitem[Feng et~al.(2024{\natexlab{b}})Feng, Zhou, Zhao, Zhang, and Qian]{feng2024lw2g}
Feng, Q., Zhou, D., Zhao, H., Zhang, C., and Qian, H.
\newblock Lw2g: Learning whether to grow for prompt-based continual learning.
\newblock \emph{arXiv preprint arXiv:2409.18860}, 2024{\natexlab{b}}.
\bibitem[Fu et~al.(2025)Fu, Zhao, Dong, Zhang, and Qian]{fu2025iap}
Fu, H., Zhao, H., Dong, J., Zhang, C., and Qian, H.
\newblock Iap: Improving continual learning of vision-language models via instance-aware prompting.
\newblock \emph{arXiv preprint arXiv:2503.20612}, 2025.
\bibitem[Gohberg et~al.(1990)Gohberg, Goldberg, Kaashoek, Gohberg, Goldberg, and Kaashoek]{gohberg1990hilbert}
Gohberg, I., Goldberg, S., Kaashoek, M.~A., Gohberg, I., Goldberg, S., and Kaashoek, M.~A.
\newblock Hilbert-schmidt operators.
\newblock \emph{Classes of Linear Operators Vol. I}, pp.\ 138--147, 1990.
\bibitem[Ho \& Salimans(2022)Ho and Salimans]{ho2022classifier}
Ho, J. and Salimans, T.
\newblock Classifier-free diffusion guidance.
\newblock \emph{arXiv preprint arXiv:2207.12598}, 2022.
\bibitem[Ho et~al.(2020)Ho, Jain, and Abbeel]{ho2020denoising}
Ho, J., Jain, A., and Abbeel, P.
\newblock Denoising diffusion probabilistic models.
\newblock \emph{Advances in Neural Information Processing Systems}, 33:\penalty0 6840--6851, 2020.
\bibitem[Ho et~al.(2022)Ho, Saharia, Chan, Fleet, Norouzi, and Salimans]{JMLR:v23:21-0635}
Ho, J., Saharia, C., Chan, W., Fleet, D.~J., Norouzi, M., and Salimans, T.
\newblock Cascaded diffusion models for high fidelity image generation.
\newblock \emph{Journal of Machine Learning Research}, 23\penalty0 (47):\penalty0 1--33, 2022.
\bibitem[Kantorovich(1960)]{kantorovich1960mathematical}
Kantorovich, L.~V.
\newblock Mathematical methods of organizing and planning production.
\newblock \emph{Management science}, 6\penalty0 (4):\penalty0 366--422, 1960.
\bibitem[Karras et~al.(2022)Karras, Aittala, Aila, and Laine]{karras2022elucidating}
Karras, T., Aittala, M., Aila, T., and Laine, S.
\newblock Elucidating the design space of diffusion-based generative models.
\newblock \emph{arXiv preprint arXiv:2206.00364}, 2022.
\bibitem[Khrulkov et~al.(2023)Khrulkov, Ryzhakov, Chertkov, and Oseledets]{khrulkov2023understanding}
Khrulkov, V., Ryzhakov, G., Chertkov, A., and Oseledets, I.
\newblock Understanding {DDPM} latent codes through optimal transport.
\newblock In \emph{International Conference on Learning Representations}, 2023.
\newblock URL \url{https://openreview.net/forum?id=6PIrhAx1j4i}.
\bibitem[Kingma et~al.(2021)Kingma, Salimans, Poole, and Ho]{NEURIPS2021_b578f2a5}
Kingma, D., Salimans, T., Poole, B., and Ho, J.
\newblock Variational diffusion models.
\newblock In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J.~W. (eds.), \emph{Advances in Neural Information Processing Systems}, volume~34, pp.\ 21696--21707. Curran Associates, Inc., 2021.
\newblock URL \url{https://proceedings.neurips.cc/paper_files/paper/2021/file/b578f2a52a0229873fefc2a4b06377fa-Paper.pdf}.
\bibitem[Lang(2012)]{lang2012fundamentals}
Lang, S.
\newblock \emph{Fundamentals of differential geometry}, volume 191.
\newblock Springer Science \& Business Media, 2012.
\bibitem[Lavenant \& Santambrogio(2022)Lavenant and Santambrogio]{lavenant2022flow}
Lavenant, H. and Santambrogio, F.
\newblock The flow map of the fokker--planck equation does not provide optimal transport.
\newblock \emph{Applied Mathematics Letters}, 133:\penalty0 108225, 2022.
\bibitem[Lin et~al.(2014)Lin, Maire, Belongie, Hays, Perona, Ramanan, Doll{\'a}r, and Zitnick]{lin2014microsoft}
Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Doll{\'a}r, P., and Zitnick, C.~L.
\newblock Microsoft coco: Common objects in context.
\newblock In \emph{Computer Vision--ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13}, pp.\ 740--755. Springer, 2014.
\bibitem[Lipman et~al.(2022)Lipman, Chen, Ben-Hamu, Nickel, and Le]{lipman2022flow}
Lipman, Y., Chen, R.~T., Ben-Hamu, H., Nickel, M., and Le, M.
\newblock Flow matching for generative modeling.
\newblock \emph{arXiv preprint arXiv:2210.02747}, 2022.
\bibitem[Liu et~al.(2022)Liu, Ren, Lin, and Zhao]{liu2022pseudo}
Liu, L., Ren, Y., Lin, Z., and Zhao, Z.
\newblock Pseudo numerical methods for diffusion models on manifolds.
\newblock \emph{arXiv preprint arXiv:2202.09778}, 2022.
\bibitem[Loshchilov \& Hutter(2019)Loshchilov and Hutter]{loshchilov2018decoupled}
Loshchilov, I. and Hutter, F.
\newblock Decoupled weight decay regularization.
\newblock In \emph{International Conference on Learning Representations (ICLR)}, 2019.
\bibitem[Lu et~al.(2022{\natexlab{a}})Lu, Zheng, Bao, Chen, Li, and Zhu]{pmlr-v162-lu22f}
Lu, C., Zheng, K., Bao, F., Chen, J., Li, C., and Zhu, J.
\newblock Maximum likelihood training for score-based diffusion {ODE}s by high order denoising score matching.
\newblock In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), \emph{Proceedings of the 39th International Conference on Machine Learning}, volume 162 of \emph{Proceedings of Machine Learning Research}, pp.\ 14429--14460. PMLR, 17--23 Jul 2022{\natexlab{a}}.
\newblock URL \url{https://proceedings.mlr.press/v162/lu22f.html}.
\bibitem[Lu et~al.(2022{\natexlab{b}})Lu, Zhou, Bao, Chen, Li, and Zhu]{lu2022dpm}
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J.
\newblock Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.
\newblock \emph{Advances in Neural Information Processing Systems}, 35:\penalty0 5775--5787, 2022{\natexlab{b}}.
\bibitem[Lu et~al.(2022{\natexlab{c}})Lu, Zhou, Bao, Chen, Li, and Zhu]{lu2022dpm++}
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J.
\newblock Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models.
\newblock \emph{arXiv preprint arXiv:2211.01095}, 2022{\natexlab{c}}.
\bibitem[Masuyama(2016)]{masuyama2016limit}
Masuyama, H.
\newblock Limit formulas for the normalized fundamental matrix of the northwest-corner truncation of markov chains: Matrix-infinite-product-form solutions of block-hessenberg markov chains.
\newblock \emph{arXiv preprint arXiv:1603.07787}, 2016.
\bibitem[Monge(1781)]{monge1781memoire}
Monge, G.
\newblock M{\'e}moire sur la th{\'e}orie des d{\'e}blais et des remblais.
\newblock \emph{Mem. Math. Phys. Acad. Royale Sci.}, pp.\ 666--704, 1781.
\bibitem[Nichol et~al.(2022)Nichol, Dhariwal, Ramesh, Shyam, Mishkin, Mcgrew, Sutskever, and Chen]{pmlr-v162-nichol22a}
Nichol, A.~Q., Dhariwal, P., Ramesh, A., Shyam, P., Mishkin, P., Mcgrew, B., Sutskever, I., and Chen, M.
\newblock {GLIDE}: Towards photorealistic image generation and editing with text-guided diffusion models.
\newblock In \emph{Proceedings of the 39th International Conference on Machine Learning}, volume 162, pp.\ 16784--16804, 2022.
\bibitem[Osler(1970)]{osler1970leibniz}
Osler, T.~J.
\newblock Leibniz rule for fractional derivatives generalized and an application to infinite series.
\newblock \emph{SIAM Journal on Applied Mathematics}, 18\penalty0 (3):\penalty0 658--674, 1970.
\bibitem[Pan et~al.(2023{\natexlab{a}})Pan, Liew, Tan, Feng, and Yan]{pan2023adjointdpm}
Pan, J., Liew, J.~H., Tan, V.~Y., Feng, J., and Yan, H.
\newblock Adjointdpm: Adjoint sensitivity method for gradient backpropagation of diffusion probabilistic models.
\newblock \emph{arXiv preprint arXiv:2307.10711}, 2023{\natexlab{a}}.
\bibitem[Pan et~al.(2023{\natexlab{b}})Pan, Yan, Liew, Feng, and Tan]{pan2023towards}
Pan, J., Yan, H., Liew, J.~H., Feng, J., and Tan, V.~Y.
\newblock Towards accurate guided diffusion sampling through symplectic adjoint method.
\newblock \emph{arXiv preprint arXiv:2312.12030}, 2023{\natexlab{b}}.
\bibitem[Pollini et~al.(2018)Pollini, Lavan, and Amir]{pollini2018adjoint}
Pollini, N., Lavan, O., and Amir, O.
\newblock Adjoint sensitivity analysis and optimization of hysteretic dynamic systems with nonlinear viscous dampers.
\newblock \emph{Structural and Multidisciplinary Optimization}, 57:\penalty0 2273--2289, 2018.
\bibitem[Ramesh et~al.(2022)Ramesh, Dhariwal, Nichol, Chu, and Chen]{ramesh2022hierarchical}
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M.
\newblock Hierarchical text-conditional image generation with clip latents.
\newblock \emph{arXiv preprint arXiv:2204.06125}, 2022.
\bibitem[Rockafellar(2015)]{rockafellar2015convex}
Rockafellar, R.~T.
\newblock Convex analysis:(pms-28).
\newblock 2015.
\bibitem[Rombach et~al.(2022)Rombach, Blattmann, Lorenz, Esser, and Ommer]{rombach2022high}
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.
\newblock High-resolution image synthesis with latent diffusion models.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pp.\ 10684--10695, 2022.
\bibitem[Saharia et~al.(2022)Saharia, Chan, Saxena, Li, Whang, Denton, Ghasemipour, Gontijo~Lopes, Karagol~Ayan, Salimans, Ho, Fleet, and Norouzi]{NEURIPS2022_ec795aea}
Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E.~L., Ghasemipour, K., Gontijo~Lopes, R., Karagol~Ayan, B., Salimans, T., Ho, J., Fleet, D.~J., and Norouzi, M.
\newblock Photorealistic text-to-image diffusion models with deep language understanding.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~35, pp.\ 36479--36494, 2022.
\bibitem[Salimans \& Ho(2022)Salimans and Ho]{salimans2022progressive}
Salimans, T. and Ho, J.
\newblock Progressive distillation for fast sampling of diffusion models.
\newblock \emph{arXiv preprint arXiv:2202.00512}, 2022.
\bibitem[Santambrogio(2015)]{santambrogio2015optimal}
Santambrogio, F.
\newblock Optimal transport for applied mathematicians.
\newblock \emph{Birk{\"a}user, NY}, 55\penalty0 (58-63):\penalty0 94, 2015.
\bibitem[Schuhmann et~al.(2022)Schuhmann, Beaumont, Vencu, Gordon, Wightman, Cherti, Coombes, Katta, Mullis, Wortsman, Schramowski, Kundurthy, Crowson, Schmidt, Kaczmarczyk, and Jitsev]{schuhmann2022laionb}
Schuhmann, C., Beaumont, R., Vencu, R., Gordon, C.~W., Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., Schramowski, P., Kundurthy, S.~R., Crowson, K., Schmidt, L., Kaczmarczyk, R., and Jitsev, J.
\newblock {LAION}-5b: An open large-scale dataset for training next generation image-text models.
\newblock In \emph{Advances in Neural Information Processing Systems (NeurIPS)}, pp.\ 25278--25294, 2022.
\bibitem[Skorski(2021)]{skorski2021modern}
Skorski, M.
\newblock Modern analysis of hutchinson's trace estimator.
\newblock In \emph{2021 55th CISS}. IEEE, 2021.
\bibitem[Sohl-Dickstein et~al.(2015)Sohl-Dickstein, Weiss, Maheswaranathan, and Ganguli]{sohl2015deep}
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S.
\newblock Deep unsupervised learning using nonequilibrium thermodynamics.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 2256--2265. PMLR, 2015.
\bibitem[Song \& Lai(2024)Song and Lai]{song2024fisher}
Song, K. and Lai, H.
\newblock Fisher information improved training-free conditional diffusion model.
\newblock \emph{arXiv preprint arXiv:2404.18252}, 2024.
\bibitem[Song \& Ermon(2019)Song and Ermon]{song2019generative}
Song, Y. and Ermon, S.
\newblock Generative modeling by estimating gradients of the data distribution.
\newblock \emph{Advances in neural information processing systems}, 32, 2019.
\bibitem[Song et~al.(2020)Song, Sohl-Dickstein, Kingma, Kumar, Ermon, and Poole]{song2020score}
Song, Y., Sohl-Dickstein, J., Kingma, D.~P., Kumar, A., Ermon, S., and Poole, B.
\newblock Score-based generative modeling through stochastic differential equations.
\newblock \emph{arXiv preprint arXiv:2011.13456}, 2020.
\bibitem[Song et~al.(2021)Song, Durkan, Murray, and Ermon]{song2021maximum}
Song, Y., Durkan, C., Murray, I., and Ermon, S.
\newblock Maximum likelihood training of score-based diffusion models.
\newblock \emph{Advances in neural information processing systems}, 34:\penalty0 1415--1428, 2021.
\bibitem[Theis et~al.(2015)Theis, Oord, and Bethge]{theis2015note}
Theis, L., Oord, A. v.~d., and Bethge, M.
\newblock A note on the evaluation of generative models.
\newblock \emph{arXiv preprint arXiv:1511.01844}, 2015.
\bibitem[Tu et~al.(2025{\natexlab{a}})Tu, Fu, Yang, Zhao, Zhang, and Qian]{tu2025texttoucher}
Tu, J., Fu, H., Yang, F., Zhao, H., Zhang, C., and Qian, H.
\newblock Texttoucher: Fine-grained text-to-touch generation.
\newblock In \emph{Proceedings of the AAAI Conference on Artificial Intelligence}, volume~39, pp.\ 7455--7463, 2025{\natexlab{a}}.
\bibitem[Tu et~al.(2025{\natexlab{b}})Tu, Ji, Zhao, Zhang, Zimmermann, and Qian]{tu2025driveditfit}
Tu, J., Ji, W., Zhao, H., Zhang, C., Zimmermann, R., and Qian, H.
\newblock Driveditfit: Fine-tuning diffusion transformers for autonomous driving data generation.
\newblock \emph{ACM Transactions on Multimedia Computing, Communications and Applications}, 21\penalty0 (3):\penalty0 1--29, 2025{\natexlab{b}}.
\bibitem[Villani et~al.(2009)]{villani2009optimal}
Villani, C. et~al.
\newblock \emph{Optimal transport: old and new}, volume 338.
\newblock Springer, 2009.
\bibitem[Wallace et~al.(2023)Wallace, Gokul, and Naik]{wallace2023edict}
Wallace, B., Gokul, A., and Naik, N.
\newblock Edict: Exact diffusion inversion via coupled transformations.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pp.\ 22532--22541, 2023.
\bibitem[Wang et~al.(2024{\natexlab{a}})Wang, Yin, Dong, Zhu, Zhao, Qian, Li, et~al.]{wang2024belm}
Wang, F., Yin, H., Dong, Y.-J., Zhu, H., Zhao, H., Qian, H., Li, C., et~al.
\newblock Belm: Bidirectional explicit linear multi-step sampler for exact inversion in diffusion models.
\newblock \emph{Advances in Neural Information Processing Systems}, 37:\penalty0 46118--46159, 2024{\natexlab{a}}.
\bibitem[Wang et~al.(2024{\natexlab{b}})Wang, Zhu, Zhang, Zhao, and Qian]{wang2024gad}
Wang, F., Zhu, H., Zhang, C., Zhao, H., and Qian, H.
\newblock Gad-pvi: A general accelerated dynamic-weight particle-based variational inference framework.
\newblock \emph{Entropy}, 26\penalty0 (8):\penalty0 679, 2024{\natexlab{b}}.
\bibitem[Wu et~al.(2024)Wu, Chen, Li, Wang, and Wei]{wu2024theoretical}
Wu, Y., Chen, M., Li, Z., Wang, M., and Wei, Y.
\newblock Theoretical insights for diffusion guidance: A case study for gaussian mixture models.
\newblock \emph{arXiv preprint arXiv:2403.01639}, 2024.
\bibitem[Yang et~al.(2023)Yang, Feng, Zhang, Shen, Zhu, Huang, Zhang, Liu, Zhao, Zhou, et~al.]{yang2023lipschitz}
Yang, Z., Feng, R., Zhang, H., Shen, Y., Zhu, K., Huang, L., Zhang, Y., Liu, Y., Zhao, D., Zhou, J., et~al.
\newblock Lipschitz singularities in diffusion models.
\newblock In \emph{The Twelfth International Conference on Learning Representations}, 2023.
\bibitem[Yosida(2012)]{yosida2012functional}
Yosida, K.
\newblock \emph{Functional analysis}, volume 123.
\newblock Springer Science \& Business Media, 2012.
\bibitem[Zhang et~al.(2023)Zhang, Lewis, and Kleijn]{zhang2023exact}
Zhang, G., Lewis, J.~P., and Kleijn, W.~B.
\newblock Exact diffusion inversion via bi-directional integration approximation.
\newblock \emph{arXiv preprint arXiv:2307.10829}, 2023.
\bibitem[Zhang et~al.(2024{\natexlab{a}})Zhang, Yin, Li, and Xie]{zhang2024formulating}
Zhang, P., Yin, H., Li, C., and Xie, X.
\newblock Formulating discrete probability flow through optimal transport.
\newblock \emph{Advances in Neural Information Processing Systems}, 36, 2024{\natexlab{a}}.
\bibitem[Zhang et~al.(2024{\natexlab{b}})Zhang, Yin, Li, and Xie]{zhang2024tackling}
Zhang, P., Yin, H., Li, C., and Xie, X.
\newblock Tackling the singularities at the endpoints of time intervals in diffusion models.
\newblock \emph{arXiv preprint arXiv:2403.08381}, 2024{\natexlab{b}}.
\bibitem[Zheng et~al.(2023)Zheng, Lu, Chen, and Zhu]{pmlr-v202-zheng23c}
Zheng, K., Lu, C., Chen, J., and Zhu, J.
\newblock Improved techniques for maximum likelihood estimation for diffusion {ODE}s.
\newblock In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), \emph{Proceedings of the 40th International Conference on Machine Learning}, volume 202 of \emph{Proceedings of Machine Learning Research}, pp.\ 42363--42389. PMLR, 23--29 Jul 2023.
\newblock URL \url{https://proceedings.mlr.press/v202/zheng23c.html}.
\bibitem[Zhu et~al.(2024{\natexlab{a}})Zhu, Wang, Ding, Qu, and Zhu]{zhu2024analyzing}
Zhu, H., Wang, F., Ding, T., Qu, Q., and Zhu, Z.
\newblock Analyzing and improving model collapse in rectified flow models.
\newblock \emph{arXiv preprint arXiv:2412.08175}, 2024{\natexlab{a}}.
\bibitem[Zhu et~al.(2024{\natexlab{b}})Zhu, Wang, Zhang, Zhao, and Qian]{zhu2024neural}
Zhu, H., Wang, F., Zhang, C., Zhao, H., and Qian, H.
\newblock Neural sinkhorn gradient flow.
\newblock \emph{arXiv preprint arXiv:2401.14069}, 2024{\natexlab{b}}.
\end{thebibliography}