SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
P (d \boldsymbol {x}; \boldsymbol {V}, \boldsymbol {W}) = \frac {1}{\operatorname {T r} (\boldsymbol {V} ^ {\top} \boldsymbol {V} \boldsymbol {K} _ {\Theta})} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} \boldsymbol {v} _ {., i} ^ {\top} \boldsymbol {v} _ {., j} \exp \left(\frac {1}{2} \left(\boldsymbol {w} _ {i} + \boldsymbol {w} _ {j}\right) ^ {\top} \boldsymbol {t} (\boldsymbol {x})\right) \mu (d \boldsymbol {x}), \tag {8}
$$
4Note that  $\exp (b_i)$  is independent of  $x$  and only appears as a product with  $v_{i}$
which is a mixture $^5$  of distributions  $P_e\left(\cdot; \frac{1}{2}(\boldsymbol{w}_i + \boldsymbol{w}_j)\right)$  belonging to a classical exponential family  $P_e$  [50, 32], given in the canonical form by
$$
P _ {e} (d \boldsymbol {x}; \boldsymbol {w}) = \frac {1}{\mathrm {z} _ {e} (\boldsymbol {w})} \exp \left(\boldsymbol {w} ^ {\top} \boldsymbol {t} (\boldsymbol {x})\right) \mu (d \boldsymbol {x}), \quad \mathrm {z} _ {e} (\boldsymbol {w}) = \int_ {\mathbb {X}} \exp \left(\boldsymbol {w} ^ {\top} \boldsymbol {t} (\boldsymbol {x})\right) \mu (d \boldsymbol {x}). \tag {9}
$$
It is helpful to identify the following three further cases:
1. When  $n = m = 1$ ,  $v_{11}^2$  cancels in the numerator and denominator and we obtain an exponential family with base measure  $\mu$  supported on  $\mathbb{X}$ , sufficient statistic  $t$ , canonical parameter  $w_1$  and normalising constant  $z_e(w_1)$ . Every exponential family is thus a SNEFY, but not conversely.
2. When  $m > 1$  and  $n > 1$ , we obtain a type of exponential family mixture model with coefficients  $V^{\top}V$ , some of which may be negative. Advantages of allowing negative weights in mixture models in terms of learning rates are discussed in [42]. The rank of  $V^{\top}V$  is at most  $\min(m, n)$ .
3. When  $m > 1$  and  $n > 1$  and  $V^{\top}V$  is diagonal (i.e. each column in  $V$  is orthogonal), there are at most  $n$  non-zero mixture coefficients, all of which are nonnegative. That is, we obtain a standard exponential family mixture model.
The kernel matrix  $K_{\Theta}$  in the normalising constant of (8) is tractable whenever the normalising constant of the corresponding exponential family is itself tractable.
Proposition 1. Denote by  $\mathbf{z}_e(\boldsymbol{w})$  the normalising constant of the exponential family in (9). Then
$$
k _ {\exp (\cdot / 2), \boldsymbol {t}, \mu} \left(\boldsymbol {w} _ {i}, \boldsymbol {w} _ {j}\right) = z _ {e} \left(\frac {1}{2} \left(\boldsymbol {w} _ {i} + \boldsymbol {w} _ {j}\right)\right). \tag {10}
$$
The above kernel is well defined for any collection of  $\boldsymbol{w}_i$  which belong to the canonical parameter space of  $P_e$ , since the canonical parameter space is always convex [50]. This gives us a large number of tractable instances of SNEFY which correspond to exponential family mixture models allowing negative weights – a selection of examples is given in Appendix C. It is interesting that some properties of the exponential families are retained by this generalisation belonging to SNEFYs. For example, the following proposition, the proof of which is given in Appendix A, links the derivatives of the log-normalising constant to the mean and the covariance of the sufficient statistic.
Proposition 2. Let  $\sigma(u) = \exp(u/2)$  and define the log-normalising constant as  $\Psi = \log z(V, \Theta)$ .
$$
T h e n \quad \sum_ {i = 1} ^ {n} \frac {\partial \Psi}{\partial \boldsymbol {w} _ {i}} = \mathbb {E} [ \boldsymbol {t} (\times) ] \quad a n d \quad \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} \frac {\partial^ {2} \Psi}{\partial \boldsymbol {w} _ {i} \boldsymbol {w} _ {j} ^ {\top}} = \mathbb {E} [ \boldsymbol {t} (\times) \boldsymbol {t} (\times) ^ {\top} ] - \mathbb {E} [ \boldsymbol {t} (\times) ] \mathbb {E} [ \boldsymbol {t} (\times) ] ^ {\top}.
$$
3.3 Conditional distributions under SNEFY
An attractive property of SNEFY is that, under mild conditions, the family is closed under conditioning.
Theorem 1. Let  $\mathbf{x} = (\mathbf{x}_1, \mathbf{x}_2)$  be jointly  $SNEFY_{\mathbb{X}_1 \times \mathbb{X}_2}, t, \sigma, \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(x_1), t_2(x_2))$ . Then the conditional distribution of  $\mathbf{x}_1$  given  $\mathbf{x}_2 = \mathbf{x}_2$  is  $SNEFY_{\mathbb{X}_1, t_1, \sigma, \mu_1}$  with parameters  $\mathbf{V}$  and  $\Theta_{1|2} \triangleq (W_1, W_2 t_2(\mathbf{x}_2) + b)$ .
The proof, which we detail in Appendix A follows directly by folding the dependence on the conditioning variable  $x_{2}$  into the bias term. We note that conditional density will typically be tractable if the joint density is tractable since they share the same activation function  $\sigma$ . Thus, whenever SNEFY corresponds to a tractable NNK with a non-zero bias, we can construct highly flexible conditional density models using SNEFY by taking  $t_{2}$  itself to be a jointly trained deep neural network. Crucially,  $t_{2}$  may be completely unconstrained. We use this observation in the experiments (§ 4).
Here we are concerned with the setting where every term in the mixture model belongs to the same family, i.e. an exponential family mixture model but not a mixture of distinct exponential families.
3.4 Marginal distributions under SNEFY
Marginal distributions under SNEFY model for a general activation function  $\sigma$  need not belong to the same family. In the special case  $\sigma = \exp (\cdot /2)$ , SNEFY is in fact also closed under marginalisation, which we prove in Appendix D. Even in the general  $\sigma$  case, marginal distributions are tractable and admit closed forms whenever the joint SNEFY model and the conditional SNEFY are tractable.
Theorem 2. Let  $\mathbf{x} = (\mathbf{x}_1, \mathbf{x}_2)$  be jointly  $SNEFY_{\mathbb{X}_1 \times \mathbb{X}_2, t, \sigma, \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  $\mathbf{t}(\mathbf{x}) = (\mathbf{t}_1(\mathbf{x}_1), \mathbf{t}_2(\mathbf{x}_2))$ . Then the marginal distribution of  $\mathbf{x}_1$  is
$$
P _ {1} (d \boldsymbol {x} _ {1}) = \frac {\operatorname {T r} \left(\boldsymbol {V} ^ {\top} \boldsymbol {V} \widetilde {\boldsymbol {C}} _ {\boldsymbol {\Theta}} (\boldsymbol {x} _ {1})\right)}{\mathrm {z} (\boldsymbol {V} , \boldsymbol {\Theta})} \mu_ {1} (d \boldsymbol {x} _ {1}),
$$
where  $\widetilde{C} (\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).$
The proof is given in Appendix A. Due to this tractability of the marginal distributions, it is straightforward to include the likelihood corresponding to incomplete observations (i.e. samples where we are missing some of the components of  $\mathbf{x}$ ) into the density estimation task.
3.5 Connections with kernel-based methods for nonnegative functions
SNEFY may be viewed as a neural network variant of the non-parametric kernel models for non-negative functions [25], which are constructed as follows. Let  $\psi : \mathbb{X} \to \mathbb{H}$  be a feature mapping to a (possibly infinite dimensional) Hilbert space  $\mathbb{H}$ . Let  $\mathbb{S}(\mathbb{H})$  be the set of all positive semidefinite (PSD) bounded linear operators  $A: \mathbb{H} \to \mathbb{H}$ . Then
$$
h _ {\boldsymbol {A}} (\boldsymbol {x}) = \langle \psi (\boldsymbol {x}), \boldsymbol {A} \psi (\boldsymbol {x}) \rangle_ {\mathbb {H}} \tag {11}
$$
gives an elegant model for nonnegative functions parametrised by  $A \in \mathbb{S}(\mathbb{H})$ , and their application to density modelling with respect to a base measure has also been explored [25, 42]. Note that by assuming boundedness of  $A$  and  $\psi$ , the normalizing constant [25, Proposition 4] is given by  $\operatorname{Tr}\left(A \int_{\mathbb{X}} \psi(x) \otimes \psi(x) \mu(dx)\right)$ , which is analogous to our work where the normalising constant is given by  $\mathrm{z}(V, \Theta) = \operatorname{Tr}(V^\top VK_\Theta)$ . This can be seen by replacing  $A$  with  $V^\top V$  and replacing  $\int_{\mathbb{X}} \psi(x) \otimes \psi(x) \mu(dx)$  by  $K_\Theta = \int_{\mathbb{X}} \sigma(Wt(x) + b) \sigma(Wt(x) + b)^\top \mu(dx)$ .
Figure 1: (Left) An instance of an (untrained)  $\mathrm{SNEFY}_{\mathbb{R}^2, \mathrm{Id}, \mathrm{Snake}_a, \Phi}$  density with  $n = 100$ ,  $m = 1$ ,  $v_{ij} \sim \mathcal{N}(0, 1/n)$ ,  $w_{ij} \sim \mathcal{N}(0, 4)$  and  $b = 0$ . Shown are 50 exact samples found using rejection sampling. Numerical quadrature for this and every example supported on  $\mathbb{R}^d$  in §4 returns a value of 1.00 for the integral over  $\mathbb{X}$ . (Right) A trained  $\mathrm{SNEFY}_{\mathbb{S}^2, \mathrm{Id}, \exp, dx}$  density with  $n = m = 30$ . Shown is the training and testing dataset [44] also used by [19] for point processes.
Despite feature maps being infinite-dimensional,
model (11) often reduces to an equivalent representation in finite-dimensions. [25] utilise a representer theorem when (11) is fitted to data using a regularised objective, while [42] more directly assume that the linear operator  $\mathbf{A}$  is inside the span of the features evaluated at the available data  $\{\pmb{x}_{\ell}\}_{\ell = 1}^{N}$ . The resulting model resembles SNEFY where  $n$  equals to the number  $N$  of datapoints, i.e.
$$
h _ {M} (\boldsymbol {x}) = \left[ \kappa (\boldsymbol {x}, \boldsymbol {x} _ {1}), \dots , \dots , \kappa (\boldsymbol {x}, \boldsymbol {x} _ {N}) \right] ^ {\top} M \left[ \kappa (\boldsymbol {x}, \boldsymbol {x} _ {1}), \dots , \dots , \kappa (\boldsymbol {x}, \boldsymbol {x} _ {N}) \right] \tag {12}
$$
for a PSD matrix  $M \in \mathbb{R}^{N \times N}$  and  $\kappa(\pmb{x}_i, \pmb{x}_j) = \langle \pmb{\psi}(\pmb{x}_i), \pmb{\psi}(\pmb{x}) \rangle_{\mathbb{H}}$ .
However, there are fundamental differences between (12) and SNEFY which we list below. The models can be seen as complementary and they inherit advantages and disadvantages of kernel methods and neural networks common in other settings, respectively.
- Tractability. Whereas we identify many tractable instances of SNEFY, the normalising constant of (12) requires computing  $\int \kappa(\boldsymbol{x}_i, \boldsymbol{x}) \kappa(\boldsymbol{x}, \boldsymbol{x}_j) \mu(dx)$  - this is not generally tractable apart from some
limited combinations of  $\kappa$  and  $\mu$  (e.g. for a Gaussian kernel  $\kappa$  and a Gaussian  $\mu$ ). Note that this kernel is evaluated at the datapoints, whereas SNEFY evaluates the kernel at the learned parameters  $\boldsymbol{w}_i$ . [42] focuses on the specific case where  $\kappa$  is a Gaussian kernel, studying properties of the resulting density class which is a mixture of Gaussian densities allowing for negative weights, a model equivalent to SNEFY with the exponential activation function as described in Appendix C. Note that our treatment of SNEFY as a generalisation of the exponential family goes well beyond the Gaussian case, and that tractable instances arise with many other activation functions.
- Expressivity. Crucially, the feature map  $\psi$  and consequently finite dimensional representation via  $\kappa$  in (12) are treated as fixed feature maps and are not themselves learned – instead, expressivity in (12) only comes from fitting  $M$  (and potentially lengthscale hyperparameters of  $\kappa$ ), at the expense of a more involved optimisation over the space of PSD matrices. In contrast, we learn  $W$  (analogous to learning  $\psi$ ) and  $V$  jointly using neural-network style gradient optimisers. SNEFY is fully compatible with end-to-end and jointly optimised neural network frameworks, a property we leverage heavily in our experiments in § 4.
- Conditioning. By explicitly writing parametrisation which includes the biases, we obtain a family closed under conditioning and thus a natural model for conditional densities, whereas it is less clear how to approach conditioning when given a generic feature map  $\psi$ .
Other related work After submission, we became aware of another related literature, which includes mixture models with potentially negative mixture coefficients via squaring [23] and positive semi-definite probabilistic circuits [43]. We believe a marriage of ideas from SNEFY and probabilistic circuits will lead to future developments in tractable and expressive probability models.
Implementation All experiments are conducted on a Dual Xeon 14-core E5-2690 with 30GB of reserved RAM and a single NVidia Tesla P100 GPU. Full experimental details are given in Appendix E. We build our implementation on top of normflows [45], a PyTorch package for normalising flows. SNEFYs are built as a BaseDistribution, which are base distributions inside a NormalizingFlow with greater than or equal to zero layers. We train all models via maximum likelihood estimation (MLE) i.e. minimising forward KL divergence.
2D synthetic unconditional density estimation We consider the 2 dimensional problems also benchmarked in [46]. We compare the test performance, computation time and parameter count of non-volume preserving flows (NVPs) [7] with four types of base distribution: SNEFY, resampled [46] (Resampled), diagonal Gaussian (Gauss) and Gaussian mixture model (GMM). We consider flow depths of 0, 1, 2, 4, 8 and 32 where a flow depth of 0 corresponds with the base distribution only. We use  $\sigma = \cos$ ,  $t = \mathrm{Id}$ ,  $\mathbb{X} = \mathbb{R}^d$  and a Gaussian mixture model base density. We set  $m = 1$  and  $n = 50$ . Full architectures and further experimental details are described in Appendix E.1.
Results are shown in Table 2 for the 0 and 16 layer cases, and further tables for 1, 2, 4 and 8 layers are given in Appendix E.1. Our observations are that all base distributions are able to achieve good performance, provided they are appended with a normalising flow of appropriate depth. SNEFY is able to achieve good performance with depth 0 on all three datasets, as are Resampled and GMM for the Moons dataset. The parameter count for well-performing SNEFY models is very low, but the computation time can be relatively high. However, SNEFY is the only model which consistently achieves the highest performance within one standard deviation across all normalising flow depths.
Data on the sphere We compare the three cases mentioned in § 3.2 using Kernel 3, i.e. mixtures of the von Mises Fisher (VMF) distribution, as shown in Figure 1. Over 50 runs, the unconstrained  $V$ , diagonal  $V$ , and  $n = m = 1$  cases respectively obtain test negative log likelihoods of  $1.38 \pm 9.64 \times 10^{-3}$ ,  $1.46 \pm 0.016$  and  $2.34 \pm 0.26$  each in  $111.18 \pm 2.58$ ,  $109.12 \pm 0.80$  and  $61.88 \pm 0.33$  seconds (average ± standard deviation). In this setting, allowing for a fully flexible  $V$ , going beyond the classical mixture model, shows clear benefits in performance. Results are summarised in Table 3. Full details are given in Appendix E.2.
Conditional density estimation on astronomy data Predicting plausible values for the velocity of distant astronomical objects (such as galaxies or quasars) without measuring their full spectra, but by
SNEFY 0	Resampled 0	Gauss 0	GMM 0	SNEFY 16	Resampled 16	Gauss 16	GMM 16
Moons	-1.59 ± 0.02	-1.60 ± 0.02	-3.29 ± 0.02	-1.59 ± 0.04	-1.57 ± 0.03	-1.59 ± 0.02	-1.68 ± 0.24	-1.57 ± 0.03
241	66817	4	50	19281	85857	19044	19090
1200.84 ± 109.77	541.11 ± 12.20	606.35 ± 20.29	629.11 ± 14.57	2844.09 ± 254.02	1867.26 ± 167.75	2226.37 ± 174.80	2276.78 ± 211.15
Circles	-1.92 ± 0.03	-2.42 ± 0.22	-3.55 ± 0.01	-2.07 ± 0.07	-1.92 ± 0.04	-1.96 ± 0.04	-2.22 ± 0.30	-1.93 ± 0.03
241	66817	4	50	19281	85857	19044	19090
643.80 ± 150.84	166.35 ± 20.59	52.63 ± 3.96	73.52 ± 5.47	2272.15 ± 281.76	1533.30 ± 158.90	1660.77 ± 169.02	1673.54 ± 166.84
Rings	-2.41 ± 0.05	-2.62 ± 0.01	-3.26 ± 0.01	-2.67 ± 0.03	-2.34 ± 0.10	-2.36 ± 0.14	-2.51 ± 0.22	-2.31 ± 0.04
241	66817	4	50	19281	85857	19044	19090
997.73 ± 107.80	409.77 ± 19.37	416.23 ± 18.11	433.82 ± 16.06	2654.22 ± 269.14	1774.06 ± 182.25	2045.65 ± 170.80	2070.60 ± 177.06
Table 2: The first quantity is the average  $\pm$  sample standard deviation over 20 runs of test loglikelihood. The second quantity is the parameter count. The third quantity is the average  $\pm$  sample standard deviation over 20 runs of the computation time (seconds). The number in each column header is the number of non-volume preserving flow layers appended after the base distribution. Here there are 0 or 16 NVP layers. More tables showing intermediate layer results are given in Appendix E.1.
Method	Average Test NLL ↓	Compute time (s)	Parameter count
SNEFY n = 32	-2.195 ± 0.024	495.720 ± 22.561	179748
SNEFY n = 16	-2.172 ± 0.034	404.291 ± 57.605	46356
SNEFY n = 8	-2.108 ± 0.089	390.870 ± 16.028	12300
CNF L = 4	-2.156 ± 0.018	202.1290 ± 10.380	1413
CNF L = 2	-2.163 ± 0.024	155.090 ± 15.809	1155
CNF L = 1	-2.171 ± 0.012	122.304 ± 1.194	1026
CKDE	-2.148	391.867 (Not GPU-accelerated)	Train set size = 74309 × 6
Table 3: Performance comparison of methods on astronomy dataset. Excluding CKDE which is deterministic, an average is taken over 50 random initialisations with  $\pm$  indicating standard deviation. SNEFY shows a statistically significant average increase in performance over CNF (using a two-sample t-test), and over CKDE (using a one-sample t-test). Full experimental details in Appendix E.3.
measuring shifted spectra through broadband filters, is known as photometric redshift. Photometric redshift is an important problem in modern astronomy, as large surveys have increased the amount of available data that do not directly measure spectra. We estimate distributions for cosmological redshift  $x_{1} \in \mathbb{R}$  conditional on features  $x_{2} \in \mathbb{R}^{5}$  using a public dataset [2]. The features  $x_{2}$  are the magnitude a broadband filter (red) and set of four pairwise distances between broadband colour filters (ultraviolet-green, green-red, red-infrared1, infrared1-infrared2). We benchmark our own SNEFY against a conditional kernel density estimator (CKDE) and a conditional normalising flow (CNF) [54]. We use a publically available implementation [10] of CNF. Our SNEFY here is not appended with any normalising flow layers. We use  $\sigma = \mathrm{Snake}_{\mathrm{a}}$ ,  $t = \mathrm{Id}$ ,  $\mathbb{X} = \mathbb{R}$  and a Gaussian base density. We use a deep conditional feature extractor with layer widths  $[2^{5}n, 2^{4}n, 2^{3}n, 2^{2}n, 2n]$ , and then set  $m = n/2$ . Full details are given in Appendix E.3.
While CNFs have shown great promise in modelling high dimensional image data, we expect that they are not as well-suited to nontrivial tabular data with a small to medium number of dimensions. This is because each layer is required to be an invertible transformation in the variable being modelled, so the input and output sizes must be the same, thereby significantly limiting the parameter count in each layer. We could of course increase the number of layers in the CNF model to achieve higher parameter counts, however this results in a model that is difficult to train. In our experiments, we found that increasing the depth of the CNF decreased its performance. On the other hand, SNEFY may utilise any number of parameters, as the conditioning network  $t_2$  is completely unrestricted.
Density estimation on astronomy data using partial (marginal) observations We perform joint probability density estimation using the redshift data. The original training dataset is a matrix in  $\mathbb{R}^{74309\times 6}$ . We partition this matrix into batches of size 256 and randomly set each column of each batch to NaN with probability  $q$  (i.e. the corresponding dimension is missing from observations). We then train a SNEFY model that utilises partial observations by maximising the marginal likelihood according to Theorem 2. We leave the original  $74557\times 6$  testing dataset untouched, and measure the test NLL after training. We plot the NLL as a function of  $1 - q$ , as shown in Figure 2. We also compare the performance of the SNEFY that uses partial observations with the performance of SNEFY that simply throws away incomplete observations, as well as the same normalising flow baselines that throw away incomplete observations that were used in the conditional density estimation setting. There do exist normalising flow approaches that jointly optimises conditional distributions for missing data and
Figure 2: Density estimation under partial observations. The right plot is a zoomed in version of the left plot. NF-1, NF-2 and NF-4 are respectively normalising flows of depth 1, 2 and 4. Normalising flow models and SNEFY without marginalisation discard incomplete observations, whereas SNEFY use Theorem 2 to include partial observations via maximum marginal likelihood. Marginal likelihoods allow for an improved NLL.
model parameters [3] or samples from missing data and optimises for model parameters [41], however these do not allow for maximum likelihood estimation. Our observations are twofold: including partial observations improves performance, and adding more complete observations improves performance.
We constructed a new class of probability models – SNEFY – by normalising the squared norm of a neural network with respect to a base measure. SNEFY possesses a number of convenient properties: tractable exact normalising constants in many instances we identified, closure under conditioning, tractable marginalisation, and intriguing connections with existing models. SNEFY shows promise empirically, outperforming competing (conditional) density estimation methods in experiments.
Sampling versus density estimation We focus here on the problem of density estimation, for which SNEFY is well-suited. While it is sometimes possible to obtain exact SNEFY samples using rejection sampling, sampling is more computationally expensive than in other models such as normalising flows. Future work will focus on sampling, as has been done with related models [26], where the special case of Gaussian  $\sigma$  and hyperrectangular support  $\mathbb{X}$  is considered. In [26],  $r$  approximate samples with  $\mathcal{O}\left(r\log_2|\mathbb{X}| + rd\log_2\frac{2}{\rho}\right)$  evaluations of the normalising constant are obtained. Here  $\rho$  is an approximation tolerance parameter. Note that this complexity significantly improves upon naive rejection sampling, which typically scales exponentially in dimension  $d$ . In the unconditional setting, SNEFY is constructed as only a 2-layer network, thereby limiting expressivity. However, in the conditional density estimation setting, we may use any number of layers and any architecture for the conditioning network  $t_2$ . Finally, SNEFY inherits all the usual limitations and advantages over mirroring kernel-based approaches, as discussed in § 3.5.