Rethinking Nonlinearity: Trainable Gaussian Mixture Modules for Modern Neural Architectures

TL;DR

This work targets the nonlinearity bottleneck in modern neural architectures by introducing Gaussian Mixture-Inspired Nonlinear Modules (GMNM), a differentiable integration of Gaussian mixtures that can be trained with standard gradient methods. GMNM uses Augmented Gaussian Projection blocks to form a flexible, trainable nonlinear basis, relaxing probability constraints and avoiding strict positive-definite covariance requirements. Empirically, GMNM improves function approximation and boosts performance across diverse tasks including image classification, generation, time-series forecasting, and PDE solving, with often fewer parameters and no slowdowns. The results suggest GMNM as a versatile, practical mechanism to enhance expressiveness of CNNs, LSTMs, and related architectures in a broad range of applications.

Abstract

Neural networks in general, from MLPs and CNNs to attention-based Transformers, are constructed from layers of linear combinations followed by nonlinear operations such as ReLU, Sigmoid, or Softmax. Despite their strength, these conventional designs are often limited in introducing non-linearity by the choice of activation functions. In this work, we introduce Gaussian Mixture-Inspired Nonlinear Modules (GMNM), a new class of differentiable modules that draw on the universal density approximation Gaussian mixture models (GMMs) and distance properties (metric space) of Gaussian kernal. By relaxing probabilistic constraints and adopting a flexible parameterization of Gaussian projections, GMNM can be seamlessly integrated into diverse neural architectures and trained end-to-end with gradient-based methods. Our experiments demonstrate that incorporating GMNM into architectures such as MLPs, CNNs, attention mechanisms, and LSTMs consistently improves performance over standard baselines. These results highlight GMNM's potential as a powerful and flexible module for enhancing efficiency and accuracy across a wide range of machine learning applications.

Figures (5)

• Figure 1: Architecture of the proposed GMNM. The input $\mathbf{x} = [x_1, x_2, \dots, x_d]$ is processed by $m$ parallel Augmented Gaussian Projection (AGP) modules. Each AGP applies centering $\mathbf{z} = \mathbf{x} - \boldsymbol{\mu}$ and then uses multiple linear projections (LP) to form an aggregated scalar output $y$, which is mapped through $f(y) = e^{-0.5 \, y^2}$. The outputs $f_i(y)$ from individual AGP modules are linearly combined with weights $\{\pi_i\}$ to produce the final mixture output $G(\mathbf{x})$.
• Figure 2: 2D function fitting results for progressively more complex tarrfaces (left), training losses (middle), and test losses (right). Again, GMNM (red) outperforms KAN (green) and MLPs (blue and orange).
• Figure 3: Poisson equation results. Top: total L2 error over training steps. Middle: boundary MSE. Bottom: PDE residual MSE. Lower curves indicate better performance.
• Figure 4: Architecture for Classification Tasks and Generation Tasks. Top (Classification): The green block shows the optional attention layer, while the yellow blocks denote GMNM modules that model channel-wise interactions. Bottom (Generation): Architecture follow auto-encoder frame work with convolution down sample and convolution transpose up sample. Image down sample into a latent vector via the encoder. The latent features also separately feed to the GMNM module. The outputs GMNM module element-wise multiplication operation with decoder output. These fused features are passed through a output layer(convolution) for the final images reconstructed.
• Figure 5: L2 loss for Cifar100 image generation experiment. Top: training loss; Bottom: test loss.