SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proposition 3. Let  $\mathbf{x} = (\mathbf{x}_1, \mathbf{x}_2)$  be jointly  $SNEF_{\mathbb{X}_1 \times \mathbb{X}_2, t, \exp(\cdot/2), \mu}$  with parameters  $\mathbf{V}$  and  $\Theta = ([W_1, W_2], b)$ . Assume that  $\mu(dx) = \mu_1(dx_1)\mu_2(dx_2)$  and  $\pmb{t}(\pmb{x}) = (t_1(\pmb{x}_1), t_2(\pmb{x}_2))$ . Then the marginal distribution of  $\mathbf{x}_1$  is  $SNEF_{\mathbb{X}_1, t_1, \exp(\cdot/2), \mu_1}$  with parameters  $\tilde{\mathbf{V}}$  and  $\Theta = (W_1, b)$ , for some matrix  $\tilde{\mathbf{V}} \in \mathbb{R}^{m \times n}$ .
Proof. Proposition 1 gives us the normalising constant for the parametrisation where biases are absorbed into  $\mathbf{V}$ . If we explicitly keep the biases in the parametrisation, we have
$$
k _ {\exp (\cdot / 2), \boldsymbol {t}, \mu} \left(\boldsymbol {\theta} _ {i}, \boldsymbol {\theta} _ {j}\right) = \exp \left(\frac {1}{2} \left(b _ {i} + b _ {j}\right)\right) z _ {e} \left(\frac {1}{2} \left(\boldsymbol {w} _ {i} + \boldsymbol {w} _ {j}\right)\right), \tag {17}
$$
where  $z_{e}$  is the normalizing constant of the exponential family with the sufficient statistic  $t$  and base measure  $\mu$ . By Theorem 2, we have that
$$
P _ {1} (d \boldsymbol {x} _ {1}; \boldsymbol {M}, \boldsymbol {\Theta}) = \frac {\operatorname {T r} (\boldsymbol {M C} _ {\boldsymbol {\Theta}} (\boldsymbol {x} _ {1}))}{\operatorname {T r} (\boldsymbol {M K} _ {\boldsymbol {\Theta}})} \mu_ {1} (d \boldsymbol {x} _ {1}),
$$
where  $[C_{\Theta}(\pmb{x}_1)]_{ij} = k_{\sigma, \pmb{t}_2, \mu_2}\Bigl(\bigl(\pmb{w}_{2i}, \pmb{w}_{1i}^\top \pmb{t}_1(\pmb{x}_1) + b_i\bigr), \bigl(\pmb{w}_{2j}, \pmb{w}_{1j}^\top \pmb{t}_1(\pmb{x}_1) + b_j\bigr)\Bigr)$ . But now since  $\sigma = \exp (\cdot / 2)$ , applying (17) gives
$$
[ \boldsymbol {C} _ {\Theta} (\boldsymbol {x} _ {1}) ] _ {i j} = \mathrm {z} _ {e, 2} \left(\frac {1}{2} \left(\boldsymbol {w} _ {2 i} + \boldsymbol {w} _ {2 j}\right)\right) \exp \left(\frac {1}{2} \left(\boldsymbol {w} _ {1 i} ^ {\top} \boldsymbol {t} _ {1} (\boldsymbol {x} _ {1}) + b _ {i}\right)\right) \exp \left(\frac {1}{2} \left(\boldsymbol {w} _ {1 j} ^ {\top} \boldsymbol {t} _ {1} (\boldsymbol {x} _ {1}) + b _ {j}\right)\right),
$$
where  $\mathbf{z}_{e,2}$  is the normalizing constant of the exponential family with the sufficient statistic  $t_2$  and base measure  $\mu_2$ . Thus, we can write
$$
\begin{array}{l} \operatorname {T r} \left(M C _ {\Theta} \left(\boldsymbol {x} _ {1}\right)\right) = \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} \left\{m _ {i j} z _ {e, 2} \left(\frac {1}{2} \left(\boldsymbol {w} _ {2 i} + \boldsymbol {w} _ {2 j}\right)\right) \right. \\ \left. \cdot \exp \left(\frac {1}{2} \left(\boldsymbol {w} _ {1 i} ^ {\top} \boldsymbol {t} _ {1} (\boldsymbol {x} _ {1}) + b _ {i}\right)\right) \exp \left(\frac {1}{2} \left(\boldsymbol {w} _ {1 j} ^ {\top} \boldsymbol {t} _ {1} (\boldsymbol {x} _ {1}) + b _ {j}\right)\right) \right\} \\ = \exp \left(\frac {1}{2} \left(\boldsymbol {W} _ {1} ^ {\top} \boldsymbol {t} _ {1} (\boldsymbol {x} _ {1}) + \boldsymbol {b}\right)\right) ^ {\top} \tilde {\boldsymbol {M}} \exp \left(\frac {1}{2} \left(\boldsymbol {W} _ {1} ^ {\top} \boldsymbol {t} _ {1} (\boldsymbol {x} _ {1}) + \boldsymbol {b}\right)\right) \\ \end{array}
$$
and we conclude that the marginal is in the same family with  $\tilde{M} = M\circ Z_{e,2}$ , where
$$
[ \boldsymbol {Z} _ {e, 2} ] _ {i, j} = \mathrm {z} _ {e, 2} \left(\frac {1}{2} (\boldsymbol {w} _ {2 i} + \boldsymbol {w} _ {2 j})\right).
$$
Note that  $\tilde{M}$  is PSD as an Hadamard product of two PSD matrices. Thus, we can find  $\tilde{\boldsymbol{V}}$  such that  $\tilde{M} = \tilde{\boldsymbol{V}}^{\top}\tilde{\boldsymbol{V}}$ .
SNEFY 1	Resampled 1	Gauss 1	GMM 1
Moons	-1.59 ± 0.03	-1.76 ± 0.02	-3.29 ± 0.01	-1.57 ± 0.03
1431	68007	1194	1240
1352.19 ± 134.30	615.57 ± 24.35	730.21 ± 23.60	753.16 ± 31.41
Circles	-1.92 ± 0.02	-1.94 ± 0.03	-3.37 ± 0.01	-2.00 ± 0.06
1431	68007	1194	1240
798.13 ± 196.89	291.03 ± 29.46	162.64 ± 17.19	186.65 ± 21.10
Rings	-2.35 ± 0.04	-2.31 ± 0.02	-3.16 ± 0.01	-2.43 ± 0.04
1431	68007	1194	1240
1134.90 ± 133.10	502.32 ± 24.51	528.50 ± 24.51	559.12 ± 30.90
Table 4: As in Table 2, but with 1 NVP layer.
SNEFY 2	Resampled 2	Gauss 2	GMM 2
Moons	-1.58 ± 0.02	-1.59 ± 0.03	-1.62 ± 0.03	-1.58 ± 0.02
2621	69197	2384	2430
1453.05 ± 147.67	700.98 ± 32.94	831.40 ± 49.24	855.04 ± 53.01
Circles	-1.91 ± 0.04	-2.07 ± 0.04	-2.66 ± 0.05	-1.96 ± 0.04
2621	69197	2384	2430
899.77 ± 196.54	363.91 ± 36.62	267.23 ± 32.05	289.96 ± 33.71
Rings	-2.35 ± 0.04	-2.32 ± 0.03	-2.80 ± 0.02	-2.36 ± 0.03
2621	69197	2384	2430
1255.92 ± 168.80	579.02 ± 23.70	637.73 ± 34.04	658.98 ± 37.69
Table 5: As in Table 2, but with 2 NVP layers.
E.1 2D Unconditional density estimation
Our benchmarking protocol is slightly altered compared with [46]. Firstly, we measure performance over 20 random seeds instead of 1 fixed seed. We find that sometimes the variance over random seeds can be large (e.g. Resampled 0 on Circles). Secondly, rather than computing test performance at the last epoch of training, we follow the more standard procedure of returning the test performance evaluated at the epoch corresponding with the smallest validation performance. This validation/test monitoring results in substantial performance gains in all of the models, with no extra computational cost for SNEFY, Gauss and GMM. For Resampled, monitoring the validation performance to a high precision requires estimating the normalising constant to a high precision, which is computationally challenging. We therefore only check the validation performance every 100 epochs, and compute a high precision normalising constant if the validation performance is the lowest encountered so far. We train each models for a maximum of 20000 iterations (while monitoring validation performance). We use Adam with default hyperparameters and weight decay  $10^{-3}$ . The batch size is  $2^{10}$ .
We use a SNEFY with Gaussian mixture model base measure supported on  $\mathbb{X} = \mathbb{R}^2$ , identity sufficient statistic  $t$  and activation function  $\cos$ . The base measure consists of 8 mixture components, each with a diagonal covariance matrix. We use the same Resampled architecture as in the original paper [46]. We use an MLP with layer widths [2, 256, 256, 1] and sigmoid activations for the resampling distribution. We use a discount factor of 0.1 for the exponential moving average partition function calculation. Note that for Resampled models we are only able to provide an estimate of the test log likelihood, not the exact log likelihood as in other methods. We choose  $n = 50$  and  $m = 1$ . The Gaussian mixture model has 10 mixture components, each with a diagonal covariance matrix. Each normalising flow block consists of an affine coupling block, a permutation layer, and an actnorm layer.
E.2 Modelling distributions on the sphere
We compare the performance of a VMF distribution, a "regular" VMF mixture distribution and our SNEFY construction of the VMF mixture, which allows for some negative weight coefficients as discussed in § 3.2. We use the dataset [44] also used by [19] in a different context for point
SNEFY 4	Resampled 4	Gauss 4	GMM 4
Moons	-1.58 ± 0.03	-1.60 ± 0.07	-1.60 ± 0.03	-1.57 ± 0.03
5001	71577	4764	4810
1616.21 ± 152.61	851.29 ± 33.63	1005.16 ± 30.59	1026.76 ± 28.21
Circles	-1.92 ± 0.04	-1.99 ± 0.03	-2.14 ± 0.16	-1.94 ± 0.04
5001	71577	4764	4810
1035.64 ± 142.39	514.02 ± 47.98	455.58 ± 41.19	480.15 ± 44.61
Rings	-2.33 ± 0.02	-2.31 ± 0.03	-2.59 ± 0.12	-2.32 ± 0.04
5001	71577	4764	4810
1410.36 ± 141.04	749.10 ± 52.05	813.92 ± 27.20	833.36 ± 30.77
Table 6: As in Table 2, but with 4 NVP layers.
SNEFY 8	Resampled 8	Gauss 8	GMM 8
Moons	-1.60 ± 0.02	-1.58 ± 0.02	-1.61 ± 0.08	-1.58 ± 0.03
9761	76337	9524	9570
2036.15 ± 205.52	1198.04 ± 102.71	1399.74 ± 65.78	1423.17 ± 82.84
Circles	-1.91 ± 0.03	-1.97 ± 0.05	-2.16 ± 0.29	-1.93 ± 0.03
9761	76337	9524	9570
1405.27 ± 144.72	839.23 ± 92.99	838.62 ± 82.98	862.25 ± 81.09
Rings	-2.34 ± 0.06	-2.37 ± 0.17	-2.49 ± 0.16	-2.33 ± 0.07
9761	76337	9524	9570
1795.38 ± 186.62	1071.50 ± 109.12	1201.29 ± 71.76	1218.31 ± 70.36
Table 7: As in Table 2, but with 8 NVP layers.
processes. This dataset, retrieved in 2015, is called the "Revised New General Catalogue and Index Catalogue" (RNGC/IC). The RNGC/IC consists of locations of some 10610 galaxies. We use the spherical coordinates of these galaxies and map them to the surface of a sphere.
We compare two types of  $\mathrm{SNEFY}_{\mathbb{S}^2,\exp ,\mathrm{Id},dx}$  with  $m = n = 30$ : one is constrained so that  $\pmb{V}$  is diagonal (and therefore  $\pmb{V}^{\top}\pmb{V}$  is diagonal with nonnegative entries), and the other uses a unconstrained general  $\pmb{V}$ . We expect the unconstrained model to be more expressive, and therefore obtain a better test NLL.
We use Adam with default hyperparameters and a full batch size. We randomly shuffle the data and perform an 80/20 train/test split. Each of the 50 runs uses a randomly sampled initialisation and train/test split. We train for 20000 epochs.
E.3 Conditional density estimation of photometric redshift
Our deep conditional feature extractor is an MLP that uses ReLU activations in each layer and batch normalisation between each layer. Our SNEFY model uses  $\mathrm{Snake_a}$  activations with  $a = 10$ . The CNF models utilise affine layers with tanh activations. For SNEFY and CNF, we preprocess the input features so that they have sample mean zero and sample variance one. We train both deep learning models for 100 epochs using Adam with default hyperparameters and a batch size of 256. The dataset consists of 74309 training and 74557 examples, which are fixed between each run. Each run uses a randomly sampled initialisation (except for CKDE, which is deterministic).
F Complex activation case
Here we consider an extension where we allow the neural network activation  $\sigma$  to be complex-valued, and accordingly the readout parameters  $\mathbf{V}$  to be complex, i.e.  $\mathbf{V} \in \mathbb{C}^{m \times n}$ . Note that in this case
$$
\left\| \boldsymbol {f} (\boldsymbol {t} (\boldsymbol {x}); \boldsymbol {V}, \boldsymbol {\Theta}) \right\| _ {2} ^ {2} = \| \boldsymbol {V} \sigma (\boldsymbol {W} \boldsymbol {t} (\boldsymbol {x}) + \boldsymbol {b}) \| ^ {2} = \sum_ {r = 1} ^ {m} \left| \sum_ {j = 1} ^ {n} v _ {r j} \sigma (\boldsymbol {w} _ {j} ^ {\top} \boldsymbol {t} (\boldsymbol {x}) + b _ {j}) \right| ^ {2}.
$$
Take  $\sigma(u) = \exp(iu)$  and  $\pmb{t}(\pmb{x}) = \pmb{x}$ .
Then the above takes the form
$$
\begin{array}{l} \sum_ {r = 1} ^ {m} \left| \sum_ {j = 1} ^ {n} v _ {r j} \exp \left(i \left(\boldsymbol {w} _ {j} ^ {\top} \boldsymbol {x} + b _ {j}\right)\right) \right| ^ {2} = \sum_ {r = 1} ^ {m} \sum_ {j = 1} ^ {n} \sum_ {l = 1} ^ {n} v _ {r j} \bar {v} _ {r l} e ^ {i \left(\boldsymbol {w} _ {j} ^ {\top} \boldsymbol {x} + b _ {j}\right)} e ^ {- i \left(\boldsymbol {w} _ {l} ^ {\top} \boldsymbol {x} + b _ {l}\right)} \\ = \sum_ {r = 1} ^ {m} \sum_ {j = 1} ^ {n} \sum_ {l = 1} ^ {n} v _ {r j} \bar {v} _ {r l} e ^ {i (b _ {j} - b _ {l})} e ^ {i \left(\left(\boldsymbol {w} _ {j} - \boldsymbol {w} _ {l}\right) ^ {\top} \boldsymbol {x}\right)}. \\ \end{array}
$$
As in the  $\sigma = \exp$  case, bias terms can be folded into the readout parameters  $\mathbf{V}$  (which is the reason why we also require readout parameters to be complex). We note that the case  $n = 1$  is not interesting as it simply reduces to the (normalised) base measure  $\mu$ .
In order to obtain the normalizing constant  $\mathrm{Tr}(\pmb{V}^H\pmb{V}\pmb{K}_{\Theta})$ , we need the integral of the form
$$
[ \boldsymbol {K} _ {\boldsymbol {\Theta}} ] _ {j l} = k _ {\exp (i \cdot), \operatorname {I d}, \mu} (\boldsymbol {w} _ {j}, \boldsymbol {w} _ {l}) = \int e ^ {i \left(\left(\boldsymbol {w} _ {j} - \boldsymbol {w} _ {l}\right) ^ {\top} \boldsymbol {x}\right)} \mu (d \boldsymbol {x}) =: \kappa (\boldsymbol {w} _ {j} - \boldsymbol {w} _ {l}), \tag {18}
$$
where  $\kappa$  is simply the Fourier transform of the base measure  $\mu$ , i.e. its characteristic function in the case where  $\mu$  is a probability measure. Hence many standard choices of  $\mu$  lead to tractable normalizing constants. Note that while  $\mathbf{V}$  and  $\kappa$  both may be complex-valued, the normalizing constant as well as the density itself are real-valued.
Various examples of probability measures  $\mu$  and its Fourier transforms that give rise to shift-invariant positive definite kernels  $\kappa$  have been studied in the literature on Random Fourier Features (RFF) [38]. Here too, while the functional form of the expressions is identical, the role between the data and the parameters is reversed: in RFF, one is interested in approximating a given kernel on the data instances, by considering a probability measure on the frequency space which is that kernel's inverse Fourier transform.
Remark 3. If we restrict our attention to real-valued  $V$  and thus explicitly maintain biases inside the parameterisation, this model can also be realised with stacking cos and sin activation so that they share the readout parameters since
$$
\sum_ {r = 1} ^ {m} \left| \sum_ {j = 1} ^ {n} v _ {r j} \exp \left(i \left(\boldsymbol {w} _ {j} ^ {\top} \boldsymbol {x} + b _ {j}\right)\right) \right| ^ {2} = \| \boldsymbol {V} \cos (\boldsymbol {W} \boldsymbol {x} + \boldsymbol {b}) \| ^ {2} + \| \boldsymbol {V} \sin (\boldsymbol {W} \boldsymbol {x} + \boldsymbol {b}) \| ^ {2}.
$$
G Discrete and mixed continuous/discrete SNEFYs
G.1 Discrete SNEFYs via series
Closed-form expressions for normalising constants of discrete distributions are only advantageous if their computation costs less than naively summing up unnormalised density over all possible states. This is the case if the support has very large or infinite cardinality. An example that extends the classical Poisson distribution in exponential family form is given in Table 1. Here we discuss some other settings.
NNK as a convergent series Suppose the support  $\mathbb{X}$  is discrete and let  $h(x)$  be a nonnegative function corresponding with the discrete base measure  $\mu$ . The NNK is given by
$$
k _ {\sigma , \boldsymbol {t}, \mu} (\boldsymbol {\theta} _ {i}, \boldsymbol {\theta} _ {j}) = \sum_ {\boldsymbol {x} \in \mathbb {X}} \sigma \left(\boldsymbol {w} _ {1} ^ {\top} \boldsymbol {t} (\boldsymbol {x}) + b _ {1}\right) \sigma \left(\boldsymbol {w} _ {2} ^ {\top} \boldsymbol {t} (\boldsymbol {x}) + b _ {2}\right) h (\boldsymbol {x}).
$$
Such NNKs often resemble known convergent series.
Fourier series For example, if  $\sigma = \cos$  and  $C_x = h(x)$  are some coefficients of a convergent Fourier series, then
$$
k _ {\sigma , \boldsymbol {t}, \mu} (\boldsymbol {\theta} _ {i}, \boldsymbol {\theta} _ {j}) = \frac {1}{2} \sum_ {x = 0} ^ {\infty} C _ {x} \left(\cos (t (x) (w _ {1} - w _ {2}) + (b _ {1} - b _ {2})) + \cos (t (x) (w _ {1} + w _ {2}) + (b _ {1} + b _ {2}))\right)
$$
is just a series representation of a sum of two periodic functions. For example, if  $C_x = \frac{2(-1)^{x+1}}{\pi x}$  for  $x \geq 1$  and  $C_x = 0$  otherwise, and  $t(x) = 2\pi x$ , the NNK is a sum of two sawtooth waves with frequencies  $w_1 - w_2$  and  $w_1 + w_2$  and phase offsets  $b_1 - b_2$  and  $b_1 + b_2$ . Other examples include rectified sine waves, square waves and triangular waves. This extends to periodic functions with convergent multivariate Fourier series.
G.2 Mixed continuous/discrete SNEFYs
Mixed distributions can be obtained by choosing the base measure to be a mixed continuous distribution. For example, choose  $\mathbb{X} = \mathbb{R}^d$ ,  $\mu (d\mathbf{x}) = \frac{1}{2}\bigl (\Phi (\mathbf{x}) + \delta (\mathbf{x})\bigr)d\mathbf{x}$ , and take  $\sigma$  and  $\mathbf{t}$  to be any of the combinations leading to a closed-form NNGPK (for example,  $\mathbf{t}$  as the identity and  $\sigma$  as the error function, Leaky ReLU, GELU, cosine, or snake). Then
$$
\begin{array}{l} k _ {\sigma , \mathbf {t}, \mu} (\theta_ {i}, \theta_ {j}) = \frac {1}{2} \left(\int_ {\mathbb {X}} \sigma \left(\mathbf {W} _ {i} ^ {\top} \mathbf {t} (\mathbf {x}) + b _ {i}\right) \sigma \left(\mathbf {W} _ {j} ^ {\top} \mathbf {t} (\mathbf {x}) + b _ {j}\right) \delta (\mathbf {x}) d \mathbf {x} + \right. \\ \int_ {\mathbb {X}} \sigma \left(\mathbf {W} _ {i} ^ {\top} \mathbf {t} (\mathbf {x}) + b _ {i}\right) \sigma \left(\mathbf {W} _ {j} ^ {\top} \mathbf {t} (\mathbf {x}) + b _ {j}\right) \Phi (\mathbf {x}) d \mathbf {x} \Bigg) \\ = \frac {1}{2} \left(\sigma \left(\mathbf {W} _ {i} ^ {\top} \mathbf {t} (\mathbf {0}) + b _ {i}\right) \sigma \left(\mathbf {W} _ {j} ^ {\top} \mathbf {t} (\mathbf {0}) + b _ {j}\right) + k _ {\sigma , \mathbf {t}, \Phi} \left(\theta_ {i}, \theta_ {j}\right)\right), \\ \end{array}
$$
which is a closed form.