Squared Neural Families: A New Class of Tractable Density Models
Authors: Russell Tsuchida
Cheng Soon Ong
Dino Sejdinovic
Venue: NeurIPS 2023 spotlight (2023)
Keywords: probabilistic modelling; density estimation; exponential family;
TL;DR: We can compute normalising constants of squared neural networks to construct probability distributions with interesting and useful properties.

Abstract: Flexible models for probability distributions are an essential ingredient in many machine learning tasks. We develop and investigate a new class of probability distributions, which we call a Squared Neural Family (SNEFY), formed by squaring the 2-norm of a neural network and normalising it with respect to a base measure. Following the reasoning similar to the well established connections between infinitely wide neural networks and Gaussian processes, we show that SNEFYs admit closed form normalising constants in many cases of interest, thereby resulting in flexible yet fully tractable density models. SNEFYs strictly generalise classical exponential families, are closed under conditioning, and have tractable marginal distributions. Their utility is illustrated on a variety of density estimation, conditional density estimation, and density estimation with missing data tasks.

Squared Neural Families: A New Class of Tractable Density Models
Russell Tsuchida" Data61-CSIRO
Cheng Soon Ong Data61- & Australian National University
Dino Sejdinovic School of CMS & AIML, The University of Adelaide
Abstract
Flexible models for probability distributions are an essential ingredient many machine learning tasks. We develop and investigate new class of probability distri- butions, which we call a Squared Neural Family (SNEFY), formed by squaring the 2-norm of neural network and normalising it with respect to base measure. Fol- lowing the reasoning similar to the well established connections between infinitely wide neural networks and Gaussian processes, show that SNEFYs admit closed form normalising constants in many cases of interest, thereby resulting in flexible yet fully tractable density models. SNEFYs strictly generalise classical exponential families, are closed under conditioning, and have tractable marginal distributions. Their utility is illustrated on a variety of density estimation, conditional density estimation, and density estimation with missing data tasks.
Introduction
Probabilistic modelling lies at the heart of machine learning. In both traditional and contemporary settings, ensuring the probability model is appropriately normalised (or otherwise bypassing the need to compute normalising constants) is of central interest for maximum likelihood estimation and related statistical inference procedures. Tractable normalising constants allow the use of probability models in a variety of applications, such as anomaly detection, denoising and generative modelling.
Traditional statistical approaches [27 5 rely mathematically convenient models such as exponen- tial families (50). Such models can often be normalised closed form, but are often only suitable for relatively fow-dimensional and simple data. Early approaches in deep learning use neural networks as energy functions inside Gibbs distributions [21 Equation 2]. Such distributions typically have very intractable normalising constants, and so either surrogate losses for the negative log likelihood involving the energy function are used or MCMC, score matching [16), variational or sampling
methods [15 30 are used to approximate the normalising constant. See 335 $2.6.1, $2.2] for more energy-based and other methods.
Modern computer technology and power allows more flexible models g page 68), by partially or completely relaxing the requirement for mathematical tractability. Contemporary
modelling, such generative image models, rely primarily on neural network models [18 35 34)
For example, normalising flows [40] use constrained neural network layers to transform random variables from base measures with tractable Jacobian determinants associated with appropriately constrained but flexible pushforward maps. Similarly to normalising flows, we use neural networks to define expressive densities, but model the densities directly without using constrained transfor- mations of random variables. Morcover, our approach can be readily applied in conjunction with normalising flows and other deep learning devices such as deep feature extractors.
SSfftare available at httppp://ggtaaa..cm//usesseeeeee
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Support X
Base measure
Sufficient statistic ((x)
Activation
Kernel
function o
b
of 0
ka.t.p.
Any setting mirroring previously derived NNGPK,
e.g.
erf
X
52
C. P N
X
4
Rd
cco
Id(x)
LReLU
X
48
GELU 14
X
47
COS
X
337
Any tractable exponential family (50 32 setting, e.g.
R
Pc
at
Kernel
2
S
Uniform
Id x))
exp
Rdd
C.
V
Kernel
10, 1, .....
(z!)
Kernel
New tractable integration settings, e-g.
Rd
&c
Id(æ)
cos
Kernel
00, 1]d Uniform
Φ (x))
Rd
&cmm
Td(z)
Snake,
55
39
Kernel
Table 1: Examples of settings admitting closed-form for the normalising constant (VV, Q) by leveraging a closed-form NNK kg,t.. (4). In each case, z( V, θ) = Tr(V V Vee), where the entries of the matrix Ke are described according to the NNK ki,,t. in Identity [ & m denotes the CDF of multivariate normal distribution with mean m and covariance matrix C and v/ denotes counting measure. Rows with citations have been considered previously in the context of NNGPKs, but not as normalising constants and not with a reversal of the role of input and parameter. Note the cases where b # O; this setting is not considered by others, because when the role of parameters and data is in the usual setting, = O covers a sufficiently general setting. Noticing that SWEFYs strictly generalise exponential family mixture models (see $ 3.22). fixing af-) exp(-/2) and a given base measure 1u and sufficient statistic t for which the exponential family log- peg-partition function is known also leads tractable normalising constants. More known closed-form kernels as well as approximate kernels that can be adapted to our setting are given in [2]
Contributions. Let (1) denote measure space with C Rd, sigma algebra F, and ative measure μ. Let V Rm be the readout parameters of neural network and let W e g=* and b I" be the weights and biases
hidden layer of neural network f RD - R with acti- vation function a. For some support X E F. define probability distribution /' (and corresponding probability density p with respect to base measure (4) to be proportional to squared 2-norm of the evaluation of the neural network f.
p(dx)
P(dæ; V V) A z(V, )) | ((())
-o
e (W,b), (1)
whenever the normalising constant z(V, () - ((()) || ((((()))) is finite and non-zero.
Here we call μ the base measure, t: + R the sufficient statisti and a the activation function. We will call the corresponding family of probability distributions, parametrised by (V,0), a squared neural family (SNEFY) on X, and denote it by SNEFYx,t..... SNEFYs are new flexible probability distribution models that strike good balance between mathematical tractability, expressivity and computational efficiency. When a random vector follows a SNEFY distribution indexed by parameters (V, O), we write x SNEFY t, (V. + or simply P(: V. θ), where there is no ambiguity.
Our main technical challenge is in exactly computing the normalising constant, (V O), where
z( V. ) A f [I) V, (K)(d f(t: V. θ) bo θ (W, b).
(2)
The normalising constants we consider are special cases in the sense that they apply to specific (but commonly appearing in applications) choices of activation function a, sufficient statistic t and base measure i over support X. See Table [
Our analysis both exploits and informs a connection with
We later verify that & indeed sufficient statistic,
Note t maps from d to D dimensions.

---

so-called neural network Gaussian process kernels (NNGPKs) [1 in generalised form, which we refer as neural network kernels (NNKs).
We discuss some important theoretical properties of SNEFY such as exact normalising constant calcula- tion, marginal distributions, conditional distributions, and connections with other probability models. We then consider deep learning setting, where SNEFYs can either be used as base distributions in (non-volume preserving) normalising flows [46), or may describe flexible conditional density models with deep learning feature extractors. We demonstrate SNEFY on a variety of datasets.
1.1 Background
Notation We use lower case non-bold (like a) to denote scalars, lower case bold to denote vectors (like a) and upper case bold to denote matrices (like A). Random variables (scalar or vector) are additionally typeset in sans-serif (like and a). The special zero vector .O...., 0) and identity matrix elements are 0 and We use subscripts to extract (groups of) indices, so that for example, w; is the ith row of the matrix W (as column vector), U. is the ith column of the matrix V, and b; is the ith element of the vector b. We use O (W, b) (D+(1 to denote the concatenated hidden layer weights and biases. Correspondingly, we write 0 = (w, by) E RD+1 for the ith row of O. We will use a number of special functions. &c denotes the cumulative distribution function (CDF) of a Gaussian random vector with mean m and covariance matrix C. We also use shorthand c c and t 1
Single hidden layer neural networks
We consider a feedforward neural network f R + R"
f(t: θ) Vo(Wt + b)
with activation function o, hidden weights W E Rax and biases E R" and readout parameters V Rmxa, Here is applied element-wise to its vector inputs, returning vector of the same shape.
Neural network kernels In certain theories of deep learning, one often encounters bivariate Gaussian integral called the neural network Gaussian procexs kernel NNGPK. The NNGPK first arose as the covariance function of a well-behaved single layer neural network with random weights [31]. In the limit as the width of the network grows to infinity, the neural network (3) with suitably well- behaved (say, independent Gaussian) random weights converges to zero-ment Gaussian process, so that the NNGPK characterises the law the neural network predictions. These limiting models can be used as functional priors in classical Gaussian process sense [53]
In our setting, the positive semidefinite (PSD) NNGPK appears in an entirely novel context, where the role of the hidden weights and biases 0, = (w; b:) and the data x is reversed. Instead of marginalising out the parameters and evaluating at the data, we marginalise out the data and evaluate at the parameters. The NNGPK hi,,,,,,, admits a representation of the form
ko,1d,4, (O,, 0,) A Ex [o (w x (w) x b;)], x~ (0, I).
We do not discuss in detail how it is constructed in earlier works 31
where usually b; = b, = = but not always[49) When b, b O, closed-fon expressions for the NNGPK are available for different choices of a and [52 is important in our context (as we show in 2.2 and presents additional analytical challenges.
37 47 28
13
However, the setting of o
We will require more general notion of an NNGPK which call neural network kernel (NNK). We introduce a function t which may be thought of as a warping function applied to the input data. Such warping common in kernels and covariance functions and can be used to induce desirable
analytical and practical properties (37 29 24 $5 J.33. We also integrate with respect to more general measures y instead of the standard Gaussian CDF, We define the NNK be
/ ow. t(x) + b.)o (tw) t(x) + bz)
(0,,
(5)
After accounting for the reversal of parameters and data.
3

---

2 Closed form squared neural families
2.1 Normalising constants
Observe from 2 that by swapping the order of integration and multiplication by V. the normalising constant is quadratic in elements of V. The coefficients of the quadratic depend on (W, b).
We now characterise these coefficients of the quadratic in terms of the NNK evaluated at rows 0. 0 of ( which are totally independent of V. The proof of the following is given in Appendix [A]
Identity 1. The integral 2 admits representation of the form
= =V (V)
(6)
where ka,t, is as defined in (5) and Ke is the PSD matrix whose ijth entry is to (0., 0z).
By Identity the normalised measure ① then admits the explicit representation
Tr (v (a) P(dæ; V. 0) Tr TT(VVVKa) p(daz)= vec VTVV vec (ee)
vec(VT VVT vec Ko(
/ p(dx).
where (x) is the PSD matrix whose ijth entry is (w (xx) + b:)o(w) (x))+ +b,) We used Tr [VooTVT
the eyclic property of the trace, writing the numerator as Tr (o V Tr Voo ). We emphasise again that the role of the data t(x) and the hidden weights and biases θ in the NNK ka,t, are reversed compared with how they have appeared in previous settings. We may compute evaluations of ko,t, in closed form for special cases of (a, t, μ1) using various identities in (d), where d is the dimensionality of the domain of integration, as we soon detail in $ 2.2 Combined with the trace inner product, this leads to a total cost of computing z (V. +) of O [m n. + dn²), where and m are respectively the number of neurons in the first and second layers.
Remark 1 (Alternative parameterisations) SHEFY models depend on readout parameters V only through the direction of vec V V) and not on its norm or sign. For example, one can always find another parameterisation of readout parameters that results in the same probability distribution but has normalising constant of 1. Furthermore, noticing that V only appears as a PSD matrix M 2 V of rank at most min(m, n)), one may alternatively parameterise a SNEFY by (M, θ).
2.2 Neural network kernels
In $ 2.1 we reduced computation of the integral (2) to computation of a quadratic involving evalu- ations of the NNK (5). Several closed-form: form- are known for different settings of a, all with = Id and b = 0. The motivation behind derivation of existing known results is from the perspective of inference infinitely wide Bayesian neural networks [3] or derive certain integrals involved in computing predictors of infinitely wide neural networks trained using gradient descent [17]. Here we describe some new settings that have not been investigated previously that are useful for the new setting of SNEFY. Recall from (5). that our kernel, ka is parametrised by activation function o, warping function (sufficient statistic) t. and base measure μs. All derivations are given in Appendix B
The first kernel describes how we may express the kernels with arbitrary Gaussian base measures & in terms of the kernels with isotropic Gaussian base measures 4. This means that it suffices to consider isotropic Gaussian base measures in place of arbitrary Gaussian base measures.
Kernel 1. &e.. (0,, 0z) = ko, (TO,, T0₃), where T = (WA,b. + Wm), TO, = (w.) A, b, + wi! m) and A is a matrix facror such that covariance C = AAT T
This kernel can also be used to describe kernels corresponding with Gaussian mixture model base measures. The second kernel we describe is minor extension to the case b * 0 of a previously considered kernel [37].
Kernel 2. Reoe,Id, 4 (e., 0,) cos (b. byl exp
ons bb, exp || [w 10,,,,;1'"" 113
2
2
A similar result and derivation holds for the case of Id, which we do not reproduce here. We now mention case that shares a connection with exponential families (see $ 3.2 for a detailed description of this connection). The following and more exp cases are derived in Appendix C

---

Kernel 3. Define projea- (z) a/ | to be the projection onto the unit sphere. Then
Kexp.progge & (0,, 0,) exp(b; +b,) T(d/2)22/22114 14/2-1(11ur) + w3ll) w; + w |d/22-1
where 1p is the modified Bessel finctton of the first kind of order p. In the special case d = 3, have the closed-form kexp.pmiga @ (0,, 02) exp(b; + bi) (ell +
2 w.
We end this section with analysis of the Snake, activation function, given by
1
Snake, (2) = sin' (az) = = - cos + 2a
sin
2a
The Snake, function [55 is neural network activation function that can resemble the ReLU on an interval for special choices of a, is casy to differentiate, and as we see shortly, admits certain attractive analytical tractability. We note that similar activation function has been found using reinforcement learning to search for good activation functions [39] Table and 2, row 3], to an offset and hyperparameter a 1. The required kernel is expressed in terms of the linear kernel (Kernel 4 and the kernel corresponding with the activation function of [39], i.e. snake without the offset, Snakea(.) 2a (Kernel
5). We first describe the linear kernel.
Kernel 4. hid,aa (,, 0,) w. w + byby-
now derive the kernel corresponding with Snake, activation functions up an offset. (-) (-) t. o (0,, 0ÿ) is equal to
Kernel 5. The kernel Kynake,
4a kooo Id, (2a0; (2a0,, 2a0,) + w wi sin (2nbbb))) 2a wyll sin(2ab;)) 2a 20²1|=x11'
1
b;
cos(2ab;)e 2a 2,,1| bj
cos((ab;))
2a 10.1l kia.dd o (0. 1.1))
2a
2a
The kernel corresponding with Snake, activations is then stated terms of Kernel and Kernel 6. The kernel ksmake., (O;, 0,) is equal to
bb
cos exp( 2a² w. w1,13) + by
cos(2ab;) cxp(- 2a ww .13))
2a
2a
2a
1
+ kgnake,(-) 2n Id. (8,, 0,)
4a²
Properties of squared neural families
3.1 Fisher-Neyman factorisation and sufficient statistics
If the base measure is absolutely continuous with respect to some measure v, and dy Ω + 00 x)) de the Aadon-Nikodyn derivative, then the SNEFY admits probability density function p(. V, %,
with respect to
dj
f (t(x); V θ)
p(x x V ))
a
2
2VV +
Independent of V,
Depends on only through ((x)
The Fisher-Neyman theorem (for example, see Theorem 14 of (6) says that the existence of such factorisation equivalent the fact that is sufficient statistic for the parameters V, e.
3.2 Connections with exponential families
In this section we will use the activation (u) = (2)pu We note that we can absorb the bias terms b into the V parameters and obtain as a special case the following family of distributions
Σ Σ u. exp (H(u: + w;) TT) [e((t) u(dx),
Pddæ; V, W) = Tr(V VK))
(8)
'Note that exp(b:) is independent of and only appears product with es.
5

---

which is mixture of distributions Pe (:: (w, + wwy)) belonging to a classical exponential family given the canonical form by
Pa [0 32].
exp ( ze(w)) = / exp (w (wJ
(dla; w)
(9)
Zee (w))
It is helpful identify the following three further cases:
1. When family with base measure μ supported on X, sufficient statistic t, canonical parameter wh and normalising constant ze (w.). Every exponential family is thus SNEFY, but not conversely.
1, vi cancels in the numerator and denominator and obtain exponential
When m > and n > 1, we obtain a type of exponential family mixture model with coefficients V 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 V is at most min(m, n).
3. When m > and n > and V V is diagonal (i.c. each column in V is orthogonal), there are at most non-zero mixture coefficients, all of which are nonnegative. That is, we obtain a standard exponential family mixture model.
The kernel matrix Ke in the normalising constant of (8) is tractable whenever the normalising constant of the corresponding exponential family is itself tractable.
Proposition 1. Denole by ze (w) the normalising constant of the exponential family in 9 Then
kexp(-/2).t.pp (15;, wz) = -
(10)
The above kernel is well defined for any collection of w, which belong the canonical parameter space of Pe. since the canonical parameter space is always convex [50]. This gives us a large number of tractable instances of SNEFY which correspond allowing negative weights a selection of examples is given in Appendix
exponential family mixture models
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 normalisiin constant to the mean and the covariance of the sufficient statistic.
Proposition 2. Let a (u) exp(u/2) and define the log nogrooimaliinng constant as = logz(V O).
Then
2V
E (x))
and
=E (xx)
(x)
Ow,
Tw.
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 x = (x1,x2) be jointly SNEFY Loo. with parameters V and O = (W1, Wz]. b). Assume that u(dx) = fr (1(x)))) (dxz) and t(x) = (t1 tz(x2)) Then the conditional distribution of x1 given X2 D2 SNEFYZ, with parameters V and 112 & (W1, W2te(x2)+ b).
The proof, which we detail in Appendix A follows directly by folding the dependence the conditioning variable 22 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 o. 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 ty itself to be jointly trained deep neural network. Crucially, t₂ may be completely unconstrained, We use this observation in the experiments
'eere we are concerned with the setting where every term in the mixture model belongs to the same family, ie. 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 a need not belong the same family. In the special case o (x)). SNEFY is in fact also closed under marginalisation, which we prove in Appendix D Even in the general o case, marginal distributions are tractable and admit closed forms whenever the joint SNEFY model and the conditional SNEFY are tractable.
(x1;X) be jointly SNEFYs, XXx.tx, with parameters V and θ = ( 1. W b)- Assume that p(dæ)) = (dx; 1/22 (dxz) and t(x:
Theorem 2. Let x
Then
the marginal distribution of K1P5
Tr VC
Pi(dxi)=
Pi (dxi),
<V,
where C(x1)1) ko, ta (w2,, (, + b.), (22j. w t (11 + b,))
The proof is given in Appendix Due to this tractability of the marginal distributions, it is straight- forward to include the likelihood corresponding to incomplete observations (i.e. samples where we missing some of the components of x) into the density estimation task.
3.5 Connections with kernel-based methods for nonnegative functions
SNEFY may viewed as a neural network variant the non-parametric kernel models for non-negative functions (25), which are constructed as follows. Let (possibly infinite dimensional) Hilbert space H. Let ((H) be
3 X H be feature mapping
the set of all positive semidefinite (PSD) bounded linear operators A H + H. Then
h (2) (bb(),, Av))))
(11)
gives elegant model for nonnegative functions parametrised by A E S(IHI, and their applica- tion density modelling with respect to base measure has also been explored (25] 42). Note that by assuming boundedness of and vb, the normalizing constant [25] Proposition 4] is given by Tr (A fx vxx) @ (a))u(dx)) which is anal- work where the normalising constant given by z(V,0) = Tr(V V)). This can be seen by replacing A with VV and replacing fx (xx) (xx)) by Ke [(ttt(( +
Figure 1: (Left) An instance of (untrained)
SNEFY Sdd density with n 100, fij (o, //n,, Wij ~N 4) and = 0. Shown are 50 exact samples found using rejection sampling. Numerical quadrature for this and every example sup- ported on Rd in returns value of 1.00 for the integral over X. (Right) A trained SNEFYs², Id.exp.de density with n m 30. Shown is the training and testing dataset 44 also used by [99 for point processes.
ogous
b)o b) p(dx).
Despite feature maps being infinit-ddimensionall
model [T] often reduces ter theorem when (II) is fitted to data using a regularised objective, while [42] more directly assume that the linear operator A is inside the span of the features evaluated the available data {xz N The resulting model resembles SNEFY where n equals to the number of datapoints, i.e.
equivalent representation in dinteeddmessions. 25 utilise represen-
hM (x)
[K( æ1),
n(a, xN)] M (....,,))......(((.....
(12)
for PSD matrix M NxN and x(x,, (vb( (x)))
However, there are fundamental differences between (2) and SNEFY which we list below. The models can be seen complementary and they inherit advantages and disdvantages of kernel methods and neural networks common in other settings, respectively.
Traatability. Whereas we identify many tractable instances SNEFY, the normalising constant of 12 requires computing f k( x)x(,, <d)x(dx)) this is not generally tractable apart from some
7

---

limited combinations of and (e.g. for a Gaussian kernel k and Gaussian 14), Note that this kernel is evaluated at the datapoints, whereas SNEFY evaluates the kernel the learned parameters un; 22 focuses on the specific case where is Gaussian kernel, studying properties of the resulting density class which is 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 &b and consequently finite dimensional representation via K. 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 k),, at the expense of a more involved optimiisation over the space of PSD matrices. In contrast, we learn W (analogous to learning 2/) and V jointly using netral-network style gradient optimisers. SNEFY is fully compatible with end-to-end and jointly optimised neural network frameworks, property we leverage heavily in our experiments in $
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 vb.
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 [3]. We believe a marriage of ideas from SNEFY and probabilistic circuits will lead to future developments in tractable and expressive probability models.
4 Experiments
Implementation All experiments are conducted 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 normalising flows. SNEFYs are built as a BaseDistribution which are base distributions inside NormalizingFion with greater than equal to zero layers. We train all models via maximum likelihood estimation (MLF) i.e. minimising forward KI. divergence.
top of normflows [45). PyTorch package for
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 O corresponds with the base distribution only. We use cos, t- = Id, X IRd and Gaussian mixture model base density. We set m 1 and n = 50. Full architectures and further experimente details are described in Appendix E.I
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 normalising flow of appropriate depth. SNEFY is able to achieve good performance with depth 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, shown in Figure Over 50 runs, the unconstrained V. diagonal V, and n = m = cases respectively obtain test negative log likelihoods of 1.38 + 9.64 x 10-3, 1.46 + 0.016 and 2.34 + 0.26 each in 111.18 + 2.58, 109.12 12 + 0.80 and 61.88 + 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 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 quasars) without measuring their full spectra, but by
8

---

BHEPY
RessspledO
BHEPY 16
Resuxpled
Gauza
Moons
11:50
-1.07 0.02
-3.20 0.00
-1.51
-1.5 0.00
11250
-188
T.Y DX0
243
06817
19281
80857
1901
1200.84 84 100.77
541.11 12,30
606.35
20.20
629.11 14.57
2844.00 254.02
1867 167.75
2226.37 174.80
2296.70
211
Clielss
1,32 0.105
2,42 0.22
3.55 U.00
2.07 #: 0.07
1192 1004
1.56 B00
2:22 0.30
13 1000
241
6GM17
19221
85557
19348
19029
640,90 4: 150.84
108,45 4: 20,.0
52.63 4 3.36
73.52 4 5.447
2272:15 15 251.75
1683. 30
138.90
1063777 + 100.02
1673.54 = 106.84
Kings
2.41 = 0..05
2202 0.00
5.20 0.00
-2.07 + 000
-2.34 0.10
-2.36 0.14
2:51 =0.22
-2.1 0.01
241
60817
50
10265
65857
1901
10000
507.73 4: 107.80
400.77 +: 19.57
416:23 1 18.11
453.82 1: 16.06
2051 200.14
1774.00
182.25
2045.66 =: 170,80
2070.60 = 177 06
Table 2: The first quantity is the average + sample standard deviation over 20 runs of test loglikelihood. The second quantity is the parameter count. The third quantity the average I 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 or 16 NVP layers. More tables showing intermediate layer results are given in Appendix E.1
Method
Average Test NLL
Compute time (8)
Parameter count
SNEFY n 32
2.99 + 0.024
495.720 + 22.561
179748
SNEFY n = 16
-2.17 + 0.034
404.291 + 57.605
46356
SNEFY 8
-2.108 + 0.089
390.870 + 16.028
12300
CNF L
-2.15 i 0.018
202.1290 + 10.380
1413
CNF L =
-..66 + 0.024
155.090 + 15.809
1155
CNF L=
-2.77 + 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 = indicating standard deviation. SNEFY shows statistically significant average increase in performance over CNF (using two-sampl 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 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 E IR conditional on features 22 E TR. using a public dataset (2) The features D2 are the magnitude a broadband filter (red) and set of four pairwise distances between broadband colour filters (ultraviolet-green, green-red, red-infrared infrared -infraredd). We benchmark our own SNEFY against a conditional kernel density estimator (CKDE) and 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 a Snake,, Id, X R and a Gaussian base density. We use deep conditional feature extractor with layer widths [2 n, 2 n, 22n, 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 not well-suited to nontrivial tabular data with 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 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 &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 74309 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 SNEFY model that utilises partial observations by maximising the marginal likelihood according to Theorem 2 We leave the original 74557 x 6 testing dataset untouched, and measure the test NLL after training. We plot the NLL as a function of 1 1 q. as shown in Figure We also compare the performance of the SNEFY that uses partial observations with the performance SNEFY that simply throws away incomplete observations, 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
well as the same normalising flow baselines that throw away

---

Density estimation under par- tial observations. The right plot is a zoomed in version of the left plot. NF- 1, NF-2 and NF-4 are respectívely nor- malising flows of depth 1, 2 and 4. Normalising flow models and SNEFY without marginalisation discard incom- plete observations, whereas SNEFY use Theorem 2 include partial observa- tions via maximum marginal likelihood. Marginal likelihoods allow for an im- proved NLL.
Figure 2:
S0AF7 veb-r: beriee 20127 Bobiiie
DOLPY me SHTE wiinten INGLY murghallekke
MR
The Poa
Proiiaaly of prinet lates
Prodict prrvent
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.
5 Discussion and conclusion
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 is sometimes possible to obtain exact SNEFY samples using rejection sampling, sampling is more computationally expensive than in other models such normalising flows. Future work will focus on sampling, as has been done with related models (66. where the
special case of Gaussian and hyperrectangular support X is considered. In (26]. approximate samples with (r log, X + rd log, 2) evaluations of the normalising constant are obtained. Here 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, may any number of layers and any architecture for the conditioning network t2. Finally, SNEFY inherits all the usual limitations and advantages over mirroring kernel-based approaches, as discussed in 3.5
Future work We see a number of further promising future research directions. First, as we detail in Appendix F. choosing at)) = exp(i-) and identity sufficient statistics results in kernel Fe which is the Fourier transform of a nonnegative measure. By Bochner's 's theorem, the kernel is guaranteed to be (real or complex-valuede shift-invariant. The kernel matrix is Hermitian PSD (so that the normalising constant is positive and nonnegative), and we may also allow (but do not require) the readout parameters V to be complex. We note that the same result would be obtained if one used a mixture of real-valued cos and sin activations with shared parameters (see Remark 3) Second, an alternative deep model to our deep conditional feature extractor might be to use a SNEFY model as a base measure μ for another SNEFY model: this might be repeated L times. This leads to ((z22.) integration terms for the normalising constant. The individual terms are tractable in certain cases, for example when a is exponential or trigonometric. Third, when modelling discrete distributions with trigonometric activations, the NNK can be expressed in terms of convergent Fourier series (see Appendix G Finally, our integration technique can be applied to other settings. For example, we may build a Poisson point process intensity function using squared neural network and compute the intensity function in closed-form, offering model that scales quadratically in the number of neurons (n²) instead of comparable models which scale cubically in the number of datapoints
(NN3) 9 511.
10

---

Acknowledgments and Disclosure of Funding
Russell and Cheng Soon would like acknowledge the support of the Machine Learning and Artificial Intelligence Future Science Platform, CSIRO. The authors would like to thank Jia Liu for early discussions about the idea.
References
[1] Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(19): 133,, 2017.
[2] Beck, C-A Lin, EEO Ishida, F Gieseke, RS de Souza, MV Costa-Duarte, MW Hattab, and A Krone-Martins. On the realistic validation of photometric redshifts. Monthly Notices the Royal Astronomical Sociery, 468(4))4333-4333) 2017.
[3] Edgar A Bernal. Training deep normalizing flow models in highly incomplete data scenarios with prior regularization, arXiv preprint arXiv:2104.0148) 2021.
[4] Youngmin Cho and Lawrence K. Saul. Kernel methods for deep learning. In Advances in Neural Information Processing Systems, pages 342-350, 2009.
[5] Michael Collins, Sanjoy Dasgupta, and Robert E Schapire. A generalization of principal components analysis to the exponential family. Advances in neural information processing systems, 14, 2001.
[6] Marc Peter Deisenroth, A Aldo Faisal, and Cheng Soon Ong. Mathematics for machine learning. Cambridge University Press, 2020.
[7] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations, 2017.
[8] Bradley Efron and Trevor Hastie. Computer Age Statistical Inference: Algorithms, Evidence, and Data Science, volume 6. Cambridge University Press, 2021.
[9] Seth Flaxman, Yee Whye Teh, and Dino Sejdinovic. Poisson intensity estimation with repro- ducing kernels. In Artificial Intelligence and Statistics, pages 270-279, PMLR, 2017.
[10] Thorsten Glüsenkamp. Unifying supervised learning and vaes-automating statistical inference (astro-) particle physics with amortized conditional normalizing flows. arXiv e-prints, 2020.
[11] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, volume 27, 2014.
[12] Insu Han, Amir Zandieh, Jachoon Lee, Roman Novak, Lechao Xiao, and Amin Karbasi. Fast neural kernel embeddings for general activations. In Advances in Neural Information P Processing Systems, volume 35, pages 35657-35671, 2022.
[13] Insu Han, Amir Zandich, Jachoon Lee, Roman Novak, Lechao Xiao, and Amin Karbasi. Fast neural kernel embeddings for general activations. Advances in neural information processing systems, 2022.
[14] Dan Hendryeks and Kevin Gimpel. Gaussian error linear units. arXiv preprint -xXix:6606.08411 2016.
[15] Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 1778)1777, 2002.
[16] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[17] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances in neural information processing systems, pages 8571-8580 2018
11

---

[18] Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv preprint 1rXiv:1312.6114, 2013.
[19] Thomas Joseph Lawrence. Point pattern analysis on a sphere. Master's thesis, The University of Western Australia, 2018.
[20] Nicolas Le Roux and Yoshua Bengio. Continuous neural networks. In Artificial Intelligence and Statistics, pages 404-414 2007.
[21] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy- based learning. Predicting structured data, 1(0), 2006.
[22] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. In International Conference on Learning Representations, 2018.
[23] Lorenzo Loconte, Stefan Mengel, Nicolas Gillis, and Antonio Vergari. Negative mixture models via squaring: Representation and learning. The 6th Workshop on Tractable Probabilistic Modeling, 2023.
[24] David JC MacKay. Introduction to Gaussian processes, 1998.
[25] Ulysse Marteau-Ferey Francis Bach, and Alessandro Rudi. Non-parametric models for non- negative functions. Advances in neural information processing systems, 33:12816-122 122226 2020.
[26] Ulysse Marteau-Ferey, Francis Bach, and Alessandro Rudi. Sampling from arbitrary functions via psd models. In International Conference on Artificial Intelligence and Statistics, pages 2823 2. PMLR, 2022.
[27] McCullagh and J.A. Nelder. Generalized Linear Models, Second Edition. Monographs on Statistics and Applied Probability Series. Chapman & Hall, 1989.
[28] Lassi Meronen, Christabella Irwanto, and Amno Solin. Stationary activations for uncertainty calibration in deep learning. Advances in Neural Information Processing Systems, 33:2338- 2350, 2020.
[29] Lassi Meronen, Martin Trapp, and Arno Solin. Periodic activation functions induce stationarity. Advances in Neural Information Processing Systems, 34:1673-1685, 2021.
[30] Charlie Nash and Conor Durkan. Autoregressive energy machines. In International Conference on Machine Learning, pages 1735-1744, 2019.
[31] Radford Neal. Bayesian learning for neural networks. PhD thesis, University of Toronto, 1995.
[32] Frank Nielsen and Vincent Garcia. Statistical exponential families: A digest with flash cards. arXiv preprint arXiv:09/ 1.4863, 2009.
[33] Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A Abolatia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are Gaussian processes. In The International Conference on Learning Representations, 2019.
[34] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1):2617-2680, 2021.
[35] Georgios Papamakarios. Neural density estimation and ikelihood-free inference. PhD thesis, University of Edinburgh, 2019.
[36] Jooyoung Park and Irwin W Sandherg. Universal approximation using citiial-aai--funtttii networks. Neural computation, 3(2):246-257, 1991.
12

---

[37] Tim Pearce, Russell Tsuchida, Mohamed Zaki, Alexandra Brintrup, and Andy Neely. Expressive priors Bayesian neural networks; Kernel combinations and periodic functions. In Uncertainty in Artificial Intelligence, 2019.
[38] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
[39] Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2018.
[40] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538, 2015.
[41] Trevor W Richardson, Wencheng Wu, Lei Lin, Beilei Xu, and Edgar Bernal. Mcflow: Monte carlo flow models for data imputation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 14205-14214, 2020.
[42] Alessandro Rudi and Carlo Ciliberto. PSD representations for effective probability models. In Advances in Neural Information Processing Systems, volume 34, pages 19411-19422, 2021.
[43] Aleksanteri Sladek. Positive Semi-Definie Probabilistic Circuits. Master's thesis, Aalto
University. School of Science, 2023.
[44] W. Steinicke. Revised new general catalogue and index catalogue (revised nge/ic). http:
//www. wwwka-liiaa----- luff.. www.dd///eeeiiiiee http:/eteiiiniececeeeeeeeeeeeeee e. h.m Accessed 2nd May 2015.
[45] Vincent Stimper, David Liu, Andrew Campbell, Vincent Berenz, Lukas Ryll, Bernhard Schölkopf, and José Miguel Hernández-LLobato. normflows: A PyTorch package for nor- malizing flows. arXiv preprint arXiv:2302. /2014, 2023.
[46] Vincent Stimper, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Resampling Base Distributions of Normalizing Flows. In Proceedings of The 25th International Conference Artificial Intelligence and Statistics, volume 151, pages 4915-4936 2022.
[47] Russell Tsuchida, Tim Pearce, Chris van der Heide, Fred Roosta, and Marcus Gallagher. Avoiding kernel fixed points: Computing with ELU and GELU infinite networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9967-9977, 2021.
[48] Russell Tsuchida, Fred Roosta, and Marcus Gallagher. Invariance of weight distributions in rectified MLPs. In International Conference on Machine Learning, pages 5002-501 1. 2018.
[49] Russell Tsuchida, Fred Roosta, and Marcus Gallagher. Richer priors for infinitely wide multi- 2019.
layer perceptrons. arXiv preprint
[50] Martin J Wainwright, Michael Jordan, al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1-2):1-305, 2008
[51] Christian J Walder and Adrian N Bishop. Fast bayesian intensity estimation for the permanental process. In International Conference Machine Learning, pages 3579-3588. PMLR, 2017.
[52] Christopher KI Williams. Computing with infinite networks. In Advances neural information processing systems, pages 295-301, 1997.
[53] Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning. MIT press Cambridge, MA, 2006.
[54] Christina Winkler, Daniel Worrall, Emiel Hoogeboom, and Max Welling, Learning likelihoods with conditional normalizing flows, 2019.
[55] Liu Ziyin, Tilman Hartwig, and Masahito Ueda. Neural networks fail learn periodic functions and how to fix it. In Advances in Neural Information Processing Systems, volume 33, pages 1593-1594, 2020.
13

---

A Proofs
Identity The integral 2 admits representation of the form
= (V
(6)
where ka is defined in 5 and Ke is the PSD matrix whose ijth entry is to (0s, 0j).
Proof. Let (x) be the PSD matrix whose ijth entry is a (w. f(xr) + b:) t(x:) bj) The
squared norm of the neural network evaluation is given by Tr (V V (2:)) since
Vo((((<)) - b)|l=
Big Ukja σ
(w) t(z) + وز. )o (w) t(x) bjs)
(w,, ( t(x) ري ) (w), ((x) + bj.)
J:
(V V, Ko(x))F
= Tr (V V V
where v..i, denotes the jith column of V and ( denotes the Frobenius inner product. Therefore using the definition 2 directly and the linearity of the Frobenius inner product, the normalising constant is
(V. + + + (()
Tr vko(=)) u(dxx)
/ Ko(z)u(dx)) = =Tr
-Tr(VTV
(13)
Proposition 2. Let (w) = exp(u/2) and define the log log-normalising constant as = logzoV, O).
n
n
n
82
Then Σ
E [t(((x)
and dw.w
E. (x)) - [((
Ow;
Proof. The result follows by noticing that the logarithmic derivative property holds, Σ 0x and that by writing
24
8w;
ow:
EV, /
exp ww; (ax) exp [w M(xe) a (dx))
1
25
we obtain
0z
||l,,el? f t(x) exp (w) t(x)) p (dx)
Ow;
t
Σ v.d / t(xz exx) exp
w. t(zz)
exp 2 w t(z)) a (dx)
j71
v. Σ n. f1(m)) t(xz) exp exp
_ u. t(xr)
exp [ t(xr) 1 (dx).
2
Now summing across gives
Σ
a4 dw.
t(x)) as required.
14

---

E it (x)). Then the To obtain the second result, we apply the product rule to find
Σ
0
1
Oz Swe:
1
Σ
Oz
ow,w,
ww, w
Iro
and note that
T 2 J [(=))((=) exp (gw, t(x)) exp [2 t(x)) (dx), p
if j.
8 z Tw,w
|o 12 (())((()) exp (w) t(x)) l (dar
t Enfs Vig t(x))(()) exp (5 t(x)) exp w t(x:)) p (dx),
if =-
Thus,
82g z
82 2
2 n dwmm
jti
1
1 [o t(x)
1
Σ i,
vs Vi t(x)t(()) exp
exp
2 / t(xz) a (dx)
2
/ t(x))(x)) (((d)),
as required.
Remark Given the above relationship between the log-normalising constant and the expectation of the sufficient statistic, we can also ask whether the maximum likelihood estimation of the mean parameters [ (x) proceeds in the same way as the exponential family case. The answer is positive but with two caveats. First, the log- tig-likeldoood need not be concave in W. and may have many local optima or stationary points. Second, unlike the exponential family distribution, the SNEFY distribution is not determined by its mean parameters, the MLE estimation of the mean parameters may not constitute a meaningful task in SHEEF modelling (unless we are in the case where precisely the expectation of t(x)) under the SNEFY model is of interest).
Corollary 1. Given dataset {xe f== and SNEFY model with (u) exp(su), assume that all rows of the maximum likelihood estimator of W are in the interior of the natural parameter space of the corresponding exponential family. Denote the mean parameter maximum likelihood estimate of m m E t(xe).
Proof. Since MLE is achieved at a stationary point of the log- bog-iikeihood, the proof follows by writing the log likelihood
N
N
log p(ze; V. O) const Σ
log (t(x)) V, V, - N
and concluding that at the MLE ww Jm for W, we must have
0 log (t(x))) V. O
N ov
i i= 1,...,n.
(14)
Ow
Ow
But
d Oog (I ) t(xx):
Oun
so the result follows by summing over and dividing by N. Note that maximum likelihood estimates are invariant to transformations. even if the transformation is not bijective. So if [w) y an MIE, we may construct mapping from (w) ) to corresponding MIE i for the mean parameter.
15

---

(x1.X2) be jointly SNEFYx, tt., with parameters V and O (W1, W2], b). Assume that p(dx) = pr (dxx 1 (dxz) and t(x) (t1 tz(2)) Then the conditional distribution of X1 given X2 = #2 ,N with parameters and 11 _ (W1, W2tz(02) + b).
Theorem 1. Let x
Proof.
The joint distribution of satisfies
P(dæ; V. ①) Vo(Witi(a)) + Watz(2z) +b)][(((a))
Therefore, the distribution of x₁ conditionally on X2 22, which is obtained by dividing the joint distribution by the marginal distribution of (which is independent of æ1), satisfies
P1 dx
22 V. (W1, W2tz(22) - b))
(Vo(Wit((z)) W2t2(222 + +)) μa (dx1).
That is, the term W2t2(22) + b is viewed as a constant bias term when the expression on the right hand side is an unnormalised measure with respect to the variable 21.
Theorem 2. Let x
(x1,X2) be jointly SNEFY tt2,,, with parameters V and + ([w, W2), b). Assume that p(dx) = (dxx) (dx2) and t(x) = (ti(x1), (2(2))). Then
the marginal distribution of X1
Tr (V (VV((())))
(i(dx))
M1 (dxx);
x(
where C(x1)) ko (w,,, ( + bi), (w2j. wiit (xx1 bj)).
Proof. The marginal distribution of the random variable x is obtained by marginalising out the joint distribution with respect 222,
(1 Vo((iit(())) Wat2(x12 + bb) (2(dx22 {)) (d(x1) z(V, 0)
fr (dx1).
Avg
The integral z2 takes similar form to z(V, VVV,),
= Tr (V Vče(())) where
(i,(ai)= = { Jxn ((w]) + w))(((2)) + bi)o + un7 + bi) (z)
= Ka.tai (w1z (Tth( + bi), (wzj ( bb))
B Derivation of neural network kernels
Kernel ko (0,, 0,) 1a, (To,, T0,), where T (WA,b b + Wm), TO, (w.. A, b, + wi m) and A a matrix factor such that covariance AAT
expectation with respect to Gaussian random variable x with mean m and covariance matrix C. It holds that x Az + m, where z is a zero-mean independent standard Gaussian random vector, so the kernel may be expressed in terms of an expectation over z instead. More concretely.
Proof. The NNK may be expressed
kg, (0, f,)
Ex [a _u[x x b, ( x b,)]
= [ww Az + w. m + (:) Az + w3 m + b,)]
= -oo.d (TO, TO,).
16

---

Kernel 2. I 0,)
cos 1b,
exp
I| we;ll")
cos tby || wy t t-m, |I
2
exp
Proof. First observe that the expected value of the cosine of Gaussian random variable can be evaluated by equating the real and imaginary components the characteristic function Gaussian random variable and the expected value of Euler's form. That is, if z is Gaussian with mean m and variance 2?
Eesh EE[cos(z)] + [[[[iin(z]] e 10²
= (cos μe + sis 20?
E[cos(z)) = cos(m)) 1.²
With this identity at hand, we proceed by direct evaluation of
Kons,Id, I (0;, 0,) Ex cos(ww + by) cos w b5}], ~ 10
2 Ex cos ((w; w)) x + (b. by)) + cos ((w; + w,) x + (b; + bb))]]
2 cos bb; - bylexp 1 -|l: - w;]l?) -ww;-
+ cos |b, + رyو exp マ w11, + w;)l?)
Kernel 4. Naa,d,,(0,,00 = w. w + b,by.
Proof. This is immediate from the expected value of a product of two correlated Gaussians, w x and w + bj.
Kernel 5. The kernel kssnake(()) Id. (0;, O5, 0j) is equal to
(2n,l,, (2a0), + mi un (in((ab;)) ) -20 w.ll + sin((ab;)) ) 20/1|
4a² 1e,
b;
cos (os())a 20 bj
cos(2ab;)e 20 w|;1? kia.a, (0 0
2a
2a
Proof. Choosing o = Snakez())
2a in and expanding the resulting quadratic,
KSnake, (+) + (0,, 0,)
4a F.
[coos/ cos (2a(w, x + b;)) cos (2a(w) x + b;))]
Rx (10 + bı) cos (2a(1r, x b;))]
2a
Kemel 2
44
Ex cos (2a(w) w x + b.)) (w) x+ b,)] + Ex [w x + bi) x + by)
(15)
2a
Keme
B
now evaluate and B. The terms A and B obey a symmetry, so suffices to evaluate term A. Term A can be evaluated using Stein's 's lemma,
1
E (z) + bi) cos(22 2ab;)] 22) N T (0((())
w.Tw: w; 2aw w
2aw w. 4a uo?w] w.
A
2a
- [[2 cos(z2 2ab,)] + "5[[os(22 + +abb)]]
1
_e
by
20
2a
by
= =- wy E sin(22 + 2ab,)] +
B cos(22 + 2ab;)]
2a
= =- wy sin(2ab,) exp(-2at 2d2| ( 12) + by 2a
cos(2ab,) exp(-2o0 (2 11))
17

---

Assembling all the known individual terms in 15
Ksuake, (-) - zz La 0;)
Reos,Id (2a0;; 2a0,) + w, wy (in())
22,ala11, (1,,)1 + sin(2ab,) exp( exp(-2a' 22 2 (;;1)))
4a
cos((2b)))))mo exp( 1app--2aa1la 22 - 2a
by 2a
cos()ab,)
12) + ( .0.)).
Kernel 6. The kernel kSnake,, Id, (0,, 0,) is equal
1
1 cxp(-2ppp (2 + bj 2a
1
b
cos((ab,)
cos(2ab,) exp(-220 (1-2aa11122 11,,,111 (13)
2a
2a
+ (nake,(()) Id, 05) + 4a²
Proof. Choose Snake, and note that Snake, (Snake,))
20 2/2 The Kernel
corresponds with the case (Snake 2/a) so left with three additional terms. These terms may be evaluated directly,
ksnake, ta (0,,
1
Ksnakea()) (.)- Ia (0,, 0,) +
4a2
1
bi
cos(2ab:) exp(-
1ap(-2oh -2a 12,211 + bb 2u
cos(2abj) exp(-2a 2200 11))
2a
2a
C
Examples which generalise standard exponential family models
In this section, we will study examples of the SNEFY model which use activation function auu) exp(u/2) (or equivalently, up scaling, (n) exp(u)) and such correspond notion of exponential family mixture models allowing negative weights, as discussed in Section [3.2 These examples have tractable kernels whenever the corresponding exponential family has tractable normalising constant and we can write the kernels directly using Proposition
SNEFY Von Mises-Fisher mixtures. The VMF distribution is helpful way of defining the notion of Gaussian distribution to the sphere. The following kernel may be used to define VMF distribution. Alternatively, it may be viewed way of constructing a distribution supported on IRd with sufficient statistics which are projected onto the sphere.
Kernel 3. Define projga (2:) x ||=l|| be the projection onto the unit sphere. Then
liexp hexp--roiga-s (0,00) = exp(b; + -b,) 1///2)22/22-1//2-110 + w. I|w: w d//2-1
where L, is the modified Bessel function of the first kind of order p. In the special case d = we (cll w le
have the closed-form form praga ( (0,,0,) =exp(b, + b,))
w
18

---

Proof. If x ~ (0. I), then x/|| is uniformly distributed on the sphere. From the normalizing constant of the von Disee-Fisher distribution, from Proposition
[ then follows that
kexppprrist-1 1,0 (0;, 02) Ex exp (w, x/| x|| + by + w. x + b;)] x~ ~~,,
exp(b: + bU) /
exp ((w; wwy) T xx) dæ r(d/2)
2T4/2
94-1
exp(b; + (d/22
exp (1|1s + wy a x)) dx,
where a is unit vector
2/
exp(b; by (d/22) (2 (2T)d/2Ta/22
ur, w
2/d
|00; + w/ |||11/2-1
exp(b; b,) r(d/2)24/2-1
1//2 wi + wwyl)
[w; + wy 1/2-1-1
where I is the modified Bessel function of the first kind of order p. In the special case of p = 1/2, we have (=) 2/ sinh(z) (exp(z) exp(-2)) 25z
This implies that when d = 3,
since l((32) = v/
krxxp.prigg (0;,0)) = =xxp(b; + bb,) (ellae I| 100, ww,
2 ww: + w
Note that kexp. propp.roogd-1.33 (0,, (O, where v is the uniform measure on the sphere gd-1, because if x is Gaussian then x/||x|| is uniform on the sphere. This allows one to construct SNEFYexp.la distributions, which are certain "mixtures" of VMF distributions with weights V V.
SNEFY Gaussian mixtures, fixed variance. We may similarly construct kernels corresponding to "mixtures" of Gaussian distributions. The case here corresponds to a case of known fixed variance parameter.
Kernel 7. keep, Oj) = exp(b,, + by) exp [bll + -,, 11=)
Proof. This is consequence of the moment generating function of the multivariate Gaussian distribution. More concretely, by Proposition !
Iexp. (0., 0,) Ex exp (w: x x b, + w. x5]x+b0 by)], NO, exp(b, + by [xx exp ((w; + +w,) +x)]
exp(b; by) exp ( |ll, + w/ w,113)
SNEFY Poisson mixtures. Most of our examples deal with continuous distributions but in fact SNEFY can readily be used for discrete distribution modelling. This is particularly helpful when the support naively be challenging
in the discrete setting. Let X = f0, 2. t(x) x, and the base measure ((dr) 1/x! v (dæ), where is the counting measure. A SNEFY model for a probability mass function which is a mixture of Poisson distributions allowing negative weights is given by
large infinite, for which computing normalising constants
1
p(z; V.w)
v u.s exp (ww + wy)x)
TV V K
ΣV DI (AA))) x x= 0,
Tr Tr
following the usual mean parametrisation >; = " so the individual mixture components have rates which are geometric means of Ay) pairs.
19

---

Kernel 8. Choose the base measure ((dx)) = (a!) v(dx), where v is the counting measure. We have
(x (0, 0,) = exp (b + by) exp (exp(w; + wy)).
Proof.
This again direct from Proposition
In detail.
(ox),Id,(z)) (0, 0g)
Σ
exp (wnz wyx + b bz)
x!
z=0
exp
(bg + by
Σ
exp
Wgy + wyr)
==0
The second factor involving the sum is the partition function of the Poisson distribution in canonical form, which is exp (cxp(w; + wz)))
The usual mean parameterisation of the Poisson distribution is through rate parameter x; exp(w:) which would lead to the kernel representation
(exp. (,, 0,) = eep (b; + by) exp (ArAz).
SNEFY Gaussian mixtures, unknown variance (Squared radial basis function network). We now discuss an intriguing connection between the Gaussian distribution and squared RBF networks. This connection is made possible through our machinery of SNEFY distributions. Let X R4 ((x) (. La; "i....;... (i.e. D 2d) and suppose u(dx) dæ is Lebesgue measure. Choose of exp(-/2) and consider the r-th output of our network f TR. D Rm
ft(((); V. 0)
Σ 2 exp
Σ Wis +
Σ waez? 8=1
2
where we denoted th = Wi,t+é. In this case, we require that Wat (square) integrable. Reparametrising of
< for the model to be war and absorbing the factor
2 and μii -2è
2004
exp
into readout parameters V, gives Tote
d
(+(x)) V,
Σ 1=1
Uy exp
E
40
Thus, we have recovered classical radial basis function (RBF) network (52) These models are well known to have universal approximation properties (36). For the most commonly used form of RBF network, we can restrict the parameters oz o to be the same across the dimensions, giving
(_
m4:l13) / 4
fr (t(z) =
Σ Uri exp
Ø E R.
with location parameters μa; Rd and the scale parameters of > 0. Note the unusual factor of 4 in front of of this ensures that our model fact reduces to the usual parametrisation of multivariate normal densities, since we will be modelling the density using the squared norm of f.
Since μ is the Lebesgue measure, SNEFY gives us density model with respect the Lebesgue measure as
a
n2
| mill 40?
|| p 11;11
pæ: V. VK))
0 U exp
exp
40²
1=
n 1, we recover simply a multivariate normal density A (o I). The above model essentially the same the one in 42), despite being derived in very different way.
20

---

Kernel 9. Let (x) = (ai...., Td, ai..... 222) so that D = 2d be the sufficient statistic. Partition
W [g.dd W| and suppose W < 0 element-wise. Choose p(dx) dæ to be the Lebesgue
measure. Then
( (0. = 2//2 exp(b; + +b,) II
exp (-
[wa + wa1)² 4(w² + wدn)) V /((wi + úرنt)
1=1
Proof.
As with the kernels above, this follows from Proposition since
k (exp. de (0: = / exp (w., t + w; x²
+1 + +w w,11. + 2 x² + b)) dæ
= (2T) explb: + b,)) II
d
(ww + wz))²
1
exp
4( + w
V-2 2( + w
While Proposition 1 gives us an expression for the kernel matrix Ke in terms of natural parameters, we can also express directly in terms parameters p4z oz In particular,
/
14l1 x
| - M4j
KKoli = exp((h: + by)
g=
exp
4o²
exp
40z
dx
d/2
AT
[4 11
= exp(b + b,)
of
+
exp
(o og)
We briefly state two more cases without an extended discussion.
SNEFY Gamma mixtures.
Kernel 10. Let X = (0, oo), t(x) = (log -z)) and = exp. Partition W Wi.ddd W and suppose Wi.i:d and W element-wise. Choose u(dx) = da to be the Lebesgue measure. Then
di 0;) = exp(b; + bb) P(wa + Wii +1)
(Wiz + Wy2) +
SNEFY Dirichlet mixtures.
Kernel 11. Let = D-1 a (D si-sipl of probability distributions, i.e. x E [0, 1]" Eb I; 1. Let exp. Let t(x:) = (log .1...., log zp) Choose p(dx)) = dx to be the Lebesgue measure. Suppose W elementwise. Then
1 I(wid + Wid + 1) kex (O,, = exp(b; - b1) r D Edm Wid + Wid)
D Marginalisation in the case a = exp(./2)
Let M be positive semi-definite n matrix. We will make of the following SNEFY parametri- sation
Pdd:; M, θ z(M, +)
o (Wt(z) + B) Mo (Wt(x) + ()) (dx).
(16)
Since we can always write M V for m x matrix V, m S n, we have
a (Wt(x) + +b) Mo (Wt(x) +)) = (W
Σ
Σ Uijo (w. t(xx) + +b,)
21

---

which is, as in the parametrisation given in the main text, simply the squared Euclidean norm of multi-output neural network, with V corresponding to the weights of the second layer. we denote by vi the j-th column of V, the normalizing constant is given by
u t
(0,
mikho e (0, = Tr(MKe)
where as before
No.t,u(0;,00) f ( (wx) t(x) ( t(x) + u,)x).
Now if we let exp(-/2) obtain family which is also closed under marginalisation (in addition to conditioning). The Proposition below generalises Proposition [42]. which considers the special case of the Gaussian PSD mixtures.
Proposition 3. Let Xx) be jointly SNEFY x with parameters V and + (W1, W2) b). Assume that p(dx) = (ddx)))2((zz) and t(x) (21)),)2((xz) Then the marginal distribution of x1 is SNEFYx, (xp((/2)xp(()))))) with parameters V and 6 = (W1, b), for some matrix V mmxn
Proof. Proposition D gives us the normalising constant for the parametrisation where biases absorbed into V. If we explicitly keep the biases in the parametrisation, we have
12 (0, = =exp (~000 b)) Ze (2 (as; + w.y))
(17)
where ze the normalizing constant of the exponential family with the sufficient statistic and base we have that
measure μ.. By Theorem 2
Tr(MC (x1))
Pi (d M,
Tr(M Ke) 11 (dxx),
where (Ce(ai))is = ho (w2i (T + b.). (w2j, t + bj)). But now since
exxpp (./2), applying (17 gives
| (w2 + w22)
exp ( (w + +bi) exp ( (w (w] + b,))
(Ce(z))lss = ze.2
where Ze, 2 is the normalizing constant of the exponential family with the sufficient statistic to and base measure u2z Thus, write
Tr(MC (x1))
mij 2e,2
[ (w22 +w2s)
( (wo + +i)) exp ( (w]])((Z)) + +b,)) }
exp
=exp ( (wTT((Z1)+ M exp { (WiT (z + b)
and we conclude that the marginal in the same family with M M Zo, where
ZZ ZZ..kk = Ze,2
(1
( (w2z + (22))
Note that M is PSD as an Hadamard product of two PSD matrices. Thus, we can find İ such that M V V.
22

---

SNEFY
Resampled
Gauss
GMM 1
Moons
1.59 + 0.03
-I.76 + 0.02
-3.29 E 0.01
1.57 E 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
11.94 + 0.03
3..7 + 0.01
2.000 + 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 NVP layer.
SNEFY 2
Resampled:
Gauss 2
GMM 2
Moons
1.58 + 0.02
-1.59 = 0.03
1.622 + 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
-.077 1 0.04
2266 + 0.05
- 1.99 1 0.04
2621
69197
2384
2430
899.77 =t: 196.54
363.91 +: 36.62
267.23 + 32.05
289.96 + 33.77 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
Experiments
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 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 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 high precision normalising constant if the validation performance is the lowest encountered so far. We train each models for maximum of 20000 iterations (while monitoring validation performance). We use Adam with default hyperparameters and weight decay 10-3 The batch size is 21.
We use a SNEFY with Gaussian mixture model base measure supported on X = 1R2 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 MIP with layer widths [2, 256, 256, 1] and sigmoid activations for the resampling distribution. We use 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 the test log likelihood, not the exact log likelihood as in other methods. We choose The Gaussian mixture model has 10 mixture components, each with a diagonal covariance matrix. Each normalising flow block consists of an affine coupling block, permutation layer, and an actnorm layer.
50 and 1.
F.2 Modelling distributions on the sphere
We compare the performance of VMF distribution, a "regular" VMF mixture distribution and our SNEFY construction of the VMF mixture, which allows for some negative weight coefficients discussed in 3.2) We use the dataset 44 also used by [19] in different context for point
23

---

SNEPY 4
Resampled 4
Gauss
GMM
Moons
1.58 + 0.03
-1.60 = 0.07
1.60 + 0.03
- 1.57 E 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.99 + 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
-2344 = 0.06
2.37 0.17
-2.99 = 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 As 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/CC consists of locations of some 10610 galaxies. We use the spherical coordinates of these galaxies and map them to the surface of sphere.
We compare two types of SNEFYg2 soxp,,d with m = n = 30: one is constrained so that V is diago- nal (and therefore V V diagonal with nonnegative entries), and the other uses unconstrained general V. We expect the unconstrained model be more expressive, and therefore obtain better test NLL.
We Adam with default hyperparameters and 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 MIP that uses ReI.U activations in each layer and batch normalisation between each layer. Our SNEFY model uses Snake, activations with 10. The CNF models utilise affine layers wtih 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 batch size of 256. The dataset consists of 74309 training and 74557 examples, which are fixed between each run. Each run uses 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 o to be complex-valued and accordingly the readout parameters V to be complex, i.c. V Cm Note that in this case
2
m Σ
|f V( (Wo(())))
Σ
o (w) t(x) + bb))
24

---

Take a(2) = exp(iu) and t(x) a.
Then the ahove takes the form
crj exp [w] +bb;))
Uofüne" [w] x xmb1) e -i(wi atbi)
Σ ΣΣΣ Uyyunell bb1) b wi)
As in the a = exp case, bias terms can be folded into the readout parameters (which the reason why we also require readout parameters to be complex). We note that the case 1 is not interesting as it simply reduces to the (normalised) base measure μ..
In order to obtain the normalizing constant Tr Ve), we need the integral of the form
/ (((w))-))) ()) (((x)) =: k(wj wi).
Kel kexp(())).dd.e (W, wi)
(18)
where & is simply the Fourier transform of the base measure μ. i.c. its characteristic function in the case where μ is probability measure. Hence many standard choices of ju lead to tractable normalizing constants. Note that while and x both may be complex-valued, the normalizing constant well as the density itself are real-valued.
Various examples of probability measures μ and its Fourier transforms that give rise iivirantvariant positive definite kernels x have been studied in the literature on Random Fourier Features (RFF) 38
Here too, while the functional form of the expressions identical, the role between the data and the parameters is reversed: in RFF, one is interested in approximating 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. we restrict 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
2
172
Org exp (i (ww x + 6,))
= =| cos (Wx WW++bll b)1l + V sin (Wa + b)]
==1
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 [ Here we discuss some other settings.
NNK convergent series Suppose the support X is discrete and let h(x)) be nonnegative function corresponding with the discrete base measure The NNK is given by
k 0,)
Σ
(w t(x: + (w) t(x) b2) h(x).
Such NNKs often resemble known convergent series.
25

---

Fourier series For example, if a = cos and C2 = h(x)) are some coefficients of convergent Fourier series, then
(0, 0z)
Ca cos (((x)(w) w2) + (br b2)) + cos (1())(w1 + w2) + (b) + b2)) 2=0
just series representation of sum of two periodic functions. For example, if C2 2 2(-1)"t 1) for x 2 1 and C otherwise, and t(x) = 2 the NNK is a sum of two sawtooth waves with frequencies W1 w2 and wi + w2 and phase offsets b1 b2 and bı + b2. 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/discrette SNEFYs
Mixed distributions can obtained by choosing the base measure to mixed continuous distri- bution. For example, choose X Rd (u(xx) = 1 (T(x) + 5(x)) dx, and take and to be any of the combinations leading to closed-form NNGPK (for example, as the identity and a as the error function, Leaky ReLU, GELU, cosine, snake). Then
(0c, = / a (W. ((x) (W) + (x)))xxxx+++++++
((( t(x) (x)
( 2 (() + t(0) + b,) + t (0,,0)))
which is closed form.
26


Reviews:
[{"review_id": "BfU6cUqjtQ", "reviewer": "Reviewer_tpyH", "summary": "The authors present a new class of neural network based density models. This new family, SNEFYs, is based on a Gaussian Process paradigm (kernel methods) and contrasts with other competitors like flow based density estimators by imposing less restrictions on the network dimensionality. SNEFYs are shown to be analytically tractable for some combinations of base measures and kernel choices, and the family of SNEFYs is moreover shown to be closed with respect to conditioning. A special case is discussed in which marginalization maintains the family as well. In a few empirical examples, SNEFYs are contrasted with Flow based and (a restricted example of) Mixture Density estimators and shown to have favorable performance.", "questions": "--", "limitations": "I considered the limitations section adequate, apart from a specific comment about sampling, which I suggest under \"weaknesses\" should have been made earlier in the paper.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Access to a new class of density estimators is in general a boon for the research community will very likely spur a lot of related activity. \nThe authors build on a framework for density estimation (fundamentally based on kernel methods) that is relatively under-explored, at least in the domain in which this reviewer is familiar with density models, which is the context of simulation based inference.\nThe paper is very technical, but in general quite well written.\n", "weaknesses": "The paper is written in a very technical flavor, tailored very much towards experts at the intersection of neural network and kernel methods. \n\nA clearer discussion, with examples on how SNEFYs interface with, are similar too and contrast with other methods in the density modeling domain would have been helpful to understand better how exactly (or rather when) a non-expert can utilize SNEFYs.\n\nRelatedly it is mentioned only shortly in the limitations that SNEFYs do not admit direct sampling (as flows may for example) and users have to resort to e.g. rejection sampling. This feature should be part of the initial discussion of other methods and the relative strengths and weakness of SNEFYs. \n\nThe disconnect between Flows and SNEFYs boil down to the difference in architecture limitations between the two methods, a priori not suggesting necessarily where in application domains the respective strengths may lie. However the examples gave the impression that SNEFYs seem specifically to outperform Flows on low dimensional problems (where flows are severely limited in widths and have to make up for in solely through depth of the transformation). Numerical examples (and figures) that clearly delineate where the performance gain emerges and where it (potentially) diminishes would have been illustrative and probably crucial cornerstones for practitioners to come away with an intuition concerning relevant application areas for SNEFYs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review. We are glad that you believe this area of interest (new classes of density estimators) will very likely spur a lot of related activity. We agree that this area is relatively under-explored, and are happy to hear that you found the paper well-written. We address your individual concerns below. Please let us know if you would like us to elaborate further.\n\n**Weaknesses:**\n- \" *The paper is written in a very technical flavor, tailored very much towards experts at the intersection of neural network and kernel methods.  A clearer discussion, with examples on how SNEFYs interface with, are similar too and contrast with other methods in the density modeling domain would have been helpful to understand better how exactly (or rather when) a non-expert can utilize SNEFYs.* \" \\\nThanks for your suggestion. We will add a paragraph to the updated manuscript. In short, SNEFYs admit closed-form normalising constants, even under marginalisation and conditioning. This stands in contrast with energy based models, which do not admit closed-form normalising constants. Normalising flows can be used to model joint or conditional distributions with closed form normalising constants. However, we are not aware of any normalising flow which can be conditioned or marginalised post-hoc in closed form. This means that one is not able to use partial observations in a (joint) density estimation task via maximum likelihood. The closest we are able to find for normalising flows that deal with missing data is an approach that jointly optimises conditional distributions for missing data and model parameters [50] or samples from missing data and optimises for model parameters. This is not maximum likelihood or sampling from a joint likelihood. In contrast, we probabilistically marginalise out missing data exactly, so that we may perform maximum likelihood. This can be beneficial when dealing with missing data --- see figure and text in global rebuttal. Further, we may use SNEFYs even without constraining the neural network mapping to be bijective. If the user suspects that the target random variable is not an invertible transformation of some simple base distribution, they might consider using a SNEFY. Having said that, normalising flows may be a better modelling choice when the user requires samples from the learnt distribution (see your next point below).\n\n- \"*it is mentioned only shortly in the limitations that SNEFYs do not admit direct sampling (as flows may for example) and users have to resort to e.g. rejection sampling. This feature should be part of the initial discussion of other methods and the relative strengths and weakness of SNEFYs.*\" \\\nWe agree with this point that this should be brought up earlier in the paper. See \"Sampling\" section of \"General response to reviewers\". In short, we will update the Introduction section of our paper to state that our method does not always allow for sampling (beyond naive rejection sampling), but normalising flows do allow this.\n\n- \" *The disconnect between Flows and SNEFYs boil down to the difference in architecture limitations between the two methods, a priori not suggesting necessarily where in application domains the respective strengths may lie. However the examples gave the impression that SNEFYs seem specifically to outperform Flows on low dimensional problems (where flows are severely limited in widths and have to make up for in solely through depth of the transformation). Numerical examples (and figures) that clearly delineate where the performance gain emerges and where it (potentially) diminishes would have been illustrative and probably crucial cornerstones for practitioners to come away with an intuition concerning relevant application areas for SNEFYs.* \\\nWe have added one additional experiment that shows where a performance gain emerges compared with normalising flows. This is given in the general rebuttal, and concerns situations where data is missing. SNEFYs are able to be marginalised in closed-form (Theorem 2), so they can be trained using maximum likelihood estimation using partial observations. Compared with throwing away missing datapoints, SNEFYs trained using partial observations show a large increase in predictive performance. They also improve against normalising flows that discard partial observations.\n\n[50, not cited in original submission] Training Deep Normalizing Flow Models in Highly Incomplete Data Scenarios with Prior Regularization\n[51, not cited in original submission] MCFlow: Monte Carlo Flow Models for Data Imputation"}, {"review_id": "N7CapbONeI", "reviewer": "Reviewer_Y6Kj", "summary": "This paper introduces a parametric family of densities which admit tractable normalizing constants when a related neural-network kernel can be efficiently computed (e.g., in closed form).\nThe model is parameterized by the weights of a one-hidden-layer network, is closed under conditioning and admits tractable marginals that have a similar form as the global model.\nThe model is evaluated in unconditonal density modeling on synthetic data, and conditional density modeling on a redshift estimation task.\n\nEDIT: I have read the author's rebuttal, which addressed my concerns.", "questions": "I did not understand the experiment setup in the unconditional density modeling experiments: how is the normalizing flow combined with the SNEFY model? It thought that the normalizing flow would play the role of the sufficient statistic $t$, which would prevent a closed-form computation of the kernel, but it seems not to be the case.\n\nWhat is the relationship between considering a more advanced base density model (such as SNEFY rather than GMM) and more expressive normalizing flow? As a crude estimation, it might be fairer to compare a SNEFY model with a network of depth $L$ to a GMM model with a network of depth $L + 1$. In that case, the improvements of the SNEFY model are reduced. What does one gain by using a more expressive density model as opposed to a more expressive network?", "limitations": "The authors acknowledged the limitations of their approach: limited expressivity, and the general limitations of kernel-related methods. In particular, its complexity is cubic in the dimensionality (assuming $n$, $m$ and $d$ having similar orders of magnitude).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is very well-written and clear.\nThe proposed density model bears resemblance with the one considered in prior work [23, 38], but SNEFY is linear in the kernel as opposed to quadratic because it does not rely on a representer theorem.\nIt thus enjoys the tractability properties mentioned above.\nI am not expert enough in this area to judge of the significance of the approach, but it seems interesting enough to me to be worth publishing.", "weaknesses": "The authors mention that the sufficient statistics $t_2$ are unrestricted in the unconditional case. However, it seems to me that this not a property of SNEFY per se but common to any conditional density model. For instance, one could consider a conditional GMM model where the means and variances of $x_1$ are general functions of $x_2$, which would then admit closed-form normalizing constants and marginalization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your kind review and useful feedback. In addition to the general rebuttal, we address your individual concerns below. Please let us know if you would like to discuss further.\n\n**Strengths:**\n- You are right to highlight a strength of our model is that it scales linearly in the number of datapoints, rather than other kernel based approaches which scale quadratically in the number of datapoints. We did not explicitly mention this in the paper (excluding the mention of this in the context of Poisson point processes at the end of the conclusion), so we will update the manuscript before publication.\n\n**Weaknesses:**\n- \"*The authors mention that the sufficient statistics are unrestricted in the unconditional case. However, it seems to me that this not a property of SNEFY per se but common to any conditional density model. For instance, one could consider a conditional GMM model where the means and variances of $x_1$ are general functions of $x_2$, which would then admit closed-form normalizing constants and marginalization.*\" \\\nThanks for bringing this up. You are right to say that a conditional GMM can have parameters that depend on the conditioning variable $x_2$ while still admitting a tractable normalising constant. \nOur claim was intended to be related to normalising flow models, not mixture models. Normalising flows in fact restrict dependence on the conditioning variable. In particular, in a conditional normalising flow (CNF), the mapping of each layer would have to be bijective in both $x_1$ and $x_2$ (see text above equation (7) in [48]). The point we were trying to make here is that CNFs (but not GMMs or other more flexible SNEFYs), the mapping is restricted to a bijection in both the variable of interest $x_1$ and the conditioning variable $x_2$, whereas we do not have this bijection requirement. We will update the text above section 5 to further clarify. \n\n**Questions:**\n- \"*how is the normalizing flow combined with the SNEFY model? I thought that the normalizing flow would play the role of the sufficient statistic, which would prevent a closed-form computation of the kernel, but it seems not to be the case.*\" \\\nThe SNEFY model is used as what is called the \"base distribution\" of the normalising flow (which is different to the base measure of SNEFY). For example, the most common base distribution in normalising flows is a Gaussian, so that a normalising flow likelihood represents the likelihood of a Gaussian random variable passed through a sequence of normalising flow layers. Instead of passing a Gaussian random variable through the normalising flow layers, we compute the likelihood of passing a SNEFY random variable through normalising flow layers. This allows combining the closed-form machinery of SNEFY with the closed-form machinery of normalising flows. We will update the experimental details appendix of the manuscript to clarify this point --- thanks for raising it.\n- \" *What is the relationship between considering a more advanced base density model (such as SNEFY rather than GMM) and more expressive normalizing flow? As a crude estimation, it might be fairer to compare a SNEFY model with a network of depth to a GMM model with a network of depth . In that case, the improvements of the SNEFY model are reduced. What does one gain by using a more expressive density model as opposed to a more expressive network?* \" \\\nIt is difficult to answer your first question mathematically, however we did try to investigate this empirically. For example, in table 3, we compare a zero layer normalising flow with SNEFY base (i.e. just a SNEFY) against normalising flows of various depths and Gaussian base distributions. We find that in this case, SNEFY outperforms normalising flows. As another 3 examples, compare table 2 and table 4 to tables 5 to 7. We see that a SNEFY with zero or 1 normalising flow layers appended to it is competitive with normalising flows with more layers, although a resampled base layer does beat SNEFY on Rings.\n\n\n"}, {"review_id": "67OKik2ixw", "reviewer": "Reviewer_zjNK", "summary": "The authors introduce a class of normalized probability densities that are obtained by squaring a single-layer neural network. The class of densities, dubbed Squared Neural Family (SNEFY), offers a number of attractive properties including: i) they can be understood as a generalization of exponential family distributions: ii) normalizing constants can be computed in closed form; iii) they are closed under conditioning under certain factorization conditions; and iv) they are closed under marginalization in some special cases. The authors demonstrate that SNEFY can yield competitive log likelihoods on a few relatively simple density estimation tasks.", "questions": "- What are some of the practical trade-offs that you find empirically when using different values of m and n? E.g. do you see numerical instabilities or difficulties in optimization in certain regimes?\n- What generic strategies might you recommend for choosing sufficient statistics `t(x)`? Is it generally necessary to learn these if you want a flexible model? If so does such joint learning tend to lead to difficulties in optimization? Or might it be sufficient to use random feature expansions or the like? One might suppose this might mitigate some optimization problems, should such problems occur. \n- \"For example, exploited...\" [Line 29] This is a run-on (i.e. incomplete) sentence.\n- Final line \"... offering a model that scales like O(n^2) instead...\". Do you mean O(n)?\n- Dense tables like Table 2 are difficult to parse; consider making a plot", "limitations": "The main methodological limitation is the intractability of sampling, which the authors emphasize.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "As I see it the main strengths of this submission include:\n- It is very clearly written\n- It addresses a problem (namely defining useful classes of density estimators) that is of wide interest in the NeurIPS community\n- The proposed class of models is, to the best of my knowledge, novel\n- The proposed class of models has some attractive properties, especially w.r.t. conditioning and marginalization\n- The authors do a good job of placing their work in the context of related literature and methods\n\n", "weaknesses": "As I see it the main limitations and weaknesses of this submission are:\n- The experimental evaluation is quite limited, making it difficult to judge the likely usefulness or impact of SNEFY in practice\n- The intractability of obtaining samples from SNEFY distributions is a serious limitation (as the authors acknowledge)\n\nTo elaborate on the first point, the authors only consider three experiments, the first of which is in 2D. Indeed the only neural baseline appears to be a CNF described as follows in the appendix:\n> The CNF models utilise affine layers wtih tanh activations.\n\nNot only is this an inadequate description of this particular baseline, but the reader is left to wonder whether SNEFY models are actually competitive with more traditional normalizing flows in most density estimation tasks. Perhaps the joint learning of W and V in SNEFY doesn't work very well in practice because the parameterization is difficult to optimize and there's a tendency to get stuck in bad local optima? Even if SNEFY is not competitive with more traditional normalizing flow methods on common density estimation tasks, one might hope that its real value is to be found from leveraging some of its special properties, especially w.r.t. marginalization and conditioning. Unfortunately the experiments are not setup to demonstrate anything like this to the reader. The closest we have is the conditioning experiment involving astronomy data. While this is a good start, it is hardly convincing on its own. The authors note that \"due to this tractability of the marginal distributions, it is straightforward to include the likelihood corresponding to incomplete observations\"---unfortunately this experiment, which the reader anticipates with bated breath, is not forthcoming (similarly with the point process suggestion in the discussion section).\n\nGranted, in its current form this paper makes some interesting mathematical and methodological contributions, but I believe a much stronger version of this paper would work harder to convince the reader that SNEFY might be more broadly useful. Otherwise one is left wondering whether SNEFY is just a mathematical curiosity, elegant in theory but not very practical or useful in real applications. It is of course not necessary to \"beat\" 100 baselines on 100 benchmarks, but I would love to see more evidence for the practical utility of SNEFY.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your in-depth review. We were very happy to read that you found that our model is novel, addresses a problem that is of wide interest to the community, admits attractive properties, and is well-placed in the context of related methods. We agree entirely with the limitation of sampling. Please see our \"general response to reviewers\" for more on this. Please see our comments below for responses to individual queries. \n\n**Weaknesses:**\n- *Further expeiments on partial observations* \\\nThanks for the excellent suggestion about experimenting with partial observations. We will add the following experiment. We consider the joint distribution of a random vector $\\mathsf{x} \\in \\mathbb{R}^6$, with samples from the joint inputs and outputs of the astronomy dataset. 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 $p$ (i.e. the corresponding dimension is missing from observations). We then train a SNEFY model that utilises partial observations. We leave the original testing dataset, a matrix in $\\mathbb{R}^{74557 \\times 6}$, untouched, and measure the test NLL after training. We vary $p$ between $0.01$ and $1.0$ in steps of $0.01$, plotting the NLL as a function of $p$.  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 normalising flow baselines. Our observations are twofold: (1) Including partial observations increases performance and (2) Adding more complete observations increases performance. See attached figure in the general global rebuttal. We hope that the inclusion of such a demonstration does highlight some of the special properties of SNEFY, and satisfies your suggestion. Please let us know what you think.\n\n- *Poisson point processes, optimisation properties ...* \\\nWe agree that a detailed experimental evaluation of SNEFY would be interesting, especially teasing apart the novel aspects of the optimisation problem. We do hope to evaluate SNEFYs with respect to these new features in follow-up works, where we can properly dedicate enough room to subtle concepts. Poisson point processes are a very interesting avenue. In particular we have found that the form of the likelihood under the Poisson point process model leads to some optimisation problems with properties different from those of the likelihood maximization under an i.i.d. sampling from a density model. The synthetic experiments and real experiment on astronomy data were intended to show that in some special cases, SNEFYs are a competitive alternative to kernel density estimators, naive mixture models, and normalising flows. \n\n\n**Questions:**\n- \"*What generic strategies might you recommend for choosing sufficient statistics $t(x)$? Is it generally necessary to learn these if you want a flexible model? If so does such joint learning tend to lead to difficulties in optimization? Or might it be sufficient to use random feature expansions or the like? One might suppose this might mitigate some optimization problems, should such problems occur.* \" \\\nAs in the case for classical exponential families with closed form normalising constants, in this paper we take sufficient statistics to be known functions which allow for easy integration. \nSo in some sense, the strategy is to use a table of known closed-forms (like table 1), and work with what you can use.\nInvestigating whether there are interesting cases where $t$ itself is a neural network or a random feature expansion and the normalising constant is still tractable is another important avenue of further research. This is a technical challenge because for arbitrary $t$, it is not easy to compute the closed form normalising constants via the kernel as in table 1. \n\n- \" *\"For example, exploited...\" [Line 29] This is a run-on (i.e. incomplete) sentence.* \"  \\\nThanks for catching this typo. We will update this in the manuscript.\n- \" *\"Final line \"... offering a model that scales like $\\mathcal{O}(n^2)$ instead...\".* \"  \\\nThanks for catching this error. We meant to say, and will update the manuscript to say, \"For example, we may build a Poisson point process intensity function using a squared neural network and compute the intensity function in closed-form, offering a model that scales quadratically in the number of neurons $\\mathcal{O}(n^2)$ instead of comparable models which scale cubically in the number of datapoints $\\mathcal{O}(N^3)$\" [8, 45].\"\n\n- \" *Dense tables like Table 2 are difficult to parse; consider making a plot* \" \\\nThanks for this suggestion. We will do this in the final version.\n\n"}, {"review_id": "Tq5Mrvblct", "reviewer": "Reviewer_k4Y8", "summary": "The paper introduces SNEFY, a tractable probabilistic model obtained by squaring the Euclidean norm of a neural net and normalizing it w.r.t. a base measure $\\mu$ (cf. Equation 1).\nThe paper is mainly focused on studying settings in which one can obtain a closed form tractable solution for the normalizing constant of a given neural net.\nThese settings, which are listed in Table 1, are characterized by: a specific support, a base measure, a sufficient statistic, an activation function, bias/no bias and a kernel function (KNN).\nFor these settings, the normalization constant can be computed as a trace of a matrix which is quadratic in the number of neurons of the neural net (cf. Equation 6).\nThis matrix is a PSD matrix whose elements result from the application of a kernel function (NNK) over the hidden weights and biases of the neural net.\nThe model is trained by MLE and aveluated on: three 2D-dimensional datasets, a spheirical manifold and on a low-dimensional astronomy dataset.\n", "questions": "- Why is \"conditional SNEFY\" (section 3.3) introduced before \"marginal SNEFY\" (section 3.4)? Doesn't tractable conditioning follow from tractable marginalisation ($P(A|B)=P(A \\cap B) / P(B)$)?\n- It is not clear to me whethere the method can be applied on discrete data. If so: (1) why didn't you include a  (simple) case study in your experiments? (2) Is it possible to handle both discrete and continuous features simultaneously?\n", "limitations": "I find authors to be open about the limitations of their model: They show that their model is slow to train (Tables 2-3) and that sampling is computationally expensive.\nBut probably, more emphasis and discussion about the long training time required would be appreciated.\n\n", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- new NNK are proposed\n- tractable marginalisation and conditioning\n- connection with exponential families (every exponential famility is a SNEFY) and kernel-based methods for nonnegative functions\n", "weaknesses": "- SNEFY seems to be very slow to train and, moreover, performs very similar to GMMs when used as prior of a Normalizing Flow (see Tables 2-3, Appendix F2).\n- Experiments are not really convincing and rather small scale. As a consequence, it is pretty unclear how expressive the model can actually be. Specifically, the 2D densities considered (see experiment section) are rather easy to fit. I'd suggest authors to investigate more complex 2D densities as the ones in [1].\n- Sampling seems to be possible only by rejection sampling, and is therefore very computationally expensive. Could you elaborate more about sampling and its challenges?\n- Minor: In general, I find that the exposition of the manuscript can be considerably improved. For instance, many equations/concepts are introduced more than once (compare contribution-background sections, Equation 1-2-3). Best results are not reported in bold, thus making Tables 2-3 not easy to read. There's no vspace between caption and Table 1. Table 3 is not referenced in the main text\n\n[1] Nash, Charlie, and Conor Durkan. \"Autoregressive energy machines.\" International Conference on Machine Learning. PMLR, 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Our thanks for your thoughtful review and suggestions.\n\n**Weaknesses:**\n- \" *Sampling seems to be possible only by rejection sampling, and is therefore very computationally expensive. Could you elaborate more about sampling?* \"\\\nThere are some interesting approaches (beyond rejection sampling) and challenges associated with sampling. Please see our general response to reviewers for our comments on sampling.\n- \" *SNEFY seems to be very slow to train and, moreover, performs very similar to GMMs when used as prior of a Normalizing Flow (see Tables 2-3, Appendix F2).* \" \\\nSNEFY (by itself, even without a normalising flow) outperforms the best GMM (which is allowed to have between 0 and 16 normalising flow layers) on Circles. It only takes on average 644 seconds, which is faster than the best performing GMM at at average of 862.25 seconds. \n- \" *the 2D densities considered (see experiment section) are rather easy to fit. I'd suggest authors to investigate more complex 2D densities as the ones in [1].* \"\\\nWe took a look at Einstein in [1]. In order to reproduce [1], 3 million training epochs will take roughly 20 of our GPU hours. Detailed experimentation on this data is unfortunately out of our budget at this stage. We will cite [1] as an EBM in the intro, where normalising constants are approximated rather than computed exactly.\n\n**Minor**\n- \" *Best results are not reported in bold, thus making Tables 2-3 not easy to read. There's no vspace between caption and Table 1.* \" \\\nThanks for your suggestions to improve the clarity of the tables. We will update these.\n\n- \" *Table 3 is not referenced in the main text* \" \\\nThanks for spotting this. We have updated the end of line 295 to refer to Table 3.\n\n**Questions:**\n- \" *Why is \"conditional SNEFY\" (section 3.3) introduced before \"marginal SNEFY\" (section 3.4)? Doesn't tractable conditioning follow from tractable marginalisation?* \"\\\nWe agree that tractable conditioning follows from tractable marginalisation, but the reverse also holds (by just rearranging Bayes' rule), which is the order of our proof. We don't see a strong reason to prefer one order over the other, however if you have reasons to preference one over the other, we are happy to swap the order according to your suggestions. Please let us know.\n\n- \" *It is not clear to me whether the method can be applied on discrete data. If so: (1) why didn't you include a (simple) case study in your experiments? (2) Is it possible to handle both discrete and continuous features simultaneously?* \" \\\nThanks for the great question. This made us think of some new special cases, which we will add to appendix. In detail:\n\n**Discrete and mixed continuous/discrete SNEFY:** \\\nSNEFYs can indeed be applied to discrete data.\nClosed-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.\nThis is the case if the support has very large or infinite cardinality. \nAn example that extends the classical Poisson distribution in exponential family form is given in Table 1.\nHere we discuss some other settings.\n\n**NNK as a convergent series**\nSuppose the support $\\mathbb{X}$ is discrete and let $h(\\mathbf{x})$ be a nonnegative function corresponding with the discrete base measure $\\mu$. \nThe NNK is given by\n$$ k_{\\sigma,\\mathbf{t},\\mu}(\\mathbf{\\theta}_i, \\mathbf{\\theta}_j) =  \\sum   \\sigma\\big(\\mathbf{w}_1^\\top \\mathbf{t}(\\mathbf{x}) + b_1\\big) \\sigma\\big(\\mathbf{w}_2^\\top \\mathbf{t}(\\mathbf{x}) + b_2\\big) h(\\mathbf{x}),$$\nwhere the sum runs over $\\mathbf{x} \\in \\mathbb{X}$.\n\nSuch NNKs often resemble known convergent series. \nFor example, if $d=1$, $\\sigma = \\cos$ and $C_x = h(x)$ are some coefficients of a convergent Fourier series, then\n$$ k_{\\sigma,\\mathbf{t},\\mu}(\\mathbf{\\theta}_i, \\mathbf{\\theta}_j) = \\frac{1}{2} \\sum  C_x\\Big(  \\cos \\big(t(x) (w_1 - w_2) + (b_1-b_2) \\big) + \\cos \\big(t(x) (w_1 + w_2) + (b_1+b_2)  \\big) \\Big) $$\nis just a series representation of a sum of two periodic functions. \nFor 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$.\nOther examples include rectified sine, square and triangular waves.\nThis extends to functions with convergent multivariate Fourier series.\nWe expect that the utility of such models is limited or difficult to uncover with benchmark datasets that we are aware of.\nIf the cardinality of the discrete domain is small, discrete distributions can be flexibly represented by standard histograms, so utility of SNEFY parametrisation is unclear.\nIn the SNEFY parameterisation, one biases the model away from general histograms, which may be favourable in terms of sample complexity if the data matches the bias or the discrete domain is very large, but otherwise not.\n\n**Mixed discrete/continuous:** \\\nWe did not originally think about mixed discrete/continuous distributions, but it is possible. \nWe imagine such models may find more utility than the purely discrete case, although we have not looked into this further.\nChoose a mixed discrete/continuous base measure (e.g. a spike and slab Gaussian). Such induced NNKs will then be tractable in the same setting as the continuous component of the base measure is tractable, under extremely mild conditions on the activation function. E.g. if the base measure is a sum of a Gaussian measure and a Dirac delta, the resulting kernel is just a sum of the kernel with the Gaussian base measure and the product of the activations evaluated at $\\mathbf{x} = \\mathbf{0}.$ We will add this example to the appendix."}]

Author rebuttal:
**General response to reviewers:**

We thank the reviewers for their careful, detailed feedback and overall positive evaluation. We address commonly raised concerns here and respond to individual queries separately. We look forward to further engaging with the reviewers during the reviwer-author discussion period, should they have any additional concerns.

**Sampling**

We completely agree that sampling is the biggest limitation of our work. We tried to emphasise this limitation (see reviews of zjNK, k4Y8). Following tpyH's suggestion to emphasise the sampling limitation earlier on, we will update the Introduction section of our paper to state that our method does not always allow for sampling (beyond naive rejection sampling), but normalising flows do allow this. The problem we are aiming to solve is (conditional) density estimation, different to sampling. Many of the best sampling tools (GANs, VAEs, ...) do not admit likelihood calculation. Normalising flows are to the best of our knowledge unique in this respect, paying the price of requiring a bijective transformation.

We will elaborate in our future work section, which currently states that "Future work will focus on sampling, as has been done with related models [ 24 ]," so that includes a more concrete description. In particular, the current results of [24] allow for principled sampling in the special case when the activation function $\sigma$ is Gaussian and the support $\mathbb{X}$ is a hyperrectangle. It obtains $r$ approximate samples with $\mathcal{O}\big(r \log_2 (|\mathbb{X}|)) + r d \log_2\frac{2}{\rho} \big) + 1 $ evaluations of a normalising constant (which for us, costs $\mathcal{O}(m^2n + dn^2 )$). Here $\rho$ is an approximation tolerance parameter. Note that this complexity significantly improves upon naive rejection sampling, which typically scales exponentially in dimension $d$. There is nothing special about the Gaussian activation function for the algorithm to run, but the approximation error is only computed for the Gaussian activation. Also, the support $\mathbb{X}$ need not be a hyperrectangle. In short, yes, we agree that sampling is a current limitation of SNEFY. However, like other related models [23,24,38], there are bespoke sampling algorithms in special cases and future avenues for expanding the sampling capabilities of SNEFY.

**New experiment on partial observations**

We will add the following experiment. We consider the joint distribution of a random vector $\mathsf{x} \in \mathbb{R}^6$, with samples from the joint inputs and outputs of the astronomy dataset. 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 $p$ (i.e. the corresponding dimension is missing from observations). We then train a SNEFY model that utilises partial observations. We leave the original testing dataset, a matrix in $\mathbb{R}^{74557 \times 6}$, untouched, and measure the test NLL after training. We vary $p$ between $0.01$ and $1.0$ in steps of $0.01$, plotting the NLL as a function of $p$.  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 normalising flow baselines. Our observations are twofold: (1) Including partial observations improves performance and (2) Adding more complete observations improves performance. See attached pdf.

Meta-review: The paper proposes a new type of density model which facilitates tractable marginals and conditionals. The paper is solid, original and of wide interest to the NeurIPS community. A clear accept, and I think this paper would be a good candidate for spotlight.

There were some minor criticisms by the reviewers, which, however, don't hamper the quality of the paper. As far as possible, the authors should please consider them for the camera-ready version.
Decision: Accept (spotlight)