Efficiently Access Diffusion Fisher: Within the Outer Product Span Space

Thank you for your appreciation of our work. We will answer your questions one by one regarding these suggestions/questions.

Answer to Suggestion 1: The authors' evaluations are easy to understand and clear, although it would be helpful if more examples (along the lines of Figs 2 and 4) were shown in the SI.

While we consider the theoretical analysis of the diffusion Fisher to be the primary contribution of this work, we agree that additional visual illustrations, such as those in Figures 2 and 4, would further enhance clarity. We will include more examples in the Supplementary Information (SI) as part of our revision.

Answer to Suggestion 2: line ~ 208, "detail discussions" -> "a detailed discussion"

We apologize for this oversight. We have now corrected the syntax error and conducted a thorough proofreading of our work.

Answer to Question 1: The authors' proposed methods involve some approximation (see, e.g., line 375). Are there cases where these approximations are expected to break down, and this method is not expected to give good results?

In the DF-EA method, we introduce an approximation technique for the adjoint term (around line 375). These approximations are likely to be inaccurate in the following scenarios:

• When $t$ is very large, the $x$-prediction part in the approximation is less accurate because of the large noise scale.
• When $t$ is very small, the approximation error bound stated in Proposition 6 will increase due to the singularity issue, and the approximation will lack an accuracy guarantee. We will discuss on this singularity issue in detail in Question 3

In practice, we apply DF-EA in the range from $t=0.7$ to $t=0.3$. This range also aligns with the recommended adjoint region of the original adjoint methods.

Answer to Question 2: In general, are there any obvious issues associated with using or scaling up these methods for computing the diffusion Fisher (or certain related quantities, like its trace)?

• Obvious issue: Our methods are derived exclusively on the classical diffusion process and thus cannot be directly extended to settings where the diffusion process is modified, such as cases involving additional consistency-model-type distillation [1], Schrödinger-Bridge-type endpoint alignment [2], non-Gaussian transition kernels [3], or flow rectification [4].

• Scaling-up: Theoretically, both the training and sampling phases of our ​DF-TM and ​DF-EA exhibit ​linear complexity growth with respect to dataset scale and data dimensionality. Experimentally, our approach performs well on commercial-grade diffusion models such as SD2-base and large-scale datasets like Laion-2B. Therefore, our method indeed has scalability.

[1] Frans, Kevin, et al. "One step diffusion via shortcut models." ICLR 2025.

[2] De Bortoli, Valentin, et al. "Diffusion schrödinger bridge with applications to score-based generative modeling." NeurIPS 2022.

[3] Yoon, Eun Bi, et al. "Score-based generative models with Lévy processes." NeurIPS 2023.

[4] Liu, Xingchao, et al. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. ICLR 2022.

Answer to Question 3: As a related point, Proposition 6 appears to show potentially bad behavior when $\sigma_t \approx 0$.

Note that Proposition 6 will be ill-defined due to division-by-zero when $t=0$ since $\sigma_0 = 0$. This concern is also raised by Reviewer V5LL. We argue that this problem, which has been referred to as the singularity issue in previous literature [1][2], stems from the inherent characteristics of diffusion models rather than flaws in our derivations.

Theoretically, the diffusion model has an inherent singularity property when $t_{\min} \to 0$. This implies that many quantities associated with diffusion models become ill-defined in the vicinity of $t_{\min} = 0$. This phenomenon, previously explored in [1] and [2], is once again recognized in the context of diffusion Fisher through our analysis. A deep theoretical exploration of the singularity problem remains an open question within the diffusion model field.

Practically, we implement the approach used in [3] and [4], where the diffusion sampling path is ended at $t_{\min} = 0.005$ instead of $t_{\min} = 0$. This method circumvents potential bad effects resulting from the singularity issue.

We briefly discussed this singularity issue on line 1242 of the original paper. In our revision, we will provide a more detailed discussion and clarification of our handling of this issue.

[1] Zhang, Pengze, et al. "Tackling the singularities at the endpoints of time intervals in diffusion models." CVPR 2024.

[2] Yang, Zhantao, et al. "Lipschitz singularities in diffusion models." ICLR 2024.

[3] Song Y, Dhariwal P, Chen M, et al. Consistency Models ICML 2023.

[4] Lu, Cheng, et al. "Maximum likelihood training for score-based diffusion odes by high order denoising score matching." ICML 2022.

Thank you for your valuable feedback! We will answer your concerns/questions one by one.

Concern 1: I have a problem with the paper's main claim that JVP cannot be calculated effectively, they claim that the time complexity is $O(d^2)$...

Please allow me to humbly clarify that we didn't claim a single JVP's time complexity to be $O(d^2)$. Instead, we stated that calculating the diffusion Fisher trace using JVPs has an $O(d^2)$ complexity. The reason is that to obtain each diagonal element of the diffusion Fisher matrix, one JVP operation is needed. As computing the full trace demands all $d$ diagonal elements and each JVP has $O(d)$ complexity, the total complexity for $d$ JVPs is $O(d^2)$.

In the context of accessing the diffusion Fisher trace, previous methods rely on $d$ JVP operations. In contrast, our DF-TM method only requires one forward pass through the learned trace network to access the diffusion Fisher trace, which is notably more efficient than the traditional multiple-JVP approach.

In the adjoint scenario, previous methods use a JVP to compute the adjoint term. Our DF-EA approximation method eliminates this JVP, making it marginally more efficient than the traditional method.

Furthermore, accessing the diffusion Fisher via JVP is a black-box approach without any accuracy guarantee. Conversely, our DF-TM and DF-EA methods exploit the analytical structure of the diffusion Fisher and offer theoretical guarantees.

Concern 2: how about approximating it with numerical methods?...

While numerical methods for JVP and trace evaluation are inaccurate and inefficient in high-dimensional scenarios, as discussed in line ~1281. We appreciate your suggestion and agree that incorporating these as baselines enhances the clarity of our experiments. So, we've added two numerical methods, Finite Difference and Hutchinson's Trace Estimation, as baselines in the toy NLL evaluation experiment:

Table 2: The relative error of NLL evaluation.

It is shown that our DF-TM outperforms these numerical methods.

Concern 3: Training a separate nontrivial network for the use cases listed in the paper feels too heavy...

• In adjoint applications, our DF-EA is training-free.

• In NLL evaluation, our DF-TM incorporates an additional trace matching network, only bringing moderate costs:

  • The trace network is scalar-valued and demonstrates a much faster convergence compared to the main network, as shown in Fig. 1. In Laion-2B dataset, the main network of SD-1.5 needs 100,000 GPU hours, while the trace network only 3000, just 3% of the main network's training cost.
  • The trace network shares the same input as the diffusion network, allowing it to use a shared backbone of $\epsilon_\theta$ for feature extraction, and only necessitates the training of an output head. We plan to explore and implement this promising approach in our future research.

  Thus, DF-TM can leverage off-the-shelf pretrained diffusion models and only needs to train a small trace network as a "plug".

Question 1: Is it true that the prior methods all use the highly inefficient way of calculating the JVP...

In prior adjoint methods [1][2], a JVP operation is required at each step. Usually, these methods use the official PyTorch auto-differentiation tools to compute the JVP. Our DF-EA approximation method circumvents the need for this JVP calculation, though the efficiency gain is limited.

The main advantage of DF-EA in adjoint optimization lies in its accuracy, leading to better scores and visual effects, as demonstrated by the qualitative comparison in Fig. 4 and the quantitative comparison in Fig. 3.

In the revision, we'll add more details on prior adjoint methods and their applications.

[1] Pan, Jiachun, et al. "Towards accurate guided diffusion sampling through symplectic adjoint method." arXiv preprint arXiv:2312.12030 (2023).

[2] Blasingame, Zander W., and Chen Liu. "AdjointDEIS: Efficient gradients for diffusion models." NeurIPS 2024.

Thank you for your valuable feedback! We will answer your weaknesses/questions one by one.

Weaknesses 1: The bound in Proposition 5 is vacuous when t is small...

It is true that as $t$ approaches $0$, the bound in Proposition 5 will be ill-defined and blow up due to division-by-zero. This concern is also raised by Reviewer a1EQ. We argue that this problem, which has been referred to as the singularity issue in previous literature [1][2], stems from the inherent characteristics of diffusion models rather than flaws in our derivations. A deep exploration of the singularity problem remains an open question, which we will investigate in our future works.

Practically, we adopt the approach in [3], where the diffusion sampling path ended at $t_{\min} = 0.005$ instead of $t_{\min} = 0$. This circumvents potential blow-up effects of singularity issues and results in a well-defined bounded NLL.

We briefly discussed this singularity issue on line 1242 of the original paper. In our revision, we will provide a more detailed discussion and clarification of our handling.

[1] Zhang, P, et al. Tackling the singularities at the endpoints of diffusion models. CVPR 2024.

[2] Yang, Z, et al. "Lipschitz singularities in diffusion models." ICLR 2024.

[3] Song, Y, et al. Consistency Models ICML 2023.

Weaknesses 2: In the Dirac setting, the likelihood evaluation becomes trivial...

We derive the theory within the Dirac setting to simplify the formulations and make them easier to follow. We also present the derivations under the general setting in the appendix. Remarkably, both the Dirac and general settings yield the same forms for the DF-TM and DF-EA methods.

All derivations within the Dirac setting can be straightforwardly adapted to the general setting by replacing discrete summations with general integrals. The underlying ideas and mechanisms are identical. The Dirac case can be regarded as a collection of finite Monte-Carlo samples drawn from a general distribution. In the mean-field limit, the Dirac case converges to the general case.

To prevent misunderstandings, we will revise all propositions in the main body of the paper to the general setting.

Weaknesses 3: In adjoint task, the visual quality...

In Fig. 4, the DF-EA method contributes to (left) a higher Laion-Aes score, as evidenced by a smoother visual effect, and (right) a higher pick-score, reflected in enhanced colorfulness and more detailed features. Since visual effects can be subjective, we also present a quantitative demonstration of the score improvement brought about by DF-EA in Fig. 3.

Suggestions: Line 145...Line 153...Eq. (7)...

We apologize for the oversights. We have now corrected the typographical errors and adopted standard Hessian notation for Eq. 7: $\left(F_t(x_t,t)\right)_{i,j} := - \frac{\partial^2 \log q_t(x_t, t)}{\partial x_t^{(i)} \partial x_t^{(j)}}$, where $x^{(i)}$ denotes the i-th element of vector $x$.

Question 1: ...Can the authors extend Proposition 4,5,6 for this setting?

Yes, our Proposition 3,4,5,6 can be extended to the general setting and result in the same DF-TM and DF-EA method. Suppose we have $m$ samples from a general initial distribution $q_0$.

• Proposition 3 (general version): The trace of the diffusion Fisher for the diffused distribution $q_t$, where $t\in(0,1]$, is given by: $\mathrm{tr}\left(F_t(x_t,t)\right) = \frac{d}{\sigma_t^2} - \frac{\alpha_t^2}{\sigma_t^4}\left[ \int w(y) ||y||^2dq_0(y) - \left\|\int w(y) y dq_0(y)\right\|^2 \right]$

• Proposition 4 (general version): $\forall (x_t, t) \in \mathbb{R}^d \times \mathbb{R}^{+}$, the optimal $t_\theta(x_t, t)$ trained by Algo. 1 are equal to $\frac{1}{d}\int w(y) ||{y}||^2 dq_0(y)$.

• Proposition 5 (general version): The approximation error of DF-TM is at most $\frac{\alpha_t^2}{\sigma_t^4}\delta_1 +\frac{1}{\sigma_t^2}\delta_2^2$.

• Proposition 6 (general version): The approximation error of the DF-EA linear operator, is at most $\frac{\alpha_t^2}{\sigma_t^3}(2\mathcal{D}_y^2+ \sqrt{d}\delta_2)$

Note that, in Proposition 5 & 6, the approximation error induced from the finite number of $m$ samples is absorbed into $\delta_1$ and $\delta_2$. We'll include the general-version propositions and their proofs in the revision.

Question 2: In the toy problem, how do you compute the ground truth NLL from the dataset?

In the toy experiment, we adopt the method in [1] and approximate the true NLL following Eq. 13. We use Euler discretization over 1000 timesteps for simulation and 5000 data samples per timestep to approximate the right-hand side term of Eq. 13. This setup aims to ensure NLL approximation accuracy.

We apologize that the description in line ~1096 is rather simplistic. In the revision, we will provide a more comprehensive description.

[1] Lu, Cheng, et al. Maximum likelihood training for score-based diffusion odes by high order denoising score matching. ICML 2022.