Annealed Fisher Implicit Sampler
Authors: Weijian Luo
Boya Zhang
Zhihua Zhang
Venue: Submitted to ICLR 2023 (2023)
Keywords: Implicit Generative Models
Score Matching
Learning to Sample
Sampling

Abstract: Sampling from an un-normalized target distribution is an important problem in many scientific fields. An implicit sampler uses a parametric transform $x=G_\theta(z)$ to push forward an easy-to-sample latent code $z$ to obtain a sample $x$. Such samplers are favored for fast inference speed and flexible architecture. Thus it is appealing to train an implicit sampler for sampling from the un-normalized target. In this paper, we propose a novel approach to training an implicit sampler by minimizing the Fisher Divergence between sampler and target distribution. We find that the trained sampler works well for relatively simple targets but may fail for more complicated multi-modal targets. To improve the training for multi-modal targets, we propose another adaptive training approach that trains the sampler to gradually learn a sequence of annealed distributions. We construct the annealed distribution path to bridge a simple distribution and the complicated target. With the annealed approach, the sampler is capable of handling challenging multi-modal targets. In addition, we also introduce a few MCMC correction steps after the sampler to better spread the samples. We call our proposed sampler \emph{the Annealed Fisher Implicit Sampler} (AFIS). We test AFIS on several sampling benchmarks. The experiments show that our AFIS outperforms baseline methods in many aspects. We also show in theory that the added MC correction steps get faster mixing by using the learned sampler as MCMC's initialization. 


Under review as a conference paper at ICLR 2023
ANNEALED FISHER IMPLICIT SAMPLER
Anonymous authors
Paper under double-blind review
ABSTRACT
Sampling from un-normalized target distribution is an important problem in many scientific fields. An implicit sampler uses a parametric transform x = G=) to push forward an casy-to-sample latent code to obtain a sample z. Such samplers are favored for fast inference speed and flexible architecture. Thus it is appealing to train an implicit sampler for sampling from the un-normalized target. In this paper, we propose novel approach to training an implicit sam- pler by minimizing the Fisher Divergence between sampler and target distribu- tion. We find that the trained sampler works well for relatively simple targets but may fail for more complicated multi-modal targets. To improve the training for multi-moda targets, we propose another adaptive training approach that trains the sampler to gradually learn a sequence of annealed distributions. We construct the annealed distribution path bridge simple distribution and the complicated target. With the annealed approach, the sampler is capable of handling challeng- ing multi-modal targets. In addition, we also introduce few MCMC correction steps after the sampler to better spread the samples. We call our proposed sampler the Annealed Fisher Implicit Sampler (AFIS). We test AFIS on several sampling benchmarks. The experiments show that our AFIS outperforms baseline methods in many aspects. We also show in theory that the added MC correction steps get faster mixing by using the learned sampler as MCMC's initialization.
INTRODUCTION
Sampling from an un-ooraalized distribution an important problem in many scientific fields such as Bayesian statistics (Green, 1995). biology (Schütte et al. 1999). physics simulations (Olsson 1995]. machine learning (Andrieu al. 2003), and so on. Typically, the problem is formulated as: given a known differentiable un-normalized target potential function log p(r), one wants to sample from the target distribution. Due to the success of deep neural networks, there is increasing popularity to train deep generative model to learn to sample(Hu et al.) 2018 Wu al., 2020) Matthews et al. 2022 Corenflos et al. (021). Such learned models which can approximately sample from target distribution are called samplers.
Training a neural network (i.c., a parameterized transform) x = =G((2) to push forward an easy- to-sample latent code pz(=) to obtain a sample is appealing approach. Such approaches are favored for fast sampling because they only need a single-time forward pass of neural network transform. Let (.). denote the parametric transform and q(x) the un-normalized target distribu- tion with unknown normalizing constant Z q(x)dd. Let pa p=(x) denote the sampler-induced distribution. Some previous work takes normalizing flow model as sampler, and then minimizes the KL divergence between sampler-induced and target distributions regardless of normalizing con- stant: Dxz (Pe. q) = Ex--po log log g (xx) + log Z]. Note that Z is parameter- free and can be ignored during training. However, minimizing KI. divergence relies on explicit log-likelihood of compler-induced distribution, which can not be computed in a general transform. Such transform with explicit likelihood is referred to as implicit sampler.
In this paper, we will focus on implicit samplers. Note that the annoying normalizing constant vanishes when considering the score function of a distribution, s(x) V, log p(x). Thus, we can take the score-based divergence to constructively get rid of the unknown normalizing constant for implicit samplers. Fisher divergence (FD), which is a popular based probability divergence, and its variants have obtained much success in recent years, especially in training deep generative models such energy- energy-based based models (Kingma & Cun 2010 Martens et al. 2012 Song al. 2019.

---

Under review as a conference paper at ICLR 2023
score based diffusion models (Song et al. 2020 Kingma et al. 2021 Vahdat et al. 2021 Song & Ermon 2019 Ho et al. (020), etc. Assume p(z), q(z) are two probability densities. The Fisher
Divergence between p and q defined
1 DFD(P. = - V logp(z) V log (2)))
is always no less than O and equals to 0 if and only if p(z) = g(x) a.s. under probability measure Fisher Divergence is suitable for measuring the dissimilarity between sampler and nn-nommalized target distribution. So as to be used for training the implicit sampler.
In this paper, firstly propose novel approach to learning sampler by minimizing the Fisher Divergence between sampler and un un-normalized target distributions. We call such a sampler the Fisher Implicit Sampler. We then show that the proposed sampler is capable of handling relatively simple target distribution, but would fail for more challenging multi-modal targets.
To remedy this issue and unlock the full potential of the Fisher Implicit Sampler, we additionally propose novel adaptive training approach that trains the implicit sampler gradually using a se- quence of annealed distributions instead of the target distribution. We anneal the target distribution to bridge the hard-to-sample target and an easy-to-sample prior. More precisely, we extend the target distribution (xx) to a sequence of annealed distributions {qu (x) for .0....,K where qK (x) is the target density and 9u(x) (x) is an easy-to-sane tompal prior distribution. typically normal distribution. The design of such an annealed path gradually reduces the learning difficulty for the sampler.
Learning Objective
( (G,00)))))
Sampler
z- P.(2)
Hard
Easy to Leam
Annealing Path
(e) Aglooggrs
Gredual Learning
N(O.
Figure 1: Illustration of proposed Annealed Fisher Implicit Sampler.
Moreover, we find that a few steps of MC correction after the sampler help the samples spread better with little cost, as also used in some previous work (Wu et al., 2020 Arbel et al. [202] Matthews
al. 2022). Combining all together, we call our proposed sampler the Annealed Fisher Implicit Sampler (AFIS), as illustrated in [igure We validate our AFIS sampling benchmarks, showing improvements over baseline approaches.
The main contributions of our work are summarized as follows:
We propose a novel loss function to minimize the Fisher Divergence. We show that mini- mizing the proposed loss is equivalent minimizing the Fisher Divergence between sam- pler and target distribution. Note that our objective largely different from other ones in previous work.
We provide an insightful understanding of the difficulty in learning multi-modal targets by minimizing Fisher Divergence. We facilitate the annealing technique on training samplers based on our understanding.
We bring a novel annealing technique and MC correction steps with our sampler, leading improved sampling performance with little additional cost.
2 BACKGROUND
2.1 TRAIN IMPLICIT SAMPLERS WITH SCORE-BASED DIVERGENCE
The learning-to toreg-tt--annd -ample problem arises in many application fields of machine learning. Assume only have access to un-normalized target distribution q(x:) (or its logarithm log q(x)), and the goal is to approximately sample from the target. In recent years, training a neural network- based transform approximately sample from target distribution is appealing method. Such
2

---

Under review as a conference paper at ICLR 2023
transform is called a neural sampler. Let Go denote neural network which transforms relatively simple latent code ~ Po(2) to a sample x = Gø(z). Here, (2) is an casy-to-sample latent distribution, usually the standard Normal distribution. A general neural sampler does not have an explicit expression of the log-likelihood function, which we name them implicit samplers. Because of the un-normalized target distribution and unavailable log-likelihodd, training implicit samplers by minimizing KL or related divergence always fails. An alternative way is to consider score-based divergence.
The Stein Neural Sampler Hu et al. (2018) is trained by minimizing Stein's Discrepancy between sampler and target distributions. The Stein Discrepancy (SD) (Gorham & Mackey 2015) is defined as
(pn g) = sup En-p V2 log (z), f(r)) (V., f(x))
SeF
The calculation of Stein's discrepancy relies on solving maximization problem w.r.t test function f. When the function class F is carefully chosen, the optimal f may have an explicit solution or casier formulation. For instance, Hu et al. (2018) found that F is taken to be F (f: E, [f] S
5}, the SD is equivalent to regularized representation
Dsn 4) max Emep V. log ((x), f(x)) + (Va, f(x)) - - )['
f
They used two neural networks: Ge to parametrize an implicit sampler and f, to parametrize the test function. Let po(z) denote the implicit sampler distribution induced by Go(2) with Pz (z). Stein Neural Sampler solves a minimax problem on parameter pair (0, 7)) to obtain a sampler that minimizes the SD between sampler and target by
min max L(O, = min max Exzpo a log g(x), (x)) (V (,(,) X[1 m
7
0
Here the notion ~ pe means = Ga(z) with 2 pz(2). They called the above SD the Fisher Stein Discrepancy and the corresponding sampler FSD Neural Sampler.
The Stein Neural Sampler opens the door to training implicit samplers by minimizing score-based Divergence. In fact, the FSD Neural Sampler calculates surrogate of Fisher Divergence. The FSD's test function f provides an approximation of Fisher Divergence. However, as we show in Section their calculation of Fisher Divergence only provides partial gradient updates of the sampler's parameters, thus leading to training failure even for simple target.
3.1
2.2
SCORE FUNCTION ESTIMATION
Since the implicit sampler does not have an explicit log-likelihood function or score function, train- ing with score-based divergence requires inevitably estimating the score function (or equivalent component). Score matching (Hyvärinen & Dayan 2005) and its variants provided powerful ap- proaches to estimating score function through samples. Assume one only has available samples & ~ p. and wants to use a parametric approximated distribution (x(x) to approximate p. Such an approximation can be made by minimizing the Fisher Divergence between p and qg- We can rewrite the Fisher Divergence as
DFD(p. (6) = Egwp { VV (o))()))) + V log (2) 2(V logp(x)) Ve log (p(rr)) }
Under certain conditions, the equality Ezzp V (x)) V log qo (x)) = E Exzpp log qo (x)
holds (usually referred to as Stein's Identity(Stemn 1981 Gorham & Mackey (017)) Here
A log (x) E. d log (x(x) denotes the Laplacian operator applied on log q (a). Combining
this equality and noting that the first term of FD Ez~pl V log p(2) does not rely parameter 6, we have that minimizing D (p. 90) is equivalent to minimizing the following objective
[(Φ) Ezup
||v log (z) 2A log (
This objective can be estimated only through samples from p. thus is tractable when qa well- defined. More specifically, one only needs to define score network s5(x): RDD RD instead of
3

---

may wish minimize the Under review as a conference paper at ICLR 2023
density to estimate the score function of p in some cases. This technique was proposed in Hyvärinen & Dayan 2005) named after Score Matching, Other variants of score matching were also studied
(Song a.. 2019 Vincent 2011 Pang et al. 2020 Meng et al., 2020; Lu et al. 2022 Bao et al. 2020). Score Matching related techniques have been widely used in training energy-based models and score-based diffusion models in recent years. In this paper, use score matching related techniques to estimate the score function of the sampler's distribution.
ANNEALED FISHER IMPLICIT SAMPLER
3.1
MINIMIZING THE FISHER DIVERGENCE: S2D Loss
Let Ga(.): D2 + Rx be an implicit sampler (i.e., a neural transform), pz latent distribution, sampler induced distribution x Go(z), and ((x:) un-normalized target. Our goal is pull close the FD between Po and in order to train the sampler. Recall the definition of Fisher Divergence between pa, q is
DFD g) Ez-pe VV log.olop (1)) - V. log g(x))12
For our learning-to- seariinggtt--aapppep seems work one uses an additional score network ss(.): R Px - px to approximate sampler's score function. Samples from implicit sampler is cheap to obtain, so estimating sampler's ss score function is not hard with score matching related techniques. We call this step the Score Estimation Step. With good approximated sy(2 (I) of sampler's score function, approximated Fisher Divergence to update the sampler
sample setting, the target score function V, log 7(x) is known. A direct solution
0" = arg min (x)))))))) - V log (()))
We call this step the Score Difference Minimization Step. By alternating the above two steps, one may wish the Fisher divergence will be minimized, thus the training sampler is done. We name the resulting approach the Direct Method. Interestingly, the Direct Method coincides with FSD Neural Sampler as we state in Proposition [ Proposition Estimating the sampler's 's score function sa(.) with score matching is equivalent to maximizing the Fisher Stein Discrepancy objective to obtain FSD's optimal test function. More specially, the optimal score estimation 8" and FSD optimal test function I satisfy
We put detailed proof in Appendix dud to limited pages.
(*(x)
V log g(x))(()))))() - sx)]
2
Moreover, the Direct method is equivalent FSD when training implicit Sampler.
Although the direct method seems reasonable, fails as we show in the experiment nana target in Figure 2 We find that even sam- pler's score function is estimated perfectly at each iteration, the direct method still gives only partial parameter gradient for minimizing the Fisher Diver- gence. We start by analyzing Fisher Divergence's gradient w.r.t. sampler's parameter. The Fisher Di- vergence is
simple Ba-
(a) real
(b) S2D
(c) Direct
Figure Direct method fails for simple Ba- nana distribution while S2D loss succeeds.
(DD(0)= Ex V log ((x) V. log po (1)
One wants to adjust to minimize (0) The gradient of the above objective writes
0e Ep. -0ll=s((==s(o]])(()) = Ep, (ox))) log Po ( (o. (((())))))) 2
0e so()).
00
The first gradient term coincides with the direct approach if we asynchronously estimate the sam- pler's score function perfectly. More precisely, with perfect score estimation s5(x) Va po(),, we have
0
E Eumppo V log g(x) s (x)11
00
/ V log (xx) (0( (1(00)1(p [1V, (1) 81(1))1 á
2g (x)
00
/ | lorg(ar) (x)) 12p0 00
θ
log po po (x) Exeapo I1 V log q(x) ) 8 (x))1 2
log (x)).. 00
4

---

Under review as a conference paper at ICLR 2023
The above equation reveals that the direct method only takes partial gradient to minimize the FD between sampler and target. In many cases, this partial gradient leads to training failure as we observe in Figure 2 In (Hu et al., 20118), FSD Neural Sampler used Kernelized Stein Discrepancy trained implicit sampler as initialization before training with FSD. However, such initialization limits the usage of FSD because the optimization might start from local minima which is close to KSD's local minima and can potentially be mislead the sampler.
In order to minimize the Fisher Divergence correctly, we propose novel training objective called Score Square Difference loss (S2D) which accounts for the full parameter gradient to minimize the Fisher Divergence. The S2D loss is defined as the difference of target and sampler's square score norm, where the sampler's ss score function is estimated asynchronously with score network sz(r). More precisely, our S2D loss is defined as
(0) = Rarpp Vv log g())1 (1,
where (-) is the estimated score function sampler distribution. The score function is usually estimated by score matching related techniques. The notation Pg means x Go(=), pz(z). The following proposition
_ shows that, if the sampler score function is estimated perfectly, the parameter gradient of S2D loss is the same the gradient of Fisher Divergence.
Proposition 2. Assume sœ(x) Vz loogog (x)). Then the following equality holds:
θ 30 (FD(0).
Ls2n (0) 00
We give the detailed proof in Appendix This proposition says that, if we alternate between score estimation of sampler's score function, and minimization of the S2D loss, we are actually minimizing the Fisher Divergence between sampler and target. The S2D loss a surrogate Fisher Divergence. So
of Fisher Divergence which can provide the same parameter gradient minimizing the S2D) loss gives the same results as minimizing the intractable Fisher Divergence. Figure 3 gives illustration of the relation between S2D) loss and Fisher Divergence. The black curve stands for the intractable Fisher Divergence. Green curve represents the S2D loss. The S2D loss shares the same gradient parameter as Fisher Divergence. We refer to a sampler trained with such approach the Fisher Implicit Sampler (FIS). We give an algorithm for FIS in Algorithm [ We take standard score matching as an illustration of score estimation step, but other score estimation techniques such as denoising score matching and sliced score matching also works.
Algorithm 1:
Fisher Implicit Sampler training
Input: un-normalized target log g(x), latent distribution Pz(z), implicit sampler Gg, score network Sg, mini-batch size B, max iteration M.
Randomly initialize (O(W), ()))).
for l in 0:M do
# update score network parameter
Get mini-batch 2, Ggco (zi), ~ p2(2), 1,.., B. Calculate score matching objective:
(SM(
B [|>0(Z))| -
2(V 2( so(2:))
Minimize (6)()) to get dl++1). # update sampler parameter
Get mini-batch latent code 2; ~ (2), 1,. Use re- re-parametrization trick to calculate S2D loss for sampler
(o) = 1/ E.
V log ((Go(2z))) 1| || (Go(z)))
Minimize L32D (0) to get ((++1),
end
return ().
5

---

Under review as a conference paper at ICLR 2023
Figure 2 shows that proposed FIT (S2D loss) can successfully train an implicit sampler from scratch to sample from the famous banana shape distribution. While the Direct method fails to train the correct sampler.
Loss
S2D loss 4yzp (0)
Same Gradient
Although FIT is capable of handling bench- mark targets, find that FIT fails on more challenging multi-modal targets with very sep- arated modes. To remedy the multi-moda fail- ure issues and fully unlock the potential of the S2D loss, we propose to combine the anneal- ing techniques with FIT for multi-modal tar- gets. The idea of annealing widely used in sampling and stochastic optimization litera- ture (Neal 2001] [Salimans et al., 2015 Chen
Fisher Divergence Lyp (
0
optimal parameter 0*
current parameter
Figure 3: S2D loss and Fisher Divergence. The S2D loss shares the same parameter gradient as Fisher Divergence if sampler's score is estimated perfectly asynchronously. Thus minimizing the S2D loss to update the sampler is equivalent to minimizing the Fisher Divergence between sam- pler and target.
ct al. 2016 Doucet 2001 Van Laarhoven & Aarts 1987). The technique constructs a distribution bridge between relatively sim- ple prior distribution and a complicated target. The learning (or other operations such as sam- pling or optimization) are gradually operated on each middle distribution from prior the tar- get. Typically, the annealing technique can lower the barrier of operation of the target by dispersing the difficulty to all middle distributions.
3.2 ANNEALED FISHER IMPLICIT TRAINING
By executing FIT steps repeatedly, the sampler is trained to minimize the Fisher divergence between Pe and target However, directly minimizing the Fisher divergence is problematic in practice. If the sampler's distribution is too dissimilar to the target, the Fisher divergence could be hard estimate accurately as mentioned in (Wenliang & Kanagawa 2020). The Fisher divergence can be small under any tolerance even if two distribution are largely different in terms of KI. divergence. More precisely, the Fisher Divergence is likely inaccurate if two distributions are too dissimilar. Due to this issue, the sampler might not be able to estimated the Fisher divergence accurately, making the training fail. In fact, above issue occurs lot in real applications. Sampler is often initialized to concentrate around the origin, while the target distribution rarely concentrates around the origin.
To remedy the inaccurate Score Estimation issue, need to guide the sampler to start from learning a relatively simple target, and then the more challenging one. Based on such intuition, we introduce gradual relaxation of target distribution. More precisely, distributions {gk)- 1 E (0,... K} which gradually transform a relatively simple distribution Yo to target distribution gK q. Typically, go chosen as A/ 1) for simplicity. We let the sampler gradually learn to sample from each with & increasing from = to k K. Since when k small 4s simpler than Øk. the estimation of Fisher divergence is easier. Thus the sampler can learn to approximate qs- When one gradually turns & to k = K, the sampler will gradually learn to sample from our final target 9K =
construct sequence of annealed
Such hard technique is commonly known as annealing techniques. Marinari & Parisi [1992); Geyer Thompson (1995) proved faster mixing time with temperature annealed target. Wenzel et al. (2020) utilized anneal path to connect model and posterior in Bayesian inference regime. Mandt et al. (0166); Huang et al. 22018: Fu et al. (2019) annealed the KL regularization in variational inference. D' Angelo & Fortuin (202 proposed to anneal the target when running Stein Variational Gradient Descent algorithm for better mixing speed. Perhaps the most similar anneal- ing approach to ours is Neal 2001); Wu al. 2020) which construct a geometric distribution path pa. (x) between a Gaussian prior and target density. We utilize similar anneal path as Wu et al.
2020

---

Under review as a conference paper at ICLR 2023
In this paper, we anneal the target distribution q with geometric interpolation starting with stan- dard Gaussian distribution as prior
log Q (x) = A log 9x (z) + (1 A4) log (o(x)
with go (0, I) and 0 < A I 1 pre-defined annealing schedule function with A = AK 1. The score function is then linearly interpolated with
V (x V log (x) + (1 2 V. log (o()),
where V, log QK (x) V (o) ) is the target score function and v, log (a()) the prior score. For standard Normal prior, we have 7 log 9o(x) -r. We name our FIS sampler combined with annealing technique the Annealed Fisher Implicit Training. Because the pages limitation, put the full AFIS algorithm in Appendix F By annealing the target distribution to sequence of easier-to-learn targets, we divide the difficulty of sampler to learn one final distribution to learn sequentially from less difficult targets. Thus the sampler will not be bothered by inaccurate Fisher divergence estimation and training failure. Figure ggives Annealed Fisher Implicit Sampler is trained along annealed distributions progressively.
gives a brief summary of how AFIS works. The
3.3 MONTE CARLO CORRECTION
Deterministic sampler suffers from mode-connection issue. The issue says that a deterministic trans- form can not fully disconnect two modes as studied in Wu al. (2020). Such issue limit the use of a pure deterministic sampler. Recent works show that combining stochastic corrections with de- terministic transforms could improve the sampling performance (Wu et al. 2020; Song al. 2020 Song & Ermon 2019). MCMC(Hassings 1970; Roberts & Rosenthal 1998] Xifara et al. 2014 Neal 2011 1) a commonly used stochastic transform family. By running MCMC, one can approxi- mated sample from some un-normalized target distribution. Thus few steps MCMC is nice way to serve as stochastic corrections.
In particular, after training the sampler, we take the generated samples Ga(2), po(2) as initialization and run several MCMC as correction steps to spread samples for better diversity. Both energy-based and score-based MCMC can be used. We take the Langevin MC as illustration and put more details of MC corrections Appendix corrections.
C Note that our method is not limited to these MC
Langevin Dynamic Correction
particles assumed reach q(x) as stationary distri- bution if it is driven by Langevin Dynamic with local updates
dX, Vx, , log (XX.)/2 dW.,
where W, is standard Brownian motion. The discrete scheme of Langevin Correction is given by
= = + V Vlogy(X())) + (ez(!),
where Z NO I). The Fokker-Planck equation tells that under certain conditions, g(x) is the only stationary distribution of above diffusion dynamic. About 20 updates of steps is sufficient to have good enough correction effects in practice.
The combination of deterministic sampler and stochastic correction in fact gives faster mixing for MCMC. The deterministic sampler sample particles coarsely target's high density modes. After that, the MCMC helps the particle spread better around each modes. In particular, we show that Langevin mixing time can be controlled by Fisher divergence between sampler distribution and target. Taking advantage of flexible neural network architecture, AFIS can be trained to match target score any precision. The Theorem shows that Langevin Correction's 's mixing time can be reduced by well trained sampler.
Theorem 1. Assume the larget potential log g(z) is smooth and satisfies
||V log (
A log g(xx) +00.
lim +oo
2

---

Under review as a conference paper at ICLR 2023
(a) real
(b) AFIS+MC
(c) AFIS
(d) FIS
(e) FSD-NS
Figure 4: Sample comparison on Double Well targets. (a) real samples; (b) samples from trained AFIS with steps of HMC correction; (c) samples from trained AFIS; (d) samples from trained FIS without annealing; (e) samples from trained FSD-NS, All samplers and score networks use the same architecture.
Assume generated distribution induced by AFIS G(z) is trained match Fisher divergence under precision Dr(p.q) S 6. Then there exists a positive constant X and a dimension-free pos- itive constant C which only depend on target distribution q(x)), such that under Langevin diffusion with initial distribution pa p.
dX,- = V log q(t)) /2d1 + dW..
the diffusion time
1 [ log(( log 2 = max 0. 2X
is enough to control the KL divergence between corrected distribution and target under toler- ance e.
In practice. the AFIS can be trained to achieve any precision to match target under Fisher Divergence. The above theorem says, the better AFIS is trained, the shorter time for MC correction is needed to achieve same tolerance in terms of KI. divergence. We provide the detailed proof in Appendix
3.4 COMBINING ALL. THE ANNEALED FISHER IMPLICIT SAMPLER
Combining the S2D loss, the annealed technique, and MC corrections, we obtain our final sampler: the Annealed Fisher Implicit Sampler (AFIS) with MC corrections. Figure 4 shows comparison of trained sampler's samples on Double Well distribution. Double Well is usually used viriite testing target with two separated modes. The figure shows that the AFIS with a few steps of MC correction gives the best samples. The AFIS with no MCMC correction can not fully separate two disjoint modes. The FIS (without annealing) fails to learn the two modes. The FSD-NS (or the Direct Method) also fails for training. To be concluded, the experiments show that S2D loss, annealed technique, and MC correction all contribute to successful learning.
4 EXPERIMENTS
4.1 AFIS FOR SYNTHETIC TARGET
For sanity check, we apply AFIS on some toy target distributions used in Hu et al. (2018) Rezende
& Mohamed 2015). The anneal path px (z) exp(A log (z)) + (1 X) log Pyrior (z) starts and ends with the target when À
from a Normal distribution when of max iterations, and be the current training iteration. We set X; to grow linearly from 0 to I when i < 9M/10. We train the sampler with real target log ((x) for rest M/10 iterations. The annealed path reduces the bar of learning to sample, resulting relatively accurate updating direction for the current sampler. The sampler is guided along the annealed path towards the target. We defer the detailed experiment settings and more results to Appendix E.
Let M be the number
8

---

Under review as a conference paper at ICLR 2023
Specifically, we visualize the sample results three distributions with sand-to-sample charac- teristics such as multi-modality and periodicity, as shown in Figure [ It shows that samples from AFIS+MC method perfectly match all larget distributions. For quantitative compari- son, calculate the Maximum Mean Discrep- ancy between the pure HMC samples and all samplers' samples. The FSD-NS does not con- verge when training, so we omit the result of FSD-NS comparison. Since the task focuses on training implicit samplers, we do not com- pare other explicit samplers. Table [ summa- rizes the results of the MMD evaluation of all samplers. In all datasets, our AFIS consistently performs better than FIS. With additional MC correction steps, we always get lower MMD compared to the pure AFIS method.
t1 target
(b) t2 target
(c) 13 target
(d) t1 sample (e) 12 sample (f) t3 sample
Figure 5: Target and AFIS+MC samples.
Table 1: MMD (with rbf kernel) evaluation for synthetic targets. Additional 10 Langevin MC correction steps used in AFIS+MC sampler. The lower the metric, the better the sampler.
Target
double well
FIS(ours)
1.12e-241 076-3
3.51e-143.060 060-3
4.54g-244 10s-3
1.48o-2+1.110 4.6.e-22t225333- 3.6le-2:12.66..
5.186-242.68--3
AFlS(our)
7.07e-4=:1.720
137--241337--3
3.31c-3#1113ed
1.6e--2tt1.9e--
AFTS+MC(our)
2.45e-4.1 1.20e4
5.99
2.83731
2.29e-2.11.97--3
4.2 BAYESIAN REGRESSION
We also test our Implicit Sampler on Bayesian regression tasks as in Song al. (2017). HMC is good baseline for such tasks, as pointed out in Neklyudov et al.] (2020): Neklyudov & Welling (2022). The inference of the Bayesian logistic regression model aims to sample from the posterior distribution. We compare FIS (no anneal), AFIS, and AFIS+MC on Australian, German, and Heart datasets. evaluate samples' quality, we run HMC as a baseline to obtain approximated samples from target distributions and calculate Maximum Mean Discrepancy between samples from implicit samplers and HMC baseline. Table 2 shows the results of the Bayesian inference experiments. Other than FSD-NS, which always fails during training, samples. Moreover, annealed technique and MC correction steps further improve sample quality. Experimental details can be found in Appendix E.2
generators generate high-quality
Table 2: MMD (with rbf kernel) evaluation for posterior sampling. Additional 10 Langevin MC correction steps used in AFIS+MC sampler. The lower the metric, the better the sampler.
Posterior
Australian
German
Heart
FIS(ours)
7.99e-312.81e-
1.91e-4 6.48e-6
9.84e-5 1.00e-
AFIS(ours)
6.30e-3:2.50e.44
2.42e-614.02e--
3.66e-5 1.08e-5
AFIS+MC(ours)
1.06e-3=1.00e-4
2.466-6=3.97e-7
3.64e-5=1.077
CONCLUSION
We have presented novel approach for training an implicit sampler to sample from un-normalized density. Our approach minimizes the Fisher Divergence with the aid of an asynchronous score network. We show theoretically that our method can accurately minimize the Fisher Divergence for the implicit sampler, which is the first one as far as we know. Besides, our approach uses both the annealing technique and stochastic corrections for improved sampling performance. We also prove the faster mixing for MC correction. We test our approach on commonly used synthetic target generation and Bayesian regression benchmarks and observe ideal performance.

---

Under review as a conference paper at ICLR 2023
ETHICS STATEMENT
Our work proposes an approach to train an implicit sampler by minimizing Fisher Divergence be- tween sampler and target distribution. Since the research is fundamental methodology in machine learning, the negative consequences of the methodology seem not obvious.
REPRODUCIBILITY STATEMENT
We provide details of our approach and sampler in Appendix. provide complete proofs of theoretical results also in Appendix. We also propose the python code for implementation. We state that our research is reproducible.
REFERENCES
Christophe Andrieu, Nando De Freitas, Arnaud Doucet, and Michael Jordan. An introduction to meme for machine learning. Machine learning, 50(1):5-43, 2003.
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport monte carlo. In International Conference on Machine Learning, pp. 318-330, PMLR, 2021.
Fan Bao, Chongxuan Li, Kun Xu, Hang Su, Jun Zhu, and Bo Zhang. Bi-level score matching for learning energy-based latent variable models. Advances in Neural Information Processing Systems, 33:181 10-18122, 2020.
MS Bartlett. Approximate confidence intervals. Biometrika, 40(1/2):12-19, 1953.
Changyou Chen, David Carlson, Zhe Gan, Chunyuan Li, and Lawrence Carin. Bridging the gap be- tween stochastic gradient meme and stochastic optimization. In Artificial Intelligence and Statis- tics, pp. 1051-1060. PMLR, 2016.
Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet. Differentiable par- ticle filtering via entropy-regularized optimal transport. In International Conference on Machine Learning, pp. 2110-2111. PMLR, 2021.
Francesco D' Angelo and Vincent Fortuin. Annealed stein variational gradient descent. In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
Arnaud Doucet, Nando de Freitas, and Neil Gordon. An introduction to sequential monte carlo methods. In Sequential Monte Carlo methods practice, pp. 3-14. Springer, 2001.
Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Çelikyilmaz, and Lawrence Carin. Cyclical annealing schedule: A simple approach to mitigating kl vanishing. In NAACL-HLT 2019.
Charles Geyer and Elizabeth A Thompson. Annealing markov chain monte carlo with applications to ancestral inference. Journal of the American Statistical Association. (0(43)):909-99) 1995.
Jackson Gorham and Lester Mackey, Measuring sample quality with stein's method. Advances in
Neural Information Processing Systems, 28, 2015.
Jackson Gorham and Lester Mackey. Measuring sample quality with kernels. In International Conference on Machine Learning, pp. 1292-1301. PMLR, 2017.
Peter Green. Reversible jump markov chain monte carlo computation and bayesian model deter- mination. Biometrika, 82(4):711-732, 1995.
Leonard Gross. Logarithmic sobolev inequalities. American Journal of Mathematics, 97(4):1061- 1083, 1975.
WK Hastings. Monte carlo sampling methods using markov chains and their applications. Biometrika, 57(1)997-109 1970.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:68400666 68511, 2020.
10

---

Under review as a conference paper at ICLR 2023
Tianyang Hu, Zixiang Chen, Hanxi Sun, Jincheng Bai, Mao Ye, and Guang Cheng. Stein neural sampler. arXiv preprint arXiv:1810.03545, 2018.
Chin-Wet Wei Huang, Shawn Tan, Alexandre Lacoste, and Aaron C Courville. Improving explorabil- ity in variational inference with annealed variational objectives. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Infor- mation Processing Systems, volume 31. Curran Associates, Inc., 2018.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances in neural information processing systems, 34:21696-11707, 2021.
Durk P Kingma and Yann Cun. Regularized estimation of image statistics by score matching. Ad- vances in neural information processing systems, 23, 2010.
Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score-based diffusion odes by high order denoising score matching. In International Conference on Machine Learning, pp. 14460. PMLR, 2022.
Stephan Mandt, James McInerney, Farhan Abrol, Rajesh Ranganath, and David Blei. Variational tempering. In Arthur Gretton and Christian C. Robert (eds.), Proceedings of the 19th Interna- tional Conference on Artificial Intelligence and Statistics, volume 51 of Proceedings of Machine Learning Research, pp. 704-712, Cadiz, Spain, 09-11 May 2016. PMLR.
Enzo Marinari and Giorgio Parisi. Simulated tempering: new monte carlo scheme. EPL (Euro- physics Letters), 19(6),551, 1992.
James Martens, Ilya Sutskever, and Kevin Swersky. Estimating the hessian by back- paok-propgggting curvature. In Proceedings of the 29th International Coference on International Conference on Machine Learning, pp. 963-970, 2012.
Agdg Matthews, M. Arbel, D. J. Rezende, and A. Doucet. Continual repeated annealed flow trans- port monte carlo. 2022.
Chenlin Meng, Lantao Yu, Yang Song, Jiaming Song, and Stefano Ermon. Autoregressive score matching. Advances in Neural Information Processing Systems, 33:6673-666 6683, 2020.
Radford M Neal. Annealed importance sampling. Statistics and computing, 1 11(2): 139, 2001.
Radford M Neal. Meme using hamiltonian dynamics. In Handbook of Markov Chain Monte Carlo, pp. 139-188. Chapman and Hall/CRC, 2011.
Kirill Neklyudov and Max Welling, Orbital meme. In International Conference on Artificial Intel- ligence and Statistics, pp. 5790-5814. PMLR, 2022.
Kirill Neklyudov, Max Welling, Evgenii Egorov, and Dmitry Vetrov. Involutive meme: a unifying framework. In International Conference on Machine Learning, pp. 7273-7282. PMLR, 2020.
Peter Olsson. Two phase transitions in the fully frustrated xy model. Physical review letters, 75(14): 2758, 1995.
Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu. Efficient learn- ing of generative models via finite-difierence score matching. Advances in Neural Information Processing Systems, 33:19175-19188 2020.
Grigorios A Pavliotis. Stochastic processes and applications: diffusion processes, the Tokker-Plance and Langevin equations, volume 60. Springer, 2014.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Interna- fional conference machine learning, pp. 1530-1538. PMLR, 2015.

---

Under review as a conference paper at ICLR 2023
Gareth O Roberts and Jeftrey S Rosenthal. Optimal scaling of discrete approximations to langevin diffusions. Journal of the Royal Statistical Society: Series (Statistical Methodology), 60(1): 255-268, 1998.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International conference on machine learning, pp. 1218-1226. PMLR, 2015.
Ch Schüttte, Alexander Fischer, Wilhelm Huisinga, and Peter Deuflhard. direct approach con- formational dynamics based on hybrid monte carlo. Journal of Computational Physics, 151(I): 146-168, 1999.
Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-nice-mc: Adversarial training for memc. Advances in Neural Information Processing Systems, 30, 2017.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation, In Proceedings of the Thirty-Fiftt Conference on Uncertainty in Artificial Intelligence, UAI 2019, Tel Aviv, Israel, July 22-25, 2019, pp. 204, 2019.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1.Xi::2011.13455 2020.
Charles M Stein. Estimation of the mean of multivariate normal distribution. Statistics, pp. 1135-1151, 1981.
The annals of
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 1287-11302, 2021.
Peter JM Van Laarhoven and Emile Aarts. Simulated annealing. In Simulated annealing: Theory and applications, pp. 7-15. Springer, 1987.
Pascal Vincent. connection between score matching and denoising autoencoders. Neural compu- tation, 1661-17661 2011.
Li K Wenliang and Heishiro Kanagawa. Blindness of score-based methods to isolated components and mixing proportions. arXiv e-prints, pp. arXiv-2008, 2020.
Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the Bayes posterior in deep neural networks really? In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 10248-10259, PMLR, 13-18 Jul 2020.
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. Advances in Neural Informa- fion Processing Systems, 33:5933-5944, 2020.
Tatiana Xifara, Chris Sherlock, Samuel Livingstone, Simon Bymee and Mark Girolami. Langevin diffusions and the metropolis-adjuste langevin algorithm. Statistics & Probability Letters, 91: 14-19, 2014.
Takuya Yamano. Skewed jensen-fisher divergence and its bounds. Foundations, 1(2):256-264 2021.
12

---

Under review as a conference paper at ICLR 2023
A PROOF OF PROPOSITION
We provide the proof of Proposition here.
Proof.
With fixed p and known target ạ the optimal test function f has representation
arg min L(F)
f
Where functional [f) has integral representation
C(f) =Ezup (V loggll)) + (Vx)
= / (...
/ Ux, Vf)dz.
Here 1(x Vf) f (V= log (x)) f(z)) + ((x)(V,, ((:)) (x)((r) By Fuler-
Lagrange equation, the optimal function f satisfies
Ø1
d
01
0²
0l
0.
Of
dx
00
Ox
0P"
By caleulation, have
01
Of (x) = p(x)) log (()) - 2xp(z)f(x)
de DP (x) = vxp(x)
OF (x) = O.
So the optimal f* satisfies the Euler-Lagrange equation as
V logg(z) - 2Ap(z)f(x) - v.p(x) = 0.
Divide the both side with p(z) and note that Vap(z)/p(z) = v. log the equation turns to
1
Va logg(z) V= logr(z)) (*))) 2X
Next consider optimal s* The s* obtained by minimizing the Score Matching objective, which is equivalent to minimizing the Fisher divergence between and induced family, thus the optimal (z) V. p(r). Substitute V, log p(x) with s* into f formula, have
1 V log g(x) - s-*(())) (z)].
('(x)
2A
B PROOF PROPOSITION 2
In this section, prove that the S2D loss and Fisher Divergence shares exactly the parameter gradient.
Proof. Let Pe denote sampler's distribution. so denote the true but unknown sampler's score func- tion. g denotes the known un-normalized target. For rest of the proof, the notion ||>|]] represents the L norm of a vector in Dx dimensional Euclidean space x E Dx Recall that the Fisher Divergence is defined
LFD(O) Ex V lxg so(x)
13

---

Under review as a conference paper at ICLR 2023
Thus the sampler parameter gradient of Fisher Divergence writes
0
/ IIV logg(a) so (00 (x())) 00 E v (lr) so(x)]
a
00
/ r logg(())) s0(x)11 a
0
po(x)der + po po()))
logg(z) - s0(x)) (2)
0g
00
= pa V (x) so(x)]1 log Bp, ((0 V. log (2;))
00 50(x)
00
= (1) + (2).
The first term can be estimated with
V log ss(x) 0
(1)
00 po (z)
sg V logg(z) - - po(
00
0 Ep sg
V -()
00
Here the operator sg denotes stop gradient operator with respect to parameter 8. sg(fo] stop the
parameter dependence of for function meaning that can only evaluate fo(x) point-wise but can not obtain the gradient of fz)).. Here we stop the gradient of function || log (()) sg (x) we use another score network &$ approximate so point-wise, regardless of the 0 parameter dependence. Next we consider the second term. The second term turns to
(2) = Ep. (x) - v log ( so(z)
00
(xp(((s(r)) v logg(z))) T
V pa)
00
/ (o())(80(x) v. logg
log po do
00 dr
Opo(x) dr
po(())(0(x)) Va log (x)))
00 po(x) Or
2 / (so(x)) V logg(x))²
① log po 0 log pe (x)
09 On po(z) dx 2 (o(2))(0(())- - Ve logg(x))²
dr
06
(3) + (4).
Looking at (3), we have
(3) 2 (30(x) V oog)
30 Fx po (x de
2
sg (so(x) V. Vzlog((z))
dr (o))
00
θ
/
2 00
pa(x + ev)dæ,
sg (so(x) - V Volog((t))
de
2 0 0
/ po(x + ev)dæ
00 0e
2 0 0
00 de
0.
Above equality holds because of f "olo + ev)dz holds for all c, 9, we view e as a shift strength parameter, the above equality recovers the first order Bartlett identity (Bartlett 1953).
14

---

Under review as a conference paper at ICLR 2023
Next turns term (4). Note that
(4) 2 / po(x)(so(x) V. log (
3 log po (3)
alog po (1:)
Ōr
00
2 / po(x)) (so(x)- - v log (2)) )
d log pa (x) dx
i log (x
00
2 / (80((0) Vx (z) 2 log po (x)
Opo
Ox
00
0
/ sg
log pa (x: Oz
2
(so(x) v log (x))(x)))
Po(x)
00
0
T 2 log po (1)) Oz
2
00 8g
(80(x) V (x))
Po(x))
0
2 log Pa (2:)
2 00
(m)
E
sg
(so(x) V logg(x)))
Or
0
a log (o(
2
E
Po (1))
(so(x) v log (z)))
sg
00
dr
0
T
}
2 00
E
sg
(V log ((x) - so s0()))
sg
sp((r)
Combining all above, we caleulate the parameter derivative as
V log (x)) ==
08 pu
(1)
(3) (4)
0
V x ((z) - s2(x))2 + - 2
T
E,,
po
sg
(sa(x) - V. log ((x)) sg so(x)}
08
00
0
E po
sg
V log (x) 1|3
sg
||80(1))1
00
Thus the equivalent loss function
Ls2p (0) = Ep,
sg | log g(x)
sg [
Share the same parameter gradients as the Fisher divergence which is intractable. Since we only need the x gradient of sampler score function se (because the stop gradient operator), (0(x) through another score network sz) with samples consistently obtained from sampler. With above objective function, could minimize the Fisher divergence between po and
estimate
15

---

Under review as a conference paper at ICLR 2023
C INTRODUCTION TO METROPOLIS- AND HAMILTONIAN CORRECTION
Assume the target distribution is p(x), the MH MCMC requires proposal distribution p(ā(z)) to propose candidate samples 3: ~ (()(ar) The Markov chain then accept the candidate sample with probability min p(ā)g(a #)
conditions, the chain will eventually reach p(æ) as stationary distribution. The proposal distribution can be symmetric non-symmetric. Conditional gaussian g[ā(2z) = (2; o²) is a usual choice. Proposals based on score function g(0[x;) NNx + , log)(z)) o is also popular (Xifara et al. 2014). If one consider an auxiliary state space of (x, v) and execute the proposal in such space, the MC schedule is called Hamiltonian Monte Carlo. The Hamiltonian Monte Carlo execute Monte Carlo dynamic in auxiliary space. With current sample X The HMC sample momentum vector from an auxiliary distribution V(t) exp(-o" M u/2) The joint sample (x(9, v(5) updated by running Hamiltonian Dynamics in joint space via
(( æ) 13. Under
dX. dt OV
OH dv. OH
0x
Here Hxx, v) - logp(()) + 10 10 M- is the Hamiltonian of such mechanical system. HMC has many advantage that it mixes well for high-dimensional targets, and travels in joints space thus not easy to be trapped in local minima. Leap frog integrator is usually practical choice for numerical updates (Neal 201T). To make Markov Chain detail balanced, additional Metropolis correction also needed for Hamiltonian proposal. In short words, HMC iteratively accepts new position and
H(8,0) where (ø, ü) LeapFrog(a; v) as approximated momentum pair (z, 0) with rate min 1, ((a,)) Hamiltonian proposal.
D PROOF OF THEOREM
We give the proof of Theorem here. To begin with, we give a lemma to bound KL divergence with Fisher divergence shown in Yamano (202l)
For fixed there exists a dimension-free positive constant such that for every distri- bution p which both integral and log-integral with respect to q, and p has same support as q, we have
Lemma 2.
DKL S DF(p.9)
proof of lemma. For every order smooth function f, assume both | and // are integrable with respect to g, the log-Sobolev's inequality (Gross) 1975) shows that there exist dimension-free positive constant c, such that
L I | log (ll(())ddd sc / [ +1 log II
Here | = (zd,, Replace = v /pq.. q. we have
LHS = / (1//)) lo//////////a//-= E. log(p/)) (xi(p.9).
So we have
Vp vp
44
vp
2 2
P VVP
9
P Vq Va q
=
2/
P
logp
Vlogg
Thus the first term in RHS
/ [VF1IP(()) =c
/ / v
p 21(q(x)) = /
V Vlogp- V 2
q
= Ep V Voop
V 2 =
16

---

Under review as a conference paper at ICLR 2023
Note that |1 1 l(p/9)q fp 1. We combine both sides to conclude
2 DkL(P.1) - 2 DD +
So we have
DKL(D, S Dr(p. 4),
where e be another positive constant.
The above lemma shows that KL divergence is upper bounded with Fisher divergence, which we are using train the sampler. With above lemma, we can calculate mixing time for Langevin correction in proof below
Proof. Assume target satisfies
V logg(x))
lim ||=|| 1
A Algg(())) +00.
then their exits a constant À > 0, such that Poincare inequality holds for each f E C=C(((RR) (7)
with f Theorem 4.3 in Pavliotis (2014)
Let po denotes the ASS distribution, which is trained to be bounded with Dp(po, 9) I 8. By lemma, the KL between initial distribution po and target is bounded by Fisher divergence with dimension- free constant
DKL. 9) S DD 2) S S +o0
With Poincare's inequality holds, the KL along Langevin diffusion dX, V llog((X.)) dW, decays exponentially fast as in Theorem 4.6 in Pavliotis (0014)
(KL q) -
I (xx) C Dr (p(,,))
2
I exp( exp(
Thus if we want (KI(Pe, q) to be controlled under tolerance we only need diffused time / to satisfies
2x log( log( 2
1
Cllo
log( 6
2x
where we place C' = log( { to be another constant. The diffusion time must be positive, thus we take
1 + log( 8 9)}
T max 2A
and finish the proof.
E EXPERIMENTAL DETAILS AND MORE RESULTS
E. SYNTHETIC TARGET
For toy 2-dimensional data experiments, use 3-layer MLP neural network with 200 hidden units in each layer as the sampler. The activation of the sampler is chosen as LeakyReLU non-linen non-lincarity with 0.2 coefficient. The score network is 3-layer MLP with 200 hidden units in each layer. The activation of the score network is GELU non-lincarity.
When reporting the numbers in Tab T compute MMD metrics based total of 2000 samples. We run 20 independent experiments for each target and algorithm to calculate the mean and standard deviation.
Figure 6 visualizes the model capabilities of FIS, AFIS, and AFIS+MC samplers for matching three 2-dimensional target energy functions.
17

---

Under review as a conference paper at ICLR 2023
O
(a) Ta-Teteet
(b) tI-FIS
(C) II-FIS
(d) 1IAFFISMC
(e) 12- (2-Target
(0) 12-FIS
(g) 12-AFIS
(h) M2-AFIS+MC
(i) 13 13-target
(i) 13-FIS
(k) 13-AFIS
() 1-AFFI+MC
Figure 6: Comparison between samples generated by FIS, AFIS and AFIS+MC on three 2D energy functions.
E.2 BAYESIAN REGRESSION
For nigh-dimensional data experiments, we also use 3-layer MLP neural networks as the sampler and score network, respectively. The activation of the sampler is chosen LeakyReLU non-linearity with a 0.2 coefficient. The activation of the score network is GELU non-linearity. For Australian and Heart distributions, use 400 hidden units in each layer and 600 hidden units for German distribution.
When reporting the numbers in Tab we compute the MMD metric based on a total of 2000 samples. We run 20 independent experiments for each target and algorithm to calculate the mean and standard deviation. For the basic settings of Bayesian Regression problems, readers could refer to Song et al. (2017) for more details.
18

---

Under review as a conference paper at ICLR 2023
F
FULL AFIS ALGORITHM
This section gives the full Annealed Fisher Implicit Sampler training algorithm.
Algorithm 2: Annealed Fisher Implicit Sampler training algorithm
Input: un-normalized target log g(2), annealed schedule {{ prior distribution log Aprior (x) latent distribution pz (2), implicit sampler Gg, score network Ss, mini-batch size B, max iteration M.
Randomly initialize (8("), (.)).
for k in 1:K do
# anneal the larget
set log (x) = X log g(x) (1 Xx) log yprior (I)
for in 1:M do
# update score network parameter
Get mini-batch from sampler 2; = Goco ( I,., B. Calculate score matching objective
(o
Σ
[ss(()))1 so(x)))
Minimize (Ø) to get g(t+1). # update sampler parameter
Get mini-batch latent code i= -,...,, B. Use re-parametrization trick to calculate S2D loss for sampler
LS2 = B
Σ
1 log gu (Go - Ilsaar 1 (Go
i=1
Minimize (0) get ((++1)
end end
return (0, o).
19


Reviews:
[{"review_id": "wjAAy7kq87", "reviewer": "Reviewer_mcoE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a method to sample from a distribution with a known un-normalized posterior, named Annealed Fisher Implicit Sampler.\nMore specifically, they introduce the loss with the same gradient as the Fisher Divergence, provided a perfect estimate of the sampler's score function.\nThey train the sampler via alternatively estimating the sampler's score function and updating sampler parameters via the loss.\nThey employ annealing technique to alleviate sampler's training, with the rationale that the sampler is trained easily on easier target distributions.\nThey show the annealed sampler performance on a variety of benchmarks. They provide proofs of the propositions and a theorem, as well as experiment details in the appendix.", "summary_of_the_review": "To sum up, the paper provides a novel method to sample from a distribution with a known unnormalized posterior. The contribution has potential impact, as sampling from such distribution has wide applications. The paper could use some improvement clarity-wise. It would be benefitial to provide comparison against other methods. Initial score: borderline.\n", "strength_and_weaknesses": "*Strengths*:\n- The proposed method is theoretically well-backed. The loss enjoys the same gradient, as (intractable) Fisher divergence, provided a perfect estimate of the sampler's score function.\n- The empirical study of the method's performance show that it indeed can capture target densities.\n\n*Weaknesses*:\n- The experimental results do not compare to any other methods.\n- The clarity of the paper could use some improvements (concrete examples are given in the Clarity section).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JWxIC8wPNG", "reviewer": "Reviewer_enF9", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposed a method to design a neural sampler that minimizes the fisher score discrepancy.", "summary_of_the_review": "There are several aspects which are unclear to me, and hence I can not recommend acceptance. Further, I am worried about the relevance to the iclr community. I will reserve my judgement till rebuttals.", "strength_and_weaknesses": "\nWeaknesses: \n-- Some aspects of the paper are not clear to me. It seems these things are accepted in the community though. For example, why the need for a separate scoring network and the score estimation step in general ? If I am modeling p_\\theta anyways, why can I not just use it directly? The authors jump between using s() and log p as the score function, which makes sense but the two notations and statements like \"if we asynchronously estimate the sampler’s score function perfectly.\" tells me there is more to this that is not clear from the way it is written.\n\n-- \n\n-- Please explain \"asynchronous\" estimation. When I hear \"asynchronous\", I tend to think of it in context of parallel or distributed optimization. \n\n-- The derivative using s_d (x) below Figure 2 is not clear to me. What is s_d (x) ? Did I miss where it is defined ? \n\n-- Minor: There are some typos and clarifications required, please proof-read .e.g. \n\"seems work\" --> \"seems to work\". \n\"More specially\"\nWhat is \\lambda in Prop 1 ?\nPlease mention what FSD stands for /before/ using it.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "HDF3m703qzp", "reviewer": "Reviewer_gNTU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents the annealed fisher implicit sampler which trains an implicit sampler (i.e GAN style model) to sample from an unnormalized probability distribution. The model is trained to minimize the Fisher divergence (i.e the score matching objective) between the model samples and the unnormalized distribution. This is an appealing objective for unnormalized distributions as the normalizing constant of the target is not necessary. However, we must be able to compute the the score function (i.e log-probability gradient) of the implicit model which is not tractable. The authors propose to use score matching techniques to learn a score network to approximate this quantity. They then propose a surrogate objective involving the target distribution, the score network, and the implicit sampler to minimize the fisher divergence.\n\nThe authors advocate for an annealed training objective which anneals the target from a simple distribution to the desired distribution which improves performance. Also the authors advocate to use a few steps of Langevin MCMC on top of the implicit sampler to correct for errors which they claim improves performance.", "summary_of_the_review": "This paper proposes the Annealed Fisher Implicit Sampler which is a neural implicit sampling method training to minimize a novel objective  which approximates the fisher divergence between the target distribution and an implicit neural sampler. The method is interesting but the empirical evaluation is severely lacking. No baselines are presented besides methods proposed in this paper. As well, the method is not sufficiently explored and its problems (such as bias from score approximation) are not considered. \n\nThis method is interesting, but I believe a more thorough evaluation is necessary to understand how this approach will scale and how it compares to traditional sampling techniques and competing neural approaches. ", "strength_and_weaknesses": "Strengths:\n\nThe method presented in the paper in interesting and the provided theory is detailed and appears correct. I was intrigued but the surrogate objective and its equivalence to the original goal of score matching. The theory was pretty easy to follow as well. \n\n\nWeaknesses:\n\nOverall the paper has 2 main weaknesses. First, I do not think the potential issues of the approach are explored in enough detail. Second the empirical evaluation is insufficient and lacking baselines.\n\nPotential Issues:\n\nI am concerned by utilizing an approximate to the score function of the model. This incorporates bias into the training objective. The authors show that when the score network is perfect then we are minimizing FD. How much error (in terms of FD maybe) is tolerable? Is it possible to bound the bias using the FD between the score network and implicit model? I am worried as well that this optimization may run into saddle point issues like with GANs or the neural stein sampler. It might be interesting to see how this approach works in a setting where we can get the exact score function (such as with a normalizing flow model). We could then compare how the normalizing flow performance works when trained with reverse-KL, true Fisher Divergence, and approximated Fisher Divergence (using the score network). Right now its not so clear how these things compare. \n\nEvaluation: \n\nOnly two tasks are used for evaluation and these are relatively small scale in nature. Existing sampling approaches often present results on distributions with thousands of dimensions. I would be more convinced if the authors included at least one such distribution such as a pre-trained normalizing flow or a log gaussian cox process (two benchmarks which are becoming more standard in the field). Next the, paper is severly lacking baselines. In fact, there are no baselines beyond methods presented in this paper. The authors compare their method to the neural stein sampler. Why did they not compare with it? As well, the direct competitor to neural samplers are interactive methods such as MCMC and SVGD. We should also compare with these kinds of approaches. Sequential Monte-Carlo is currently the gold standard in the field and some such method should be used to give an indication of how well this approach performs to the best out there. Besides this there are a number of alternative neural sampling approaches such as the Path Integral Sampler which are beginning to gain traction. It would also be useful to provide one such comparison. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "0CavKwf64Y", "reviewer": "Reviewer_G8ZU", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed an implicit sampler by minimizing the Fisher divergence and then propose to improve it by using the annealing technique. Specifically, the author proposed an alternative objective such that minimizing it is equivalent to minimizing the Fisher divergence. Later, the author tried to solve the distant mode problem by creating a sequence of annealed distributions from simple Gaussian to target distribution. Theoretically, the author showed the equivalence of minimizing S2D (the proposed one) to Fisher divergence and also shows that the proposed implicit sampler can be used as an initialization for MCMC chains. \n\nEmpirically, the author conduct 1 synthetic experiment and 1 Bayesian regression to demonstrate the effectiveness of the proposed approach.", "summary_of_the_review": "This paper provides an implicit sampler by minimizing Fisher divergence. My major concern is its usefulness due to the property of Fisher divergence and clarity of Proposition 2. Empirically, I think more large experiments should be conducted to see its effectiveness to support the claim made in the introduction. ", "strength_and_weaknesses": "## Strength\nThe paper is clearly written apart from the proof of Proposition 2. So in general, the paper is easy to read and understand. Although it is well-known that the Langevin sampler can be controlled by KL divergence and Fisher divergence is stronger (i.e. upper bounding) than other discrepancies, it is still good the formalize it in theorem 1. The proof strategy of theorem 1 is relatively easy to understand. \n\n## Weakness\nAfter reading the paper, I have several concerns: (1) the usefulness of the proposed approach; (2) the clarity of proposition 2, and (3) the empirical evaluation. \n\nFirst, theoretically speaking, Fisher divergence is stronger than many of the well-known divergences and should provide better control over the convergence. However, in practice, it has been found that it is not the case. There have been many works investigating the theoretical properties of using Fisher divergence as the estimator (see [1], [2]). Especially in [2], it has been found that score matching may fail even under the simple student-t distribution. So, I wonder is your proposed approach works under the following unimodal setting: \n1. Implicit sampler is a student t distribution with mean $\\theta$\n2. target distribution is another student t distribution with mean $\\theta_0$, and $\\theta$ and $\\theta_0$ are far away from each other. \nBy minimizing the Fisher divergence, $\\theta$ should approach $\\theta_0$, is this true in practice (without using the MCMC step afterward)?\n\n\nFor proposition 2, I am not sure I understand it correctly. In the end of the proof, the equivalence is established with the $sg(\\cdot)$ operator but in algorithm 1, the $\\theta$ is back-propagated through the score estimator and target score. So why there isn't a $sg(\\cdot)$ in this case?\n\nFor the empirical evaluation, the tasks are too simple to demonstrate the usefulness of the proposed approach. In the introduction, you mentioned the importance of the neural sampler, but I fail to see this based on the current empirical evaluation. For example, sampling from EBM is an interesting experiment since it is well-known MCMC sampling is too slow to train EBM. Another possible experiment is to train VAE where the encoder is the neural sampler. More baseline can also be considered, at least with the neural Stein sampler. \n\n\n\n\n[1] Barp, A., Briol, F. X., Duncan, A., Girolami, M., & Mackey, L. (2019). Minimum stein discrepancy estimators. Advances in Neural Information Processing Systems, 32.\n\n[2] Gong, W., & Li, Y. (2021). Interpreting diffusion score matching using normalizing flow. arXiv preprint arXiv:2107.10072.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject