Squared Neural Families: A New Class of Tractable Density Models

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.

Paper

Similar papers

Peer review

Reviewer k4Y85/10 · confidence 3/52023-06-13

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). The 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. These 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). For 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). This 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. The model is trained by MLE and aveluated on: three 2D-dimensional datasets, a spheirical manifold and on a low-dimensional astronomy dataset.

Strengths

- new NNK are proposed - tractable marginalisation and conditioning - connection with exponential families (every exponential famility is a SNEFY) and kernel-based methods for nonnegative functions

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). - 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]. - Sampling seems to be possible only by rejection sampling, and is therefore very computationally expensive. Could you elaborate more about sampling and its challenges? - 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 [1] Nash, Charlie, and Conor Durkan. "Autoregressive energy machines." International Conference on Machine Learning. PMLR, 2019.

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)$)? - 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?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

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. But probably, more emphasis and discussion about the long training time required would be appreciated.

Reviewer zjNK6/10 · confidence 4/52023-06-29

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.

Strengths

As I see it the main strengths of this submission include: - It is very clearly written - It addresses a problem (namely defining useful classes of density estimators) that is of wide interest in the NeurIPS community - The proposed class of models is, to the best of my knowledge, novel - The proposed class of models has some attractive properties, especially w.r.t. conditioning and marginalization - The authors do a good job of placing their work in the context of related literature and methods

Weaknesses

As I see it the main limitations and weaknesses of this submission are: - The experimental evaluation is quite limited, making it difficult to judge the likely usefulness or impact of SNEFY in practice - The intractability of obtaining samples from SNEFY distributions is a serious limitation (as the authors acknowledge) To 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: > The CNF models utilise affine layers wtih tanh activations. Not 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). Granted, 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.

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? - 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. - "For example, exploited..." [Line 29] This is a run-on (i.e. incomplete) sentence. - Final line "... offering a model that scales like O(n^2) instead...". Do you mean O(n)? - Dense tables like Table 2 are difficult to parse; consider making a plot

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The main methodological limitation is the intractability of sampling, which the authors emphasize.

Reviewer Y6Kj7/10 · confidence 3/52023-07-06

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). The 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. The model is evaluated in unconditonal density modeling on synthetic data, and conditional density modeling on a redshift estimation task. EDIT: I have read the author's rebuttal, which addressed my concerns.

Strengths

The paper is very well-written and clear. The 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. It thus enjoys the tractability properties mentioned above. I 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.

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. 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 $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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

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).

Reviewer tpyH7/10 · confidence 2/52023-07-10

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.

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. The 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. The paper is very technical, but in general quite well written.

Weaknesses

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. Relatedly 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. 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.

Questions

--

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

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.

Reviewer k4Y82023-08-10

Thanks for your clarifications, and specifically for elaborating on discrete/mixed input variables. I did appreciate it. I read all reviews, and specifically the review from Reviewer zjNK, and I should say that I agree with basically everything they say: Sampling and limited evaluation (also due to long training times) are relevant weaknesses to me as well (as I had already pointed out). I like the work and will keep an eye on it but I'm still comfortable with my score.

Authorsrebuttal2023-08-13

Thanks again for your feedback, especially regarding mixed and discrete settings, and confirming you have read the reviews and discussions.

Reviewer zjNK2023-08-10

I thank the authors for their response, in particular the new experiment with partially observed data. While the empirical story is still a bit weak here, this goes some way in making the practical advantages of SNEFY more concrete. As such I will raise my score to a 6. I urge the authors to revise the way they present some of their empirical results: in particular the already-mentioned difficult-to-parse table but also the figure for the new experiment which is also difficult to parse.

Authorsrebuttal2023-08-13

Thanks for your constructive feedback on our manuscript. We are glad that you urged us to consider the partial observation setting as an experiment. Regarding presentation of empirical results, we will add a figure for the tables and put in bold numbers for the final revision.

Reviewer Y6Kj2023-08-11

Thank you for their detailed answer. - I still disagree with the claim that "in a conditional normalising flow (CNF), the mapping of each layer would have to be bijective in both $x_1$ and $x_2$". Indeed, for maximum-likelihood training, one needs to compute the determinant of $\partial_{x_1} f(x_1 | x_2)$ only, so that the network should be invertible in $x_1$ for every $x_2$. In fact, in [48], the authors "the conditioning term $x$ [(standing for $x_2$ in this paper's notations)] is transformed into a rich representation $h = g(x)$ using a large network $g$ [which does not need to be invertible]" (last paragraph of section 3.1). Am I missing something here? - I like the new missing data experiment, which showcases the usefulness of tractable marginalization in practical applications. To me, this is the strongest point of approach. I thus recommend acceptance.

Authorsrebuttal2023-08-13

Thanks for following up

You're right about point 1. The paragraph on line 295 is about bijectivity in the variable we are modelling, not the variable we are conditioning on. The last sentence of this paragraph is confusing because, as you point out, conditional normalising flows are not required to be bijective in the variable they are conditioning on. We should have been talking about the variable we are modelling, in which case normalising flows are restricted. We will update this paragraph in the final manuscript. We are glad that you like the new experiment. Thanks again for your feedback on the paper!

Reviewer tpyH2023-08-15

Thank you for your clarifications. I considered this paper worth publishing a priori, however the improvements delineated will make it more readable to help especially practitioners to see the potential of this work. The emphasis on missing data scenarios is useful in distinguishing the power of flows and SNEFY's in respective applications. Sampling remains a weakness (now more prominently stated in the paper), but this work is a significant step forward.

Program Chairsdecision2023-09-21

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC