Title: Squared Neural Families: A New Class of Tractable Density Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: probabilistic modelling; density estimation; exponential family;
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:

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.

1 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 on mathematically convenient models such as exponential families [50]. Such models can often be normalised in closed form, but are often only suitable for relatively low-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 [35, §2.6.1, §2.2] for more energy-based and other methods.

Modern computer technology and power allows us more flexible models [8, page 68], by partially or completely relaxing the requirement for mathematical tractability. Contemporary high-dimensional modelling, such as generative image models, rely primarily on neural network models [18, 11, 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 we model the densities directly without using constrained transformations of random variables. Moreover, our approach can be readily applied in conjunction with normalising flows and other deep learning devices such as deep feature extractors.

̸

Support

X

Base

measure

µ

Sufficient

statistic

t

(

x

)

Activation

function

σ

b

=

0

Kernel

k

σ,

t

,µ

Any setting mirroring a previously derived NNGPK, e.g.

R

d

Φ

C

,

0

Id(

x

)

erf

✗

[52]

(

·

)

p

+

, p

∈

N

✗

[4]

LReLU

✗

[48]

GELU

[14]

✗

[47]

cos

✗

[37]

Any tractable exponential family [50, 32] setting, e.g.

R

d

Φ

C

,

m

x

∥

x

∥

2

exp

✓

Kernel 3

S

d

-

1

Uniform

Id(

x

)

✓

R

d

Φ

C

,

m

✓

Kernel 7

{

0

,

1

,

2

, . . .

}

(

x

!)

-

1

ν

✓

Kernel 8

New tractable integration settings, e.g.

R

d

Φ

C

,

m

Id(

x

)

cos

✓

Kernel 2

[0

,

1]

d

Uniform

Φ

-

1

(

x

)

R

d

Φ

C

,

m

Id(

x

)

Snake

a

[55, 39]

✓

Kernel 6

̸

Contributions. Let (Ω , F , µ ) denote a measure space with Ω ⊆ R d , sigma algebra F , and nonnegative measure µ . Let V ∈ R m × n be the readout parameters of a neural network and let W ∈ R n × D and b ∈ R n be the weights and biases of a hidden layer of a neural network f : R D → R m with activation function σ . For some support X ∈ F , define a probability distribution P (and corresponding probability density p with respect to base measure µ ) to be proportional to squared 2 -norm of the evaluation of the neural network f ,

whenever the normalising constant z( V , Θ ) ≜ ∫ X ∥ ∥ f ( t ( x ); V , Θ ) ∥ ∥ 2 2 µ ( d x ) is finite and non-zero. Here we call µ the base measure, t : X → R D the sufficient statistic 2 and σ the activation function. We will call the corresponding family of probability distributions, parametrised by ( V , Θ ) , a squared neural family ( SNEFY ) on X , and denote it by SNEFY X , t ,σ,µ . SNEFY s are new flexible probability distribution models that strike a good balance between mathematical tractability, expressivity and computational efficiency. When a random vector x follows a SNEFY distribution indexed by parameters ( V , Θ ) , we write x ∼ SNEFY X , t ,σ,µ ( V , Θ ) or simply x ∼ P ( · ; V , Θ ) , where there is no ambiguity.

Our main technical challenge is in exactly computing the normalising constant, z( V , Θ ) , where

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 σ , sufficient statistic t and base measure µ over support X . See Table 1. Our analysis both exploits and informs a connection with

so-called neural network Gaussian process kernels (NNGPKs) [31] in a generalised form, which we refer to as neural network kernels (NNKs).

We discuss some important theoretical properties of SNEFY such as exact normalising constant calculation, marginal distributions, conditional distributions, and connections with other probability models. We then consider a deep learning setting, where SNEFY s 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 a and a ). The special zero vector (0 , . . . , 0) and identity matrix elements are 0 and I . We use subscripts to extract (groups of) indices, so that for example, w i is the i th row of the matrix W (as a column vector), v · ,i is the i th column of the matrix V , and b i is the i th element of the vector b . We use Θ = ( W , b ) ∈ R n × ( D +1) to denote the concatenated hidden layer weights and biases. Correspondingly, we write θ i = ( w i , b i ) ∈ R D +1 for the i th row of Θ . We will use a number of special functions. Φ C , m denotes the cumulative distribution function (CDF) of a Gaussian random vector with mean m and covariance matrix C . We also use a shorthand Φ C = Φ C , 0 and Φ = Φ I , 0 .

Single hidden layer neural networks We consider a feedforward neural network f : R D → R m ,

with activation function σ , hidden weights W ∈ R n × D and biases b ∈ R n and readout parameters V ∈ R m × n . Here σ is applied element-wise to its vector inputs, returning a vector of the same shape.

Neural network kernels In certain theories of deep learning, one often encounters a bivariate Gaussian integral called the neural network Gaussian process 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 wellbehaved (say, independent Gaussian) random weights converges to a zero-mean Gaussian process, so that the NNGPK characterises the law of the neural network predictions. These limiting models can be used as functional priors in a 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 θ i = ( w i , b i ) 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 k σ, Id , Φ I admits a representation of the form

̸

We do not discuss in detail how it is constructed in earlier works [31, 22, 33, 17], where usually b i = b j = 0 3 , but not always[49]. When b i = b j = 0 , closed-form expressions for the NNGPK are available for different choices of σ and t [52, 20, 4, 48, 37, 47, 28, 13]. However, the setting of b = 0 is important in our context (as we show in §2.2) and presents additional analytical challenges.

We will require a more general notion of an NNGPK which we call a 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 is common in kernels and covariance functions and can be used to induce desirable analytical and practical properties [37, 29, 24, §5.4.3]. We also integrate with respect to more general measures µ instead of the standard Gaussian CDF, Φ . We define the NNK to be

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 θ i , θ j of Θ , which are totally independent of V . The proof of the following is given in Appendix A.

Identity 1. The integral (2) admits a representation of the form

where k σ, t ,µ is as defined in (5) , and K Θ is the PSD matrix whose ij th entry is k σ, t ,µ ( θ i , θ j ) .

By Identity 1, the normalised measure (1) then admits the explicit representation

where ˜ K Θ ( x ) is the PSD matrix whose ij th entry is σ ( w ⊤ i t ( x ) + b i ) σ ( w ⊤ j t ( x ) + b j ) ⊤ . We used the cyclic property of the trace, writing the numerator as Tr ( σ ⊤ V ⊤ V σ ) = Tr ( V σσ ⊤ V ⊤ ) = Tr ( V ⊤ V σσ ⊤ ) . We emphasise again that the role of the data t ( x ) and the hidden weights and biases Θ in the NNK k σ, t ,µ are reversed compared with how they have appeared in previous settings. We may compute evaluations of k σ, t ,µ in closed form for special cases of ( σ, t , µ ) using various identities in O ( 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 2 n + dn 2 ) , where n and m are respectively the number of neurons in the first and second layers.

Remark 1 (Alternative parameterisations) . SNEFY 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 a normalising constant of 1 . Furthermore, noticing that V only appears as a PSD matrix M ≜ V ⊤ 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 evaluations of the NNK (5). Several closed-forms are known for different settings of σ , all with t = Id and b = 0 . The motivation behind derivation of existing known results is from the perspective of inference in infinitely wide Bayesian neural networks [31] or to 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, k σ, t ,µ , is parametrised by activation function σ , warping function (sufficient statistic) t , and base measure µ . All derivations are given in Appendix B.

The first kernel describes how we may express the kernels with arbitrary Gaussian base measures Φ C , m in terms of the kernels with isotropic Gaussian base measures Φ . This means that it suffices to consider isotropic Gaussian base measures in place of arbitrary Gaussian base measures.

Kernel 1. k σ, Id , Φ C , m ( θ i , θ j ) = k σ, Id , Φ ( T θ i , T θ j ) , where T Θ = ( WA , b + Wm ) , T θ i = ( w ⊤ i A , b i + w ⊤ i m ) and A is a matrix factor such that covariance C = AA ⊤ .

̸

This kernel can also be used to describe kernels corresponding with Gaussian mixture model base measures. The second kernel we describe is a minor extension to the case b = 0 of a previously considered kernel [37].

A similar result and derivation holds for the case of k sin , Id , Φ , which we do not reproduce here. We now mention a 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 proj S d -1 ( x ) ≜ x / ∥ x ∥ to be the projection onto the unit sphere. Then

where I p is the modified Bessel function of the first kind of order p . In the special case d = 3 , we have the closed-form k exp , proj S 2 , Φ ( θ i , θ j ) = exp( b i + b j ) ( e ∥ w i + w j ∥ -e -∥ w i + w j ∥ ) 2 ∥ w i + w j ∥ .

We end this section with a new analysis of the Snake a activation function, given by

The Snake a function [55] is a neural network activation function that can resemble the ReLU on an interval for special choices of a , is easy to differentiate, and as we see shortly, admits certain attractive analytical tractability. We note that a similar activation function has been found using reinforcement learning to search for good activation functions [39, Table 1 and 2, row 3], up 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, Snake a ( · ) -1 2 a (Kernel 5). We first describe the linear kernel.

We now derive the kernel corresponding with Snake a activation functions up to an offset.

Kernel 5. The kernel k Snake a ( · ) -1 2 a , Id , Φ ( θ i , θ j ) is equal to

The kernel corresponding with Snake a activations is then stated in terms of Kernel 4 and 5.

Kernel 6. The kernel k Snake a , Id , Φ ( θ i , θ j ) is equal to

3 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 ν , and dµ dν : Ω → [0 , ∞ ) is the Radon-Nikodym derivative, then the SNEFY admits a probability density function p ( · | V , Θ ) with respect to ν ,

The Fisher-Neyman theorem (for example, see Theorem 6.14 of [6]) says that the existence of such a factorisation is equivalent to the fact that t is a sufficient statistic for the parameters V , Θ .

3.2 Connections with exponential families

In this section we will use the activation σ ( u ) = exp( u/ 2) . We note that we can absorb the bias terms b into the V parameters 4 and obtain as a special case the following family of distributions

which is a mixture 5 of distributions P e ( · ; 1 2 ( w i + w j ) ) belonging to a classical exponential family P e [50, 32], given in the canonical form by

It is helpful to identify the following three further cases:

When n = m = 1 , v 2 11 cancels in the numerator and denominator and we obtain an exponential family with base measure µ supported on X , sufficient statistic t , canonical parameter w 1 and normalising constant z e ( w 1 ) . Every exponential family is thus a SNEFY , but not conversely.

When m> 1 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 ) .

When m> 1 and n > 1 and V ⊤ V is diagonal (i.e. each column in V is orthogonal), there are at most n non-zero mixture coefficients, all of which are nonnegative. That is, we obtain a standard exponential family mixture model.

The kernel matrix K Θ in the normalising constant of (8) is tractable whenever the normalising constant of the corresponding exponential family is itself tractable.

Proposition 1. Denote by z e ( w ) the normalising constant of the exponential family in (9) . Then

The above kernel is well defined for any collection of w i which belong to the canonical parameter space of P e , since the canonical parameter space is always convex [50]. This gives us a large number of tractable instances of SNEFY which correspond to exponential family mixture models allowing negative weights - a selection of examples is given in Appendix C. It is interesting that some properties of the exponential families are retained by this generalisation belonging to SNEFY s. For example, the following proposition, the proof of which is given in Appendix A, links the derivatives of the log-normalising constant to the mean and the covariance of the sufficient statistic.

Proposition 2. Let σ ( u ) = exp( u/ 2) and define the log-normalising constant as Ψ = log z( V , Θ ) .

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 = ( x 1 , x 2 ) be jointly SNEFY X 1 × X 2 , t ,σ,µ with parameters V and Θ = ([ W 1 , W 2 ] , b ) . Assume that µ ( d x ) = µ 1 ( d x 1 ) µ 2 ( d x 2 ) and t ( x ) = ( t 1 ( x 1 ) , t 2 ( x 2 ) ) . Then the conditional distribution of x 1 given x 2 = x 2 is SNEFY X 1 , t 1 ,σ,µ 1 with parameters V and Θ 1 | 2 ≜ ( W 1 , W 2 t 2 ( x 2 ) + b ) .

The proof, which we detail in Appendix A follows directly by folding the dependence on the conditioning variable x 2 into the bias term. We note that conditional density will typically be tractable if the joint density is tractable since they share the same activation function σ . Thus, whenever SNEFY corresponds to a tractable NNK with a non-zero bias, we can construct highly flexible conditional density models using SNEFY by taking t 2 itself to be a jointly trained deep neural network. Crucially, t 2 may be completely unconstrained . We use this observation in the experiments (§ 4).

3.4 Marginal distributions under SNEFY

Marginal distributions under SNEFY model for a general activation function σ need not belong to the same family. In the special case σ = exp( · / 2) , SNEFY is in fact also closed under marginalisation, which we prove in Appendix D. Even in the general σ case, marginal distributions are tractable and admit closed forms whenever the joint SNEFY model and the conditional SNEFY are tractable.

Theorem 2. Let x = ( x 1 , x 2 ) be jointly SNEFY X 1 × X 2 , t ,σ,µ with parameters V and Θ = ([ W 1 , W 2 ] , b ) . Assume that µ ( d x ) = µ 1 ( d x 1 ) µ 2 ( d x 2 ) and t ( x ) = ( t 1 ( x 1 ) , t 2 ( x 2 ) ) . Then the marginal distribution of x 1 is

The proof is given in Appendix A. Due to this tractability of the marginal distributions, it is straightforward to include the likelihood corresponding to incomplete observations (i.e. samples where we are missing some of the components of x ) into the density estimation task.

3.5 Connections with kernel-based methods for nonnegative functions

SNEFY may be viewed as a neural network variant of the non-parametric kernel models for non-negative functions [25], which are constructed as follows. Let ψ : X → H be a feature mapping to a (possibly infinite dimensional) Hilbert space H . Let S ( H ) be the set of all positive semidefinite (PSD) bounded linear operators A : H → H . Then

gives an elegant model for nonnegative functions parametrised by A ∈ S ( H ) , and their application to density modelling with respect to a base measure has also been explored [25, 42]. Note that by assuming boundedness of A and ψ , the normalizing constant [25, Proposition 4] is given by Tr ( A ∫ X ψ ( x ) ⊗ ψ ( x ) µ ( d x ) ) , which is analogous to our work where the normalising constant is given by z( V , Θ ) = Tr( V ⊤ V K Θ ) . This can be seen by replacing A with V ⊤ V and replacing ∫ X ψ ( x ) ⊗ ψ ( x ) µ ( d x ) by K Θ = ∫ X σ ( Wt ( x ) + b ) σ ( Wt ( x ) + b ) ⊤ µ ( d x ) .

Despite feature maps being infinite-dimensional, model (11) often reduces to an equivalent representation in finite-dimensions. [25] utilise a representer theorem when (11) is fitted to data using a regularised objective, while [42] more directly assume that the linear operator A is inside the span of the features evaluated at the available data { x ℓ } N ℓ =1 . The resulting model resembles SNEFY where n equals to the number N of datapoints, i.e.

However, there are fundamental differences between (12) and SNEFY which we list below. The models can be seen as complementary and they inherit advantages and disdvantages of kernel methods and neural networks common in other settings, respectively.

Tractability. Whereas we identify many tractable instances of SNEFY , the normalising constant of (12) requires computing ∫ κ ( x i , x ) κ ( x , x j ) µ ( d x ) - this is not generally tractable apart from some

limited combinations of κ and µ (e.g. for a Gaussian kernel κ and a Gaussian µ ). Note that this kernel is evaluated at the datapoints, whereas SNEFY evaluates the kernel at the learned parameters w i . [42] focuses on the specific case where κ is a Gaussian kernel, studying properties of the resulting density class which is a mixture of Gaussian densities allowing for negative weights, a model equivalent to SNEFY with the exponential activation function as described in Appendix C. Note that our treatment of SNEFY as a generalisation of the exponential family goes well beyond the Gaussian case, and that tractable instances arise with many other activation functions.

Expressivity. Crucially, the feature map ψ and consequently finite dimensional representation via κ 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 κ ), at the expense of a more involved optimiisation over the space of PSD matrices. In contrast, we learn W (analogous to learning ψ ) and V jointly using neural-network style gradient optimisers. SNEFY is fully compatible with end-to-end and jointly optimised neural network frameworks, a property we leverage heavily in our experiments in § 4.

Conditioning. By explicitly writing parametrisation which includes the biases, we obtain a family closed under conditioning and thus a natural model for conditional densities, whereas it is less clear how to approach conditioning when given a generic feature map ψ .

Other related work After submission, we became aware of another related literature, which includes mixture models with potentially negative mixture coefficients via squaring [23] and positive semi-definite probabilistic circuits [43]. We believe a marriage of ideas from SNEFY and probabilistic circuits will lead to future developments in tractable and expressive probability models.

4 Experiments

Implementation All experiments are conducted on a Dual Xeon 14-core E5-2690 with 30GB of reserved RAM and a single NVidia Tesla P100 GPU. Full experimental details are given in Appendix E. We build our implementation on top of normflows [45], a PyTorch package for normalising flows. SNEFY s are built as a BaseDistribution , which are base distributions inside a NormalizingFlow with greater than or equal to zero layers. We train all models via maximum likelihood estimation (MLE) i.e. minimising forward KL divergence.

2D synthetic unconditional density estimation We consider the 2 dimensional problems also benchmarked in [46]. We compare the test performance, computation time and parameter count of non-volume preserving flows (NVPs) [7] with four types of base distribution: SNEFY , resampled [46] ( Resampled ), diagonal Gaussian ( Gauss ) and Gaussian mixture model ( GMM ). We consider flow depths of 0, 1, 2, 4, 8 and 32 where a flow depth of 0 corresponds with the base distribution only. We use σ = cos , t = Id , X = R d and a Gaussian mixture model base density. We set m = 1 and n = 50 . Full architectures and further experimental details are described in Appendix E.1.

Results are shown in Table 2 for the 0 and 16 layer cases, and further tables for 1 , 2 , 4 and 8 layers are given in Appendix E.1. Our observations are that all base distributions are able to achieve good performance, provided they are appended with a normalising flow of appropriate depth. SNEFY is able to achieve good performance with depth 0 on all three datasets, as are Resampled and GMM for the Moons dataset. The parameter count for well-performing SNEFY models is very low, but the computation time can be relatively high. However, SNEFY is the only model which consistently achieves the highest performance within one standard deviation across all normalising flow depths.

Data on the sphere We compare the three cases mentioned in § 3.2 using Kernel 3, i.e. mixtures of the von Mises Fisher (VMF) distribution, as shown in Figure 1. Over 50 runs, the unconstrained V , diagonal V , and n = m = 1 cases respectively obtain test negative log likelihoods of 1 . 38 ± 9 . 64 × 10 -3 , 1 . 46 ± 0 . 016 and 2 . 34 ± 0 . 26 each in 111 . 18 ± 2 . 58 , 109 . 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 3. Full details are given in Appendix E.2.

Conditional density estimation on astronomy data Predicting plausible values for the velocity of distant astronomical objects (such as galaxies or quasars) without measuring their full spectra, but by

measuring shifted spectra through broadband filters, is known as photometric redshift. Photometric redshift is an important problem in modern astronomy, as large surveys have increased the amount of available data that do not directly measure spectra. We estimate distributions for cosmological redshift x 1 ∈ R conditional on features x 2 ∈ R 5 using a public dataset [2]. The features x 2 are the magnitude a broadband filter (red) and set of four pairwise distances between broadband colour filters (ultraviolet-green, green-red, red-infrared1, infrared1-infrared2). We benchmark our own SNEFY against a conditional kernel density estimator (CKDE) and a conditional normalising flow (CNF) [54]. We use a publically available implementation [10] of CNF. Our SNEFY here is not appended with any normalising flow layers. We use σ = Snake a , t = Id , X = R and a Gaussian base density. We use a deep conditional feature extractor with layer widths [2 5 n, 2 4 n, 2 3 n, 2 2 n, 2 n ] , and then set m = n/ 2 . Full details are given in Appendix E.3.

While CNFs have shown great promise in modelling high dimensional image data, we expect that they are not as well-suited to nontrivial tabular data with a small to medium number of dimensions. This is because each layer is required to be an invertible transformation in the variable being modelled, so the input and output sizes must be the same, thereby significantly limiting the parameter count in each layer. We could of course increase the number of layers in the CNF model to achieve higher parameter counts, however this results in a model that is difficult to train. In our experiments, we found that increasing the depth of the CNF decreased its performance. On the other hand, SNEFY may utilise any number of parameters, as the conditioning network t 2 is completely unrestricted.

Density estimation on astronomy data using partial (marginal) observations We perform joint probability density estimation using the redshift data. The original training dataset is a matrix in R 74309 × 6 . We partition this matrix into batches of size 256 and randomly set each column of each batch to NaN with probability q (i.e. the corresponding dimension is missing from observations). We then train a SNEFY model that utilises partial observations by maximising the marginal likelihood according to Theorem 2. We leave the original 74557 × 6 testing dataset untouched, and measure the test NLL after training. We plot the NLL as a function of 1 -q , as shown in Figure 2. We also compare the performance of the SNEFY that uses partial observations with the performance of SNEFY that simply throws away incomplete observations, as well as the same normalising flow baselines that throw away incomplete observations that were used in the conditional density estimation setting. There do exist normalising flow approaches that jointly optimises conditional distributions for missing data and

SNEFY without marginalisation

SNEFY without marginalisation

SNEFY with marginalisation

SNEFY with marginalisation

NF-4

NF-4 -

NF-2

NF-2

NF-1

NF-1

TIN

TTN

ISL

Probability of present column in training batch

OU

00

Probability of present column in training batch

Figure 2: Density estimation under partial observations. The right plot is a zoomed in version of the left plot. NF1, NF-2 and NF-4 are respectively normalising flows of depth 1, 2 and 4. Normalising flow models and SNEFY without marginalisation discard incomplete observations, whereas SNEFY use Theorem 2 to include partial observations via maximum marginal likelihood. Marginal likelihoods allow for an improved NLL.

model parameters [3] or samples from missing data and optimises for model parameters [41], however these do not allow for maximum likelihood estimation. Our observations are twofold: including partial observations improves performance, and adding more complete observations improves performance.

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 Wefocus here on the problem of density estimation, for which SNEFY is well-suited. While it is sometimes possible to obtain exact SNEFY samples using rejection sampling, sampling is more computationally expensive than in other models such as normalising flows. Future work will focus on sampling, as has been done with related models [26], where the special case of Gaussian σ and hyperrectangular support X is considered. In [26], r approximate samples with O ( r log 2 | X | + rd log 2 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, we may use any number of layers and any architecture for the conditioning network t 2 . Finally, SNEFY inherits all the usual limitations and advantages over mirroring kernel-based approaches, as discussed in § 3.5.

Future work We see a number of further promising future research directions. First, as we detail in Appendix F, choosing σ ( · ) = exp( i · ) and identity sufficient statistics results in a kernel k exp( i · ) , Id ,µ which is the Fourier transform of a nonnegative measure. By Bochner's theorem, the kernel is guaranteed to be (real or complex-valued) 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 O ( n 2 L ) integration terms for the normalising constant. The individual terms are tractable in certain cases, for example when σ 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 a squared neural network and compute the intensity function in closed-form, offering a model that scales quadratically in the number of neurons O ( n 2 ) instead of comparable models which scale cubically in the number of datapoints O ( N 3 ) [9, 51].

Acknowledgments and Disclosure of Funding

Russell and Cheng Soon would like to 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

Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research , 18(19):1-53, 2017.

R 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 of the Royal Astronomical Society , 468(4):4323-4339, 2017.

Edgar A Bernal. Training deep normalizing flow models in highly incomplete data scenarios with prior regularization. arXiv preprint arXiv:2104.01482 , 2021.

Youngmin Cho and Lawrence K. Saul. Kernel methods for deep learning. In Advances in Neural Information Processing Systems , pages 342-350, 2009.

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.

Marc Peter Deisenroth, A Aldo Faisal, and Cheng Soon Ong. Mathematics for machine learning . Cambridge University Press, 2020.

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations , 2017.

Bradley Efron and Trevor Hastie. Computer Age Statistical Inference: Algorithms, Evidence, and Data Science , volume 6. Cambridge University Press, 2021.

Seth Flaxman, Yee Whye Teh, and Dino Sejdinovic. Poisson intensity estimation with reproducing kernels. In Artificial Intelligence and Statistics , pages 270-279. PMLR, 2017.

Thorsten Glüsenkamp. Unifying supervised learning and vaes-automating statistical inference in (astro-) particle physics with amortized conditional normalizing flows. arXiv e-prints , 2020.

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.

Insu Han, Amir Zandieh, Jaehoon Lee, Roman Novak, Lechao Xiao, and Amin Karbasi. Fast neural kernel embeddings for general activations. In Advances in Neural Information Processing Systems , volume 35, pages 35657-35671, 2022.

Insu Han, Amir Zandieh, Jaehoon Lee, Roman Novak, Lechao Xiao, and Amin Karbasi. Fast neural kernel embeddings for general activations. Advances in neural information processing systems , 2022.

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units. arXiv preprint arXiv:1606.08415 , 2016.

Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation , 14(8):1771-1800, 2002.

Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems , pages 8571-8580, 2018.

Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114 , 2013.

Thomas Joseph Lawrence. Point pattern analysis on a sphere. Master's thesis, The University of Western Australia, 2018.

Nicolas Le Roux and Yoshua Bengio. Continuous neural networks. In Artificial Intelligence and Statistics , pages 404-411, 2007.

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energybased learning. Predicting structured data , 1(0), 2006.

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.

Lorenzo Loconte, Stefan Mengel, Nicolas Gillis, and Antonio Vergari. Negative mixture models via squaring: Representation and learning. In The 6th Workshop on Tractable Probabilistic Modeling , 2023.

David JC MacKay. Introduction to Gaussian processes, 1998.

Ulysse Marteau-Ferey, Francis Bach, and Alessandro Rudi. Non-parametric models for nonnegative functions. Advances in neural information processing systems , 33:12816-12826, 2020.

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-2861. PMLR, 2022.

P. McCullagh and J.A. Nelder. Generalized Linear Models, Second Edition . Monographs on Statistics and Applied Probability Series. Chapman & Hall, 1989.

Lassi Meronen, Christabella Irwanto, and Arno Solin. Stationary activations for uncertainty calibration in deep learning. Advances in Neural Information Processing Systems , 33:23382350, 2020.

Lassi Meronen, Martin Trapp, and Arno Solin. Periodic activation functions induce stationarity. Advances in Neural Information Processing Systems , 34:1673-1685, 2021.

Charlie Nash and Conor Durkan. Autoregressive energy machines. In International Conference on Machine Learning , pages 1735-1744, 2019.

Radford M Neal. Bayesian learning for neural networks . PhD thesis, University of Toronto, 1995.

Frank Nielsen and Vincent Garcia. Statistical exponential families: A digest with flash cards. arXiv preprint arXiv:0911.4863 , 2009.

Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are Gaussian processes. In The International Conference on Learning Representations , 2019.

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.

Georgios Papamakarios. Neural density estimation and likelihood-free inference . PhD thesis, University of Edinburgh, 2019.

Jooyoung Park and Irwin W Sandberg. Universal approximation using radial-basis-function networks. Neural computation , 3(2):246-257, 1991.

Tim Pearce, Russell Tsuchida, Mohamed Zaki, Alexandra Brintrup, and Andy Neely. Expressive priors in Bayesian neural networks: Kernel combinations and periodic functions. In Uncertainty in Artificial Intelligence , 2019.

Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems , 20, 2007.

Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2018.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning , pages 1530-1538, 2015.

Trevor W Richardson, Wencheng Wu, Lei Lin, Beilei Xu, and Edgar A Bernal. Mcflow: Monte carlo flow models for data imputation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14205-14214, 2020.

Alessandro Rudi and Carlo Ciliberto. PSD representations for effective probability models. In Advances in Neural Information Processing Systems , volume 34, pages 19411-19422, 2021.

Aleksanteri Sladek. Positive Semi-Definite Probabilistic Circuits. Master's thesis, Aalto University. School of Science, 2023.

W. Steinicke. Revised new general catalogue and index catalogue (revised ngc/ic). http: //www.klima-luft.de/steinicke/index_e.htm . Accessed 2nd May 2015.

Vincent Stimper, David Liu, Andrew Campbell, Vincent Berenz, Lukas Ryll, Bernhard Schölkopf, and José Miguel Hernández-Lobato. normflows: A PyTorch package for normalizing flows. arXiv preprint arXiv:2302.12014 , 2023.

Vincent Stimper, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Resampling Base Distributions of Normalizing Flows. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151, pages 4915-4936, 2022.

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.

Russell Tsuchida, Fred Roosta, and Marcus Gallagher. Invariance of weight distributions in rectified MLPs. In International Conference on Machine Learning , pages 5002-5011, 2018.

Russell Tsuchida, Fred Roosta, and Marcus Gallagher. Richer priors for infinitely wide multilayer perceptrons. arXiv preprint arXiv:1911.12927 , 2019.

Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning , 1(1-2):1-305, 2008.

Christian J Walder and Adrian N Bishop. Fast bayesian intensity estimation for the permanental process. In International Conference on Machine Learning , pages 3579-3588. PMLR, 2017.

Christopher KI Williams. Computing with infinite networks. In Advances in neural information processing systems , pages 295-301, 1997.

Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning . MIT press Cambridge, MA, 2006.

Christina Winkler, Daniel Worrall, Emiel Hoogeboom, and Max Welling. Learning likelihoods with conditional normalizing flows, 2019.

Liu Ziyin, Tilman Hartwig, and Masahito Ueda. Neural networks fail to learn periodic functions and how to fix it. In Advances in Neural Information Processing Systems , volume 33, pages 1583-1594, 2020.

A Proofs

Identity 1. The integral (2) admits a representation of the form

where k σ, t ,µ is as defined in (5) , and K Θ is the PSD matrix whose ij th entry is k σ, t ,µ ( θ i , θ j ) .

Proof. Let ˜ K Θ ( x ) be the PSD matrix whose ij th entry is σ ( w ⊤ i t ( x ) + b i ) σ ( w ⊤ j t ( x ) + b j ) ⊤ . The squared norm of the neural network evaluation is given by Tr ( V ⊤ V ˜ K Θ ( x ) ) , since

where v · ,j 1 denotes the j 1 th column of V and ⟨· , ·⟩ F denotes the Frobenius inner product. Therefore using the definition (2) directly and the linearity of the Frobenius inner product, the normalising constant is

Proposition 2. Let σ ( u ) = exp( u/ 2) and define the log-normalising constant as Ψ = log z( V , Θ ) .

Proof. The result follows by noticing that the logarithmic derivative property holds, ∑ n i =1 ∂ Ψ ∂ w i = 1 z ∑ n i =1 ∂ z ∂ w i , and that by writing

we obtain

̸

Now summing across i gives ∑ n i =1 ∂ Ψ ∂ w i = ∫ t ( x ) P ( d x ) as required.

To obtain the second result, we apply the product rule to find

and note that

Thus,

as required.

Remark 2. 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 E [ t ( x )] proceeds in the same way as in the exponential family case. The answer is positive but with two caveats. First, the log-likelihood 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, so the MLE estimation of the mean parameters may not constitute a meaningful task in SNEFY modelling (unless we are in the case where precisely the expectation of t ( x ) under the SNEFY model is of interest).

Corollary 1. Given a dataset { x ℓ } N ℓ =1 , and a SNEFY model with σ ( u ) = exp( 1 2 u ) , 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 as m = E [ t ( x )] . Then the maximum likelihood estimate of m is ˆ m = 1 N ∑ N ℓ =1 t ( x ℓ ) .

Proof. Since MLE is achieved at a stationary point of the log-likelihood, the proof follows by writing the log likelihood as

and concluding that at the MLE { w ∗ i } n i =1 for W , we must have

But

so the result follows by summing (14) over i and dividing by N . Note that maximum likelihood estimates are invariant to transformations, even if the transformation is not bijective. So if { w ∗ i } n i =1 is an MLE, we may construct a mapping from { w ∗ i } n i =1 to a corresponding MLE ˆ m for the mean parameter.

̸

̸

̸

Theorem 1. Let x = ( x 1 , x 2 ) be jointly SNEFY X 1 × X 2 , t ,σ,µ with parameters V and Θ = ([ W 1 , W 2 ] , b ) . Assume that µ ( d x ) = µ 1 ( d x 1 ) µ 2 ( d x 2 ) and t ( x ) = ( t 1 ( x 1 ) , t 2 ( x 2 ) ) . Then the conditional distribution of x 1 given x 2 = x 2 is SNEFY X 1 , t 1 ,σ,µ 1 with parameters V and Θ 1 | 2 ≜ ( W 1 , W 2 t 2 ( x 2 ) + b ) .

Proof. The joint distribution of x satisfies

Therefore, the distribution of x 1 conditionally on x 2 = x 2 , which is obtained by dividing the joint distribution by the marginal distribution of x 2 (which is independent of x 1 ), satisfies

That is, the term W 2 t 2 ( x 2 ) + 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 x 1 .

Theorem 2. Let x = ( x 1 , x 2 ) be jointly SNEFY X 1 × X 2 , t ,σ,µ with parameters V and Θ = ([ W 1 , W 2 ] , b ) . Assume that µ ( d x ) = µ 1 ( d x 1 ) µ 2 ( d x 2 ) and t ( x ) = ( t 1 ( x 1 ) , t 2 ( x 2 ) ) . Then the marginal distribution of x 1 is

Proof. The marginal distribution of the random variable x 1 is obtained by marginalising out the joint distribution with respect to x 2 ,

The integral z 2 takes a similar form to z( V , Θ ) ,

B Derivation of neural network kernels

Kernel 1. k σ, Id , Φ C , m ( θ i , θ j ) = k σ, Id , Φ ( T θ i , T θ j ) , where T Θ = ( WA , b + Wm ) , T θ i = ( w ⊤ i A , b i + w ⊤ i m ) and A is a matrix factor such that covariance C = AA ⊤ .

Proof. The NNK may be expressed as an expectation with respect to a Gaussian random variable x with mean m and covariance matrix C . It holds that x d = A z + 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. First observe that the expected value of the cosine of a Gaussian random variable can be evaluated by equating the real and imaginary components of the characteristic function of a Gaussian random variable and the expected value of Euler's form. That is, if z is Gaussian with mean m and variance v 2 ,

With this identity at hand, we proceed by direct evaluation of (4).

Proof. This is immediate from the expected value of a product of two correlated Gaussians, w ⊤ i x + b i and w ⊤ j x + b j .

Kernel 5. The kernel k Snake a ( · ) -1 2 a , Id , Φ ( θ i , θ j ) is equal to

Proof. Choosing σ = Snake a ( · ) -1 2 a in (4), and expanding the resulting quadratic,

We now evaluate A and B . The terms A and B obey a symmetry, so it suffices to evaluate term A . Term A can be evaluated using Stein's lemma,

Assembling all the known individual terms in (15),

Kernel 6. The kernel k Snake a , Id , Φ ( θ i , θ j ) is equal to

Proof. Choose σ = Snake a and note that Snake a ( · ) = ( Snake a ( · ) -1 2 a ) + 1 2 a . The Kernel 5 corresponds with the case ( Snake a ( · ) -1 2 a ) , so we are left with three additional terms. These terms may be evaluated directly,

C Examples which generalise standard exponential family models

In this section, we will study examples of the SNEFY model which use activation function σ ( u ) = exp( u/ 2) (or equivalently, up to scaling, σ ( u ) = exp( u ) ) and as such correspond to a 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 a tractable normalising constant and we can write the kernels directly using Proposition 1.

SNEFY Von Mises-Fisher mixtures. The VMF distribution is a helpful way of defining the notion of a Gaussian distribution to the sphere. The following kernel may be used to define a VMF distribution. Alternatively, it may be viewed as a way of constructing a distribution supported on R d with sufficient statistics which are projected onto the sphere.

Kernel 3. Define proj S d -1 ( x ) ≜ x / ∥ x ∥ to be the projection onto the unit sphere. Then

where I p is the modified Bessel function of the first kind of order p . In the special case d = 3 , we have the closed-form k exp , proj S 2 , Φ ( θ i , θ j ) = exp( b i + b j ) ( e ∥ w i + w j ∥ -e -∥ w i + w j ∥ ) 2 ∥ w i + w j ∥ .

Proof. If x ∼ N (0 , I ) , then x / ∥ x ∥ is uniformly distributed on the sphere. From the normalizing constant of the von Mises-Fisher distribution, from Proposition 1, it then follows that

where I p is the modified Bessel function of the first kind of order p . In the special case of p = 1 / 2 , we have I 1 / 2 ( z ) = √ 2 πz sinh( z ) = ( exp( z ) -exp( -z ) ) √ 1 2 πz . This implies that when d = 3 , since Γ(3 / 2) = √ π 2 ,

Note that k exp , proj S d -1 , Φ ( θ i , θ j ) = k exp , Id ,ν ( θ i , θ j ) , where ν is the uniform measure on the sphere S d -1 , because if x is Gaussian then x / ∥ x ∥ is uniform on the sphere. This allows one to construct SNEFY exp , Id ,ν 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.

Proof. This is a consequence of the moment generating function of the multivariate Gaussian distribution. More concretely, by Proposition 1,

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 is large or infinite, for which computing normalising constants can naively be challenging even in the discrete setting. Let X = { 0 , 1 , 2 , . . . } , t ( x ) = x , and the base measure µ ( dx ) = 1 /x ! ν ( dx ) , 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

following the usual mean parametrisation λ i = e w i , so the individual mixture components have rates which are geometric means of ( λ i , λ j ) pairs.

Kernel 8. Choose the base measure µ ( dx ) = ( x !) -1 ν ( dx ) , where ν is the counting measure. We have

Proof. This is again direct from Proposition 1. In detail,

The second factor involving the sum is the partition function of the Poisson distribution in canonical form, which is exp ( exp( w i + w j ) ) .

The usual mean parameterisation of the Poisson distribution is through a rate parameter λ i = exp( w i ) , which would lead to the kernel representation

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 = R d , t ( x ) = ( x 1 , . . . , x d , x 2 1 , . . . , x 2 d ) (i.e. D = 2 d ) and suppose µ ( d x ) = d x is Lebesgue measure. Choose σ ( · ) = exp( · / 2) and consider the r -th output of our network f : R D → R m

where we denoted ˜ w iℓ = w i,d + ℓ . In this case, we require that ˜ w eℓ < 0 for the model to be (square) integrable. Reparametrising σ 2 iℓ = -1 2 ˜ w iℓ and µ iℓ = -w iℓ 2 ˜ w iℓ and absorbing the factor exp ( -∑ d ℓ =1 µ 2 iℓ 4 σ 2 iℓ ) into readout parameters V , gives

Thus, we have recovered a 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 the RBF network, we can restrict the parameters σ 2 iℓ = σ 2 i to be the same across the dimensions, giving

with location parameters µ i ∈ R d and the scale parameters σ 2 i > 0 . Note the unusual factor of 4 in front of σ 2 i - this ensures that our model in 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 a density model with respect to the Lebesgue measure as

If n = 1 , we recover simply a multivariate normal density N ( µ 1 , σ 2 1 I ) . The above model is essentially the same as the one in [42], despite being derived in a very different way.

Kernel 9. Let t (2) ( x ) = ( x 1 , . . . , x d , x 2 1 , . . . , x 2 d ) so that D = 2 d be the sufficient statistic. Partition W = [ W [: , 1: d ] , ˜ W ] and suppose ˜ W < 0 element-wise. Choose µ ( d x ) = d x to be the Lebesgue measure. Then

Proof. As with the kernels above, this follows from Proposition 1, since

While Proposition 1 gives us an expression for the kernel matrix K Θ in terms of natural parameters, we can also express it directly in terms of parameters µ i , σ 2 i . In particular,

We briefly state two more cases without an extended discussion.

SNEFY Gamma mixtures.

Kernel 10. Let X = (0 , ∞ ) , t ( x ) = (log x, -x ) and σ = exp . Partition W = [ W [: , 1: d ] , ˜ W ] and suppose W [: , 1: d ] > -1 and ˜ W > 0 element-wise. Choose µ ( d x ) = d x to be the Lebesgue measure. Then

SNEFY Dirichlet mixtures.

Kernel 11. Let X = ∆ D -1 , a ( D -1) -simplex of probability distributions, i.e. x ∈ [0 , 1] D , ∑ D i =1 x i = 1 . Let σ = exp . Let t ( x ) = (log x 1 , . . . , log x D ) . Choose µ ( d x ) = d x to be the Lebesgue measure. Suppose W > -1 elementwise. Then

D Marginalisation in the case σ = exp( · / 2)

Let M be a positive semi-definite n × n matrix. We will make use of the following SNEFY parametrisation

Since we can always write M = V ⊤ V , for an m × n matrix V , m ≤ n , we have

which is, as in the parametrisation given in the main text, simply the squared Euclidean norm of a multi-output neural network, with V corresponding to the weights of the second layer. If we denote by v · j the j -th column of V , the normalizing constant is given by

where as before

Now if we let σ = exp( · / 2) we obtain a family which is also closed under marginalisation (in addition to conditioning). The Proposition below generalises Proposition 1 of [42], which considers the special case of the Gaussian PSD mixtures.

Proposition 3. Let x = ( x 1 , x 2 ) be jointly SNEFY X 1 × X 2 , t , exp( · / 2) ,µ with parameters V and Θ = ([ W 1 , W 2 ] , b ) . Assume that µ ( d x ) = µ 1 ( d x 1 ) µ 2 ( d x 2 ) and t ( x ) = ( t 1 ( x 1 ) , t 2 ( x 2 ) ) . Then the marginal distribution of x 1 is SNEFY X 1 , t 1 , exp( · / 2) ,µ 1 with parameters ˜ V and Θ = ( W 1 , b ) , for some matrix ˜ V ∈ R m × n .

Proof. Proposition 1 gives us the normalising constant for the parametrisation where biases are absorbed into V . If we explicitly keep the biases in the parametrisation, we have

where z e is the normalizing constant of the exponential family with the sufficient statistic t and base measure µ . By Theorem 2, we have that

where [ C Θ ( x 1 )] ij = k σ, t 2 ,µ 2 ( ( w 2 i , w ⊤ 1 i t 1 ( x 1 ) + b i ) , ( w 2 j , w ⊤ 1 j t 1 ( x 1 ) + b j ) ) . But now since σ = exp( · / 2) , applying (17) gives

where z e, 2 is the normalizing constant of the exponential family with the sufficient statistic t 2 and base measure µ 2 . Thus, we can write

and we conclude that the marginal is in the same family with ˜ M = M ◦ Z e, 2 , where

Note that ˜ M is PSD as an Hadamard product of two PSD matrices. Thus, we can find ˜ V such that ˜ M = ˜ V ⊤ ˜ V .

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 1 fixed seed. We find that sometimes the variance over random seeds can be large (e.g. Resampled 0 on Circles). Secondly, rather than computing test performance at the last epoch of training, we follow the more standard procedure of returning the test performance evaluated at the epoch corresponding with the smallest validation performance. This validation/test monitoring results in substantial performance gains in all of the models, with no extra computational cost for SNEFY , Gauss and GMM . For Resampled , monitoring the validation performance to a high precision requires estimating the normalising constant to a high precision, which is computationally challenging. We therefore only check the validation performance every 100 epochs, and compute a high precision normalising constant if the validation performance is the lowest encountered so far. We train each models for a maximum of 20000 iterations (while monitoring validation performance). We use Adam with default hyperparameters and weight decay 10 -3 . The batch size is 2 10 .

We use a SNEFY with Gaussian mixture model base measure supported on X = R 2 , identity sufficient statistic t and activation function cos . The base measure consists of 8 mixture components, each with a diagonal covariance matrix. We use the same Resampled architecture as in the original paper [46]. We use an MLP with layer widths [2 , 256 , 256 , 1] and sigmoid activations for the resampling distribution. We use a discount factor of 0 . 1 for the exponential moving average partition function calculation. Note that for Resampled models we are only able to provide an estimate of the test log likelihood, not the exact log likelihood as in other methods. We choose n = 50 and m = 1 . The Gaussian mixture model has 10 mixture components, each with a diagonal covariance matrix. Each normalising flow block consists of an affine coupling block, a permutation layer, and an actnorm layer.

E.2 Modelling distributions on the sphere

We compare the performance of a VMF distribution, a 'regular' VMF mixture distribution and our SNEFY construction of the VMF mixture, which allows for some negative weight coefficients as discussed in § 3.2. We use the dataset [44] also used by [19] in a different context for point

processes. This dataset, retrieved in 2015, is called the 'Revised New General Catalogue and Index Catalogue' (RNGC/IC). The RNGC/IC consists of locations of some 10610 galaxies. We use the spherical coordinates of these galaxies and map them to the surface of a sphere.

We compare two types of SNEFY S 2 , exp , Id ,d x with m = n = 30 : one is constrained so that V is diagonal (and therefore V ⊤ V is diagonal with nonnegative entries), and the other uses a unconstrained general V . We expect the unconstrained model to be more expressive, and therefore obtain a better test NLL.

We use Adam with default hyperparameters and a full batch size. We randomly shuffle the data and perform an 80/20 train/test split. Each of the 50 runs uses a randomly sampled initialisation and train/test split. We train for 20000 epochs.

E.3 Conditional density estimation of photometric redshift

Our deep conditional feature extractor is an MLP that uses ReLU activations in each layer and batch normalisation between each layer. Our SNEFY model uses Snake a activations with a = 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 a batch size of 256 . The dataset consists of 74309 training and 74557 examples, which are fixed between each run. Each run uses a randomly sampled initialisation (except for CKDE, which is deterministic).

F Complex activation case

Here we consider an extension where we allow the neural network activation σ to be complex-valued, and accordingly the readout parameters V to be complex, i.e. V ∈ C m × n . Note that in this case

Then the above takes the form

As in the σ = exp case, bias terms can be folded into the readout parameters V (which is the reason why we also require readout parameters to be complex). We note that the case n = 1 is not interesting as it simply reduces to the (normalised) base measure µ .

In order to obtain the normalizing constant Tr( V H V K Θ ) , we need the integral of the form

where κ is simply the Fourier transform of the base measure µ , i.e. its characteristic function in the case where µ is a probability measure. Hence many standard choices of µ lead to tractable normalizing constants. Note that while V and κ both may be complex-valued, the normalizing constant as well as the density itself are real-valued.

Various examples of probability measures µ and its Fourier transforms that give rise to shift-invariant positive definite kernels κ have been studied in the literature on Random Fourier Features (RFF) [38]. Here too, while the functional form of the expressions is identical, the role between the data and the parameters is reversed: in RFF, one is interested in approximating a given kernel on the data instances, by considering a probability measure on the frequency space which is that kernel's inverse Fourier transform.

Remark 3. If we restrict our attention to real-valued V and thus explicitly maintain biases inside the parameterisation, this model can also be realised with stacking cos and sin activation so that they share the readout parameters since

G Discrete and mixed continuous/discrete SNEFY s

G.1 Discrete SNEFY s via series

Closed-form expressions for normalising constants of discrete distributions are only advantageous if their computation costs less than naively summing up unnormalised density over all possible states. This is the case if the support has very large or infinite cardinality. An example that extends the classical Poisson distribution in exponential family form is given in Table 1. Here we discuss some other settings.

NNK as a convergent series Suppose the support X is discrete and let h ( x ) be a nonnegative function corresponding with the discrete base measure µ . The NNK is given by

Such NNKs often resemble known convergent series.

Fourier series For example, if σ = cos and C x = h ( x ) are some coefficients of a convergent Fourier series, then

is just a series representation of a sum of two periodic functions. For example, if C x = 2( -1) x +1 πx for x ≥ 1 and C x = 0 otherwise, and t ( x ) = 2 πx , the NNK is a sum of two sawtooth waves with frequencies w 1 -w 2 and w 1 + w 2 and phase offsets b 1 -b 2 and b 1 + b 2 . Other examples include rectified sine waves, square waves and triangular waves. This extends to periodic functions with convergent multivariate Fourier series.

G.2 Mixed continuous/discrete SNEFY s

Mixed distributions can be obtained by choosing the base measure to be a mixed continuous distribution. For example, choose X = R d , µ ( d x ) = 1 2 ( Φ( x ) + δ ( x ) ) d x , and take σ and t to be any of the combinations leading to a closed-form NNGPK (for example, t as the identity and σ as the error function, Leaky ReLU, GELU, cosine, or snake). Then

which is a closed form.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\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).\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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nAs 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\nWEAKNESSES:\nAs 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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nThe main methodological limitation is the intractability of sampling, which the authors emphasize.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nThe 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).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nAccess 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n--\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}