TLDR: This research investigates how the distributions of weight matrices in neural networks change during training, moving beyond simple initial Gaussian assumptions. It finds that 13-parameter permutation-invariant Gaussian matrix models (PIGMMs) effectively capture the correlated Gaussianity in these weight matrices throughout the training process for MNIST classification, even when simple Gaussian models fail. The study also explores the effects of regularization and increasing layer width, demonstrating the robustness and interpretability of PIGMMs as a framework for understanding neural network parameter spaces.
Neural networks, the backbone of modern artificial intelligence, are incredibly powerful tools capable of approximating complex functions. However, their success often comes with a challenge: they are typically vastly over-parameterized, making it difficult to understand exactly how they arrive at their solutions. Each time a neural network is initialized and trained, its final set of parameters can vary significantly, leading to a need for better ways to interpret and model these internal workings.
This research delves into the fascinating world of neural network weight matrices, specifically examining how their distributions evolve throughout the training process. Moving beyond the common assumption that weights remain simple, independently distributed Gaussians, the study explores a more sophisticated approach using “matrix models” grounded in “permutation symmetry.”
The core of this investigation lies in Permutation Invariant Gaussian Matrix Models (PIGMMs). These models, characterized by 13 parameters, are designed to capture more complex, correlated Gaussian patterns within the weight matrices. Unlike simpler Gaussian models where each weight is treated as an independent variable, PIGMMs account for the intricate relationships and symmetries that emerge as a network learns.
To test the effectiveness of PIGMMs, the researchers focused on the MNIST classification problem, a standard benchmark in machine learning involving handwritten digit recognition. They trained neural networks with three hidden layers, each containing 10 neurons, generating square weight matrices. The training involved different initialization schemes (Gaussian and Uniform) and was repeated 1000 times to gather robust statistical data across 50 epochs.
A key finding was that while simple Gaussian models adequately describe weight distributions at initialization, they quickly become poor fits as training progresses. This is where PIGMMs shine. The study demonstrated that these more general 13-parameter models effectively represent the correlated Gaussianity in the weight matrices, not just at the initial setup but consistently throughout the entire training journey. This suggests that even as a network learns and its weights adjust to solve a problem, the underlying statistical patterns can still be described by a generalized Gaussian framework, albeit one that accounts for complex correlations.
To quantify these changes, the researchers calculated the Wasserstein distance, a metric that measures the distance between probability distributions. This allowed them to track how the fitted PIGMMs moved in their parameter space, smoothly diverging from the initial simple Gaussian models as training advanced. They observed that deeper layers and uniform initialization schemes led to greater deviations from the simple Gaussian, indicating more significant changes in weight distributions.
The research also explored the impact of architectural modifications. When L2 regularization was introduced, PIGMMs remained good models, though some specific higher-order invariants showed increased deviations, hinting at how regularization might subtly break certain symmetries. In the context of increasing layer width (from 10 to 640 neurons), the study found that while PIGMMs became less universally representative for very wide networks, they still provided valuable insights into where the Gaussianity assumption was most violated, pointing towards specific higher-order correlations that could be incorporated into even more advanced models.
Also Read:
- Scaling Neural Networks to Millions of Classes with Latent Space Configuration
- Finding the Sweet Spot: How Training Noise Shapes Model Merging
In essence, this work highlights the power of PIGMMs as an interpretable and scalable framework for understanding the statistical properties of neural network weight matrices. By providing a reduced-degree-of-freedom model, PIGMMs offer a pathway to better analyze, optimize, and potentially design neural networks, especially as architectures grow in complexity. For more technical details, you can refer to the original research paper. Read the full research paper here.
