It's not clear from the formulas how x=[x1,...,xN] relates to y, μ, and σ since these are defined without x. Assuming y = Wx + b, and μ, σ, and π are learnable parameters for each output dimension. The symbol π seems to mean both weight and the constant 3.14159 in the same formula.
Overall it looks similar to radial basis activations, but the activations look to be log of weighted "stochastic" sums (weights sum to one) of a set of radial basis functions.
Clearly I wasn't in neural net mode. I take it then the learned parameters, the means, variances and mixing coefficients, are effectively functions of the output of the previous layer.
Thanks - That's correct, the Gaussian mixture parameters (mu, sigma, pi) are learned as functions of the input from the previous layer. So it’s still a feedforward net: the activations from layer x determine the mixture parameters for the next layer.
The reason the neuron’s output is written as a log-density Pj(y) is just to emphasize the probabilistic view: each neuron is modeling how likely a latent variable y would be under its mixture distribution.
Thank you for your work! I would be interested to see what this means to a CNN architecture. Maybe it wouldn't actually be needed to have the whole architecture based on uGMM-NNs but only the last layers?
Thanks - good question, in theory, the uGMM layer could complement CNNs in different ways - for example, one could imagine (as you mentioned):
using standard convolutional layers for feature extraction,
then replacing the final dense layers with uGMM neurons to enable probabilistic inference and uncertainty modeling on top of the learned features.
My current focus, however, is exploring how uGMMs translate into Transformer architectures, which could open up interesting possibilities for probabilistic reasoning in attention-based models.
uGMM-NN is a novel neural architecture that embeds probabilistic reasoning directly into the computational units of deep networks. Unlike traditional neurons, which apply weighted sums followed by fixed nonlinearities, each uGMM-NN node parameterizes its activations as a univariate Gaussian mixture, with learnable means, variances, and mixing coefficients.
Upshot: Gaussian sampling along the parameters of nodes rather than a fixed number. This might offer one of the following:
* Better inference time accuracy on average
* Faster convergence during training
It probably costs additional inference and training compute.
The paper demonstrates worse results on MNIST, and shows the architecture is more than capable of dealing with the Iris test (which I hadn’t heard of; categorizing types of irises, I presume the flower, but maybe the eye?)
The paper claims to keep the number of parameters and depth the same, but it doesn’t report as to
* training time/flops (probably more I’d guess?)
* inference time/flops (almost certainly more)
Intuitively if you’ve got a mean, variance and mix coefficient, then you have triple the data space per parameter — no word as to whether the networks were normalized as to total data taken by the NN or just the number of “parameters”.
Upshot - I don’t think this paper demonstrates any sort of benefit here or elucidates the tradeoffs.
Quick reminder, negative results are good, too. I’d almost rather see the paper framed that way.
Thanks for the comment. Just to clarify, the uGMM-NN isn't simply "Gaussian sampling along the parameters of nodes."
Each neuron is a univariate Gaussian mixture with learnable mean, variance, and mixture weights. This gives the network the ability to perform probabilistic inference natively inside its architecture, rather than approximating uncertainty after the fact.
The work isn’t framed as "replacing MLPs." The motivation is to bridge two research traditions:
- probabilistic graphical models and probabilistic circuits (relatively newer)
- deep learning architectures
That's why the Iris dataset (despite being simple) was included - not as a discriminative benchmark, but to show the model could be trained generatively in a way similar to PGMs, something a standard MLP cannot do. Hence, the other benefits of the approach mentioned in the paper.
Thanks for writing back! I appreciate the plan to integrate the two architectures. On that front, it might be interesting to have a future research section - like what would be uniquely good about this architecture if scaled up?
On ‘usefulness’ I think I’m still at my original question - it seems like an open theoretical q to say that the combination of a tripled-or-greater training budget, data size budget of the NN, and probably a close to triple or greater inference budget, the costs of the architecture you described, cannot be closely approximated by the “fair equivalent”-ly sized MLP.
I hear you that the architecture can do more, but can you talk about this fair size question I have? That is, if a PGM of the same size as your original network in terms of weights and depth is as effective, then we’d still have a space savings to just have the two networks (MLP and PGM) side by side.
That’s a fair question. You’re right that on paper a uGMM neuron looks like it “costs” ~3× an MLP weight. But there are levers to balance that. For example, the paper discusses parameter tying, where the Gaussian component means are tied directly to the input activations. In that setup, each neuron only learns the mixture weights and variances, which cuts parameters significantly while still preserving probabilistic inference. The tradeoff may be reduced expressiveness, but it shows the model doesn’t have to be 3x heavier.
More broadly: traditional graphical models were largely intractable at deep learning scale until probabilistic circuits, which introduced tractable probabilistic semantics without exploding parameter counts. Circuits do this by constraining model structure. uGMM-NN sits differently: it brings probabilistic reasoning inside dense architectures.
So while compute cost is real, the “fair comparison” isn’t just params-per-weight, it’s also about what kinds of inference the model can do at all, and the added interpretability of mixture-based neurons, which traditional MLP neurons don’t provide - it shares some spirit with recent work like KAN, but tackles the problem through probabilistic modeling rather than spline-based function fitting.
It's not clear from the formulas how x=[x1,...,xN] relates to y, μ, and σ since these are defined without x. Assuming y = Wx + b, and μ, σ, and π are learnable parameters for each output dimension. The symbol π seems to mean both weight and the constant 3.14159 in the same formula.
Overall it looks similar to radial basis activations, but the activations look to be log of weighted "stochastic" sums (weights sum to one) of a set of radial basis functions.
The biggest difference is probably log outputs.
I'm having a very dense moment I think, and it's been far to long since the statistics courses.
They state the output of a neuron j is a log density P_j(y), where y is a latent variable.
But how does the output from the previous layer, x, come into play?
I guess I was expecting some kind of conditional probabilities, ie the output is P_j given x or something.
Again, perhaps trivial. Just struggling to figure out how it works in practice.
Clearly I wasn't in neural net mode. I take it then the learned parameters, the means, variances and mixing coefficients, are effectively functions of the output of the previous layer.
Thanks - That's correct, the Gaussian mixture parameters (mu, sigma, pi) are learned as functions of the input from the previous layer. So it’s still a feedforward net: the activations from layer x determine the mixture parameters for the next layer.
The reason the neuron’s output is written as a log-density Pj(y) is just to emphasize the probabilistic view: each neuron is modeling how likely a latent variable y would be under its mixture distribution.
Thank you for your work! I would be interested to see what this means to a CNN architecture. Maybe it wouldn't actually be needed to have the whole architecture based on uGMM-NNs but only the last layers?
Thanks - good question, in theory, the uGMM layer could complement CNNs in different ways - for example, one could imagine (as you mentioned):
using standard convolutional layers for feature extraction,
then replacing the final dense layers with uGMM neurons to enable probabilistic inference and uncertainty modeling on top of the learned features.
My current focus, however, is exploring how uGMMs translate into Transformer architectures, which could open up interesting possibilities for probabilistic reasoning in attention-based models.
uGMM-NN is a novel neural architecture that embeds probabilistic reasoning directly into the computational units of deep networks. Unlike traditional neurons, which apply weighted sums followed by fixed nonlinearities, each uGMM-NN node parameterizes its activations as a univariate Gaussian mixture, with learnable means, variances, and mixing coefficients.
Meh. Well, at least, possibly “meh”.
Upshot: Gaussian sampling along the parameters of nodes rather than a fixed number. This might offer one of the following:
* Better inference time accuracy on average
* Faster convergence during training
It probably costs additional inference and training compute.
The paper demonstrates worse results on MNIST, and shows the architecture is more than capable of dealing with the Iris test (which I hadn’t heard of; categorizing types of irises, I presume the flower, but maybe the eye?)
The paper claims to keep the number of parameters and depth the same, but it doesn’t report as to
* training time/flops (probably more I’d guess?)
* inference time/flops (almost certainly more)
Intuitively if you’ve got a mean, variance and mix coefficient, then you have triple the data space per parameter — no word as to whether the networks were normalized as to total data taken by the NN or just the number of “parameters”.
Upshot - I don’t think this paper demonstrates any sort of benefit here or elucidates the tradeoffs.
Quick reminder, negative results are good, too. I’d almost rather see the paper framed that way.
Thanks for the comment. Just to clarify, the uGMM-NN isn't simply "Gaussian sampling along the parameters of nodes."
Each neuron is a univariate Gaussian mixture with learnable mean, variance, and mixture weights. This gives the network the ability to perform probabilistic inference natively inside its architecture, rather than approximating uncertainty after the fact.
The work isn’t framed as "replacing MLPs." The motivation is to bridge two research traditions:
- probabilistic graphical models and probabilistic circuits (relatively newer)
- deep learning architectures
That's why the Iris dataset (despite being simple) was included - not as a discriminative benchmark, but to show the model could be trained generatively in a way similar to PGMs, something a standard MLP cannot do. Hence, the other benefits of the approach mentioned in the paper.
Thanks for writing back! I appreciate the plan to integrate the two architectures. On that front, it might be interesting to have a future research section - like what would be uniquely good about this architecture if scaled up?
On ‘usefulness’ I think I’m still at my original question - it seems like an open theoretical q to say that the combination of a tripled-or-greater training budget, data size budget of the NN, and probably a close to triple or greater inference budget, the costs of the architecture you described, cannot be closely approximated by the “fair equivalent”-ly sized MLP.
I hear you that the architecture can do more, but can you talk about this fair size question I have? That is, if a PGM of the same size as your original network in terms of weights and depth is as effective, then we’d still have a space savings to just have the two networks (MLP and PGM) side by side.
Thanks again for publishing!
That’s a fair question. You’re right that on paper a uGMM neuron looks like it “costs” ~3× an MLP weight. But there are levers to balance that. For example, the paper discusses parameter tying, where the Gaussian component means are tied directly to the input activations. In that setup, each neuron only learns the mixture weights and variances, which cuts parameters significantly while still preserving probabilistic inference. The tradeoff may be reduced expressiveness, but it shows the model doesn’t have to be 3x heavier.
More broadly: traditional graphical models were largely intractable at deep learning scale until probabilistic circuits, which introduced tractable probabilistic semantics without exploding parameter counts. Circuits do this by constraining model structure. uGMM-NN sits differently: it brings probabilistic reasoning inside dense architectures.
So while compute cost is real, the “fair comparison” isn’t just params-per-weight, it’s also about what kinds of inference the model can do at all, and the added interpretability of mixture-based neurons, which traditional MLP neurons don’t provide - it shares some spirit with recent work like KAN, but tackles the problem through probabilistic modeling rather than spline-based function fitting.