SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We here provide various additional results for the SBI benchmark. First, we compare the performance of FMPE and NPE when using the Maximum Mean Discrepancy metric (MMD). The results can be found in Fig. 6. FMPE shows superior performance to NPE for most tasks and simulation budgets. Compared to the C2ST scores in Fig. 4 the improvement shown by FMPE in MMD is more substantial.
Fig. 7 compares the FMPE results with the optimal transport path from the main text with a comparable score matching model using the Variance Preserving diffusion path [15]. The score matching results were obtained using the same batch size, network size and learning rate as the FMPE network, while optimizing for  $\beta_{\mathrm{min}} \in \{0.1, 1, 4\}$  and  $\beta_{\mathrm{max}} \in \{4, 7, 10\}$ . FMPE with the optimal transport path clearly outperforms the score-based model on almost all configurations.
In Fig. 8 we compare FMPE using the architecture proposed in Section 3.2 with  $(t,\theta)$ -conditioning via gated linear units to FMPE with a naive architecture operating directly on the concatenated  $(t,\theta,x)$  vector. For the two displayed tasks the context dimension  $\dim(x) = 100$  is much larger than the parameter dimension  $\dim(\theta) \in \{5,10\}$ , and there is a clear performance gain in using the GLU conditioning. Our interpretation is that the low dimensionality of  $(t,\theta)$  means that it is not well-learned by the network when simply concatenated with  $x$ .
hyperparameter	sweep values
hidden dimensions	2nfor n ∈ {4, . . ., 10}
number of blocks	10, . . ., 18
batch size	2nfor n ∈ {2, . . ., 9}
learning rate	1.e-3, 5.e-4, 2.e-4, 1.e-4
α (for time prior)	-0.25, -0.5, 0, 1, 4
Table 2: Sweep values for the hyperparameters for the SBI benchmark. We split the configurations according to simulation budgets, e.g. for 1000 simulations, we only swept over smaller values for network size and batch size. The network architecture has a diamond shape, with increasing layer width from smallest to largest and then decreasing to the output dimension. Each block consists of two fully-connected residual layers.
Figure 6: Comparison of FMPE and NPE performance across 10 SBI benchmarking tasks [46]. We here quantify the deviation in terms of the Maximum Mean Discrepancy (MMD) as an alternative metric to the C2ST score used in Fig. 4. MMD can be sensitive to its hyperparameters [46], so we use the C2ST score as a primary performance metric.
Fig. 9 displays the densities of the reference samples under the FMPE model as a histogram for all tasks (extended version of Fig. 3). The support of the learned model  $q(\theta | x)$  covers the reference samples  $\theta \sim p(\theta | x)$ , providing additional empirical evidence for the mass-covering behavior theoretically explored in Thm. 1. However, samples from the true posterior distribution may have a small density under the learned model, especially if the deviation between model and reference is high; see Lotka-Volterra (bottom right panel). Fig. 10 displays P-P plots for two selected tasks.
Finally, we study the impact of our time prior re-weighting for one example task in Fig. 11. We clearly see that our proposed re-weighting leads to increased performance by up-weighting samples for  $t$  closer to 1 during training.
Figure 7: Comparison of FMPE with the optimal transport path (as used throughout the main paper) with comparable models trained with a Variance Preserving diffusion path [15] by regressing on the score (SMPE). Note that the SMPE baseline shown here is not directly comparable to NPSE [8, 9], as this method uses Langevin steps, which reduces the dependence of the results on the vector field for small  $t$  (at the cost of a tractable density).
Number of Simulations
Figure 8: Comparison of the architecture proposed in Section 3.2 with gated linear units for the  $(t,\theta)$ -conditioning (red) and a naive architecture based on a simple concatenation of  $(t,\theta,x)$  (black). FMPE with the proposed architecture performs substantially better.
Figure 9: Histogram of FMPE densities  $\log q(\theta |x)$  for samples  $\theta \sim q(\theta |x)$  and reference samples  $\theta \sim p(\theta |x)$  for simulation budgets  $N = 10^{3}$  (left),  $N = 10^{4}$  (center) and  $N = 10^{5}$  (right). The reference samples  $\theta \sim p(\theta |x)$  are all within the support of the learned model  $q(\theta |x)$ , indicating mass covering FMPE results. Nonetheless, reference samples may have a small density under  $q(\theta |x)$ , if the validation loss is high, see Lotka-Volterra (LV).
$\log q(\theta |x)$
p
Figure 10: P-P plot for the marginals of the FMPE-posterior for the Two Moons (upper) and SLCP (lower) tasks for training budgets of  $10^{3}$  (left),  $10^{4}$  (center), and  $10^{5}$  (right) samples.
Number of Simulations
Figure 11: Comparison of the time prior re-weighting proposed in Section 3.3 with a uniform prior over  $t$  on the Two Moons task (Section 4). The network trained with the re-weighted prior clearly outperforms the reference on all simulation budgets.
D Gravitational-wave inference
We here provide the missing details and additional results for the gravitational wave inference problem analyzed in Section 5.
D.1 Network architecture and hyperparameters
Compared to NPE with normalizing flows, FMPE allows for generally simpler architectures, since the output of the network is simply a vector field. This also holds for NPSE (model also defined by a vector) and NRE (defined by a scalar). Our FMPE architecture builds on the embedding network developed in [7], however we extend the network capacity by adding more residual blocks (Tab. 3, top panel). For the  $(t,\theta)$ -conditioning we use gated linear units applied to each residual block, as described in Section 3.2. We also use a small residual network to embed  $(t,\theta)$  before applying the gated linear units.
In this Appendix we also perform an ablation study, using the same embedding network as the NPE network (Tab. 3, bottom panel). For this configuration, we additionally study the effect of conditioning on  $(t,\theta)$  starting from different layers of the main residual network.
D.2 Data settings
We use the data settings described in [7], with a few minor modifications. In particular, we use the waveform model IMRPhenomPv2 [76-78] and the prior displayed in Tab. 4. Generation of the training dataset with 5,000,000 samples takes around 1 hour on 64 CPUs. Compared to [7], we reduce the frequency range from [20, 1024] Hz to [20, 512] Hz to reduce the computational load for data preprocessing. We also omit the conditioning on the detector noise power spectral density (PSD) introduced in [7] as we evaluate on a single GW event. Preliminary tests show that the performance with PSD conditioning is similar to the results reported in this paper. All changes to the data settings have been applied to FMPE and the NPE baselines alike to enable a fair comparison.
D.3 Additional results
Tab. 5 displays the inference times for FMPE and NPE. NPE requires only a single network pass to produce samples and (log-)probabilities, whereas many forwards passes are needed for FMPE to solve the ODE with a specific level of accuracy. A significant portion of the additional time required for calculating (log-)probabilities in conjunction with the samples is spent on computing the divergence of the vector field, see Eq. (3).
hyperparameter	values
residual blocks	2048, 4096 × 3, 2048 × 3, 1024 × 6, 512 × 8, 256 × 10, 128 × 5, 64 × 3, 32 × 3, 16 × 3
residual blocks (t,θ) embedding	16, 32, 64, 128, 256
batch size	4096
learning rate	5.e-4
α (for time prior)	1
residual blocks	2048 × 2, 1024 × 4, 512 × 4, 256 × 4, 128 × 4, 64 × 3, 32 × 3, 16 × 3
residual blocks (t,θ) embedding	16, 32, 64, 128, 256
batch size	4096
learning rate	5.e-4
α (for time prior)	1
Table 3: Hyperparameters for the FMPE models used in the main text (top) and in the ablation study (bottom, see Fig. 12). The network is composed of a sequence of residual blocks, each consisting of two fully-connected hidden layers, with a linear layer between each pair of blocks. The ablation network is the same as the embedding network that feeds into the NPE normalizing flow.
Description	Parameter	Prior
component masses	m1, m2	[10, 120] M⊙, m1 ≥ m2
chirp mass	Mc = (m1m2)3/5 / (m1 + m2)1/5	[20, 120] M⊙ (constraint)
mass ratio	q = m2/m1	[0.125, 1.0] (constraint)
spin magnitudes	a1, a2	[0, 0.99]
spin angles	θ1, θ2, φ12, φJL	standard as in [79]
time of coalescence	tc	[-0.03, 0.03] s
luminosity distance	dL	[100, 1000] Mpc
reference phase	φc	[0, 2π]
inclination	θJN	[0, π] uniform in sine
polarization	ψ	[0, π]
sky position	α,β	uniform over sky
Fig. 12 presents a comparison of the FMPE performance using networks of the same hidden dimensions as the NPE embedding network (Tab. 3 bottom panel). This comparison includes an ablation study on the timing of the  $(t,\theta)$  GLU-conditioning. In the top-row network, the  $(t,\theta)$  conditioning is applied only after the 256-dimensional blocks. In contrast, the middle-row network receives  $(t,\theta)$  immediately after the initial residual block. With FMPE we can achieve performance comparable to NPE, while having only  $\approx 1/3$  of the network size (most of the NPE network parameters are in the flow). This suggests that parameterizing the target distribution in terms of a vector field requires less learning capacity, compared to directly learning its density. Delaying the  $(t,\theta)$  conditioning until the final layers impairs performance. However, the number of FLOPs at inference is considerably reduced, as the context embedding can be cached and a network pass only involves the few layers with the  $(t,\theta)$  conditioning. Consequently, there's a trade-off between accuracy and inference speed, which we will explore in a greater scope in future work.
Table 4: Priors for the astrophysical binary black hole parameters. Priors are uniform over the specified range unless indicated otherwise. Our models infer the mass parameters in the basis  $(M_c, q)$  and marginalize over the phase parameter  $\phi_c$ .
Network Passes	Inference Time (per batch)
FMPE (sample only)	248	26s
FMPE (sample and log prob)	350	352s
NPE (sample and log prob)	1	1.5s
Table 5: Inference times per batch for FMPE and NPE on a single Nvidia A100 GPU, using the training batch size of 4096. We solve the ODE for FMPE using the dopri5 discretization [80] with absolute and relative tolerances of 1e-7. For FMPE, generation of the (log-)probabilities additionally requires the computation of the divergence, see equation (3). This needs additional memory and therefore limits the maximum batch size that can be used at inference.
Figure 12: Jensen-Shannon divergence between inferred posteriors and the reference posteriors for GW150914 [73]. We compare two FMPE models with the same architecture as the NPE embedding network, see Tab. 3 bottom panel. For the model in the first row, the GLU conditioning of  $(\theta ,t)$  is only applied before the final 128-dim blocks. The model in the middle row is given the context after the very first 2048 block.