💻 computer science

On the Stability of the Jacobian Matrix in Deep Neural Networks

This paper establishes a general stability theorem for the Jacobian matrix in deep neural networks with sparse and weakly correlated weights, extending rigorous spectral stability guarantees beyond traditional fully connected networks with i.i.d. weights by leveraging recent advances in random matrix theory.

Original authors: Benjamin Dadoun, Soufiane Hayou, Hanan Salam, Mohamed El Amine Seddik, Pierre Youssef

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin Dadoun, Soufiane Hayou, Hanan Salam, Mohamed El Amine Seddik, Pierre Youssef

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Whispering Game" of AI

Imagine a Deep Neural Network (DNN) as a long line of people playing the "whispering game" (also known as "Telephone").

The Input: A message whispered into the first person's ear.
The Layers: Each person in the line hears the message, adds a tiny bit of their own interpretation, and whispers it to the next person.
The Output: The final message heard by the last person.

In this paper, the authors are worried about how the message changes as it travels down the line. They call this the Jacobian.

Vanishing Gradient: If the message gets quieter and quieter at every step, the last person hears nothing. The network "forgets" the input.
Exploding Gradient: If the message gets louder and louder (shouted), the last person is deafened. The network becomes chaotic and unstable.

The goal of the paper is to figure out how to set up the game so the message stays at a "Goldilocks" volume—not too quiet, not too loud—no matter how long the line of people is.

The Problem: Old Rules Don't Work Anymore

Previously, scientists knew how to set up this game if everyone was a stranger with a random voice (independent, random weights). They found a "magic setting" (called the Edge of Chaos) where the message stays stable.

However, real-world AI isn't always that simple. The authors looked at two messy, real-world scenarios where the old rules might break:

Pruned Networks (The Sparse Network): Imagine you cut out half the people in the line to save space. Does the message still travel well?
Correlated Weights (The Dependent Network): Imagine the people in the line are friends who tend to whisper in similar tones or copy each other. Does the message stay stable?

The paper asks: Can we still keep the message stable in these messy situations?

The Solution: A Universal "Volume Knob"

The authors developed a new mathematical rule (a Universality Theorem) that acts like a universal volume knob. They proved that even in messy situations, if you adjust the "volume" correctly, the message will behave just like it does in the perfect, random scenario.

Here is how they solved the two specific problems:

1. The Sparse Network (Pruning)

The Analogy: Imagine you remove 90% of the people from the whispering line. Naturally, the message will die out because there are fewer people to pass it along.
The Paper's Finding: You can fix this! If you remove people, you must turn up the volume of the remaining people to compensate.

Random Pruning: If you cut people out randomly, you need to turn up the volume by a specific amount (mathematically, a factor of $1/\sqrt{1-s}$ , where $s$ is how much you cut).
Magnitude Pruning: If you cut people out based on who is whispering the loudest (keeping only the "important" ones), the math changes slightly. You need a different volume knob setting than the random one.
The Result: If you use the wrong volume knob, the message either vanishes or explodes. If you use the right one, the network stays stable, even if it's 99% empty.

2. The Correlated Network (Dependent Weights)

The Analogy: Imagine the people in the line are a group of friends who all whisper in the exact same pitch. If they are too similar, the message might get distorted or amplified in weird ways.
The Paper's Finding: You can have friends whispering together, but they can't be too similar.

There is a strict limit on how much they can "copy" each other. If the correlation (similarity) is too high, the message breaks.
However, if the similarity is kept below a very specific, tiny threshold (related to the size of the network), the message travels perfectly fine, just as if they were strangers.

The "Magic" Discovery

The most exciting part of the paper is the Universality claim.
The authors proved that:

A network with randomly cut connections (if scaled correctly).
A network with slightly copying friends (if correlation is low enough).
A network with perfect random strangers (the old standard).

...all behave exactly the same way regarding the stability of the message. They all reach the same "Goldilocks" state.

Why This Matters (According to the Paper)

The paper doesn't claim to invent new AI models or cure diseases. Instead, it provides the theoretical safety manual for modern AI practices.

It explains why we need to re-scale weights after pruning a network (a common practice to make AI run faster on phones).
It tells us exactly how much "friendship" (correlation) we can allow between weights before the AI starts to fail.
It gives rigorous mathematical proof that these "messy" setups can be just as stable as the "perfect" theoretical setups, provided you follow the specific scaling rules they discovered.

Summary

Think of this paper as a guide for building a very long, very complex relay race.

Old Guide: "Only run this race if every runner is a stranger and stands in a perfect line."
New Guide (This Paper): "You can have runners who are friends, and you can even remove some runners to save time! But, you must adjust their running speed (scaling) and make sure the friends aren't too synchronized. If you follow our new math, the race will finish smoothly every time."

Technical Summary: On the Stability of the Jacobian Matrix in Deep Neural Networks

Problem Statement
Deep neural networks (DNNs) frequently suffer from exploding or vanishing gradients as network depth increases, a phenomenon intrinsically linked to the spectral behavior of the input-output Jacobian matrix. While prior research has established critical initialization schemes (e.g., the "Edge of Chaos" regime) that ensure Jacobian stability for fully connected networks with independent and identically distributed (i.i.d.) Gaussian weights, these analyses are limited. They do not readily extend to practical scenarios involving structured or dependent randomness, such as:

Correlated Weights: Arising from data-dependent initialization, meta-learning, or fine-tuning pretrained models.
Sparse Networks: Resulting from pruning techniques (e.g., Random Pruning, Magnitude-based Pruning) used to reduce computational footprints.

The behavior of the Jacobian in these non-standard settings remains poorly understood, particularly regarding how weight correlations and sparsity affect gradient propagation and training stability in the large-depth regime.

Methodology
The authors develop a general universality theorem for Jacobian products (Theorem 3.1) based on recent breakthroughs in random matrix theory (specifically [BvH24]). The analysis focuses on Multi-Layer Perceptrons (MLPs) in the fixed-depth, infinite-width regime.

The core technical approach involves:

Layer-by-Layer Replacement: The authors employ a comparison argument where a general random matrix layer is replaced by a corresponding i.i.d. Gaussian layer.
Hermitian Dilation: To handle non-self-adjoint matrix products, the Jacobian terms are mapped to self-adjoint matrices using Hermitian dilation.
Concentration Inequalities: The proofs rely on advanced concentration results for sums of independent self-adjoint random matrices and Gaussian matrices with dependent entries.
Variance Matching: For sparse networks, the authors identify specific rescaling factors for the retained weights to match the second moment (variance) of the critically initialized i.i.d. Gaussian model.

Key Contributions and Results

Universality Theorem for Jacobian Products (Theorem 3.1):
The paper establishes that suitably rescaled sparse networks and weakly correlated Gaussian networks possess the same fixed-depth, large-width Jacobian spectral norm as the critically initialized i.i.d. Gaussian model. This holds under two specific conditions:
- Independent Entries: The maximum entry magnitude must satisfy a specific decay condition ( $r_n \log^4 n \to 0$ ).
- Weakly Correlated Gaussian Entries: The maximum correlation between distinct entries must satisfy $\rho_n = o(n^{-5/2} / \log n)$ .
Stability in Sparse Networks:
- Random Pruning: The authors prove that for random pruning with sparsity $s_n$ , stability is preserved if the weights are rescaled by a factor of $(1-s_n)^{-1/2}$ (Theorem 4.1). Without this scaling, the Jacobian norm exhibits exponential dependence on depth, leading to instability. The paper identifies an "edge of stability" where sparsity approaches $1 - s_n \sim n^{-1} \log^9 n$ ; beyond this threshold, the universality result breaks down.
- Magnitude-Based Pruning: The study extends to score-based (magnitude) pruning (Theorems 4.2 and 4.3). Crucially, the required scaling factor differs from that of random pruning. The paper demonstrates that using the random pruning scaling factor on magnitude-pruned networks fails to guarantee stability, necessitating a method-specific rescaling derived from the threshold distribution of the weights.
Stability with Dependent Weights (Theorem 5.1):
The paper quantifies the tolerance for weight correlations. It shows that if the correlation between weight entries is $o(n^{-3/2} / \log n)$ , the Jacobian spectral norm converges to the same deterministic limit as the i.i.d. case. This provides a rigorous bound on the level of correlation permissible before Jacobian stability is compromised.
Empirical Verification:
Theoretical findings are supported by empirical results on MLPs trained with Stochastic Gradient Descent (SGD). Experiments confirm that:
- Approximations regarding the independence of diagonal activation gates and weights hold for i.i.d. weights but degrade under high correlation.
- Applying the correct rescaling factors to pruned networks restores trainability and prevents Jacobian explosion/vanishing.
- Networks with weakly correlated weights behave similarly to i.i.d. networks, while high correlations lead to divergent Jacobian norms.

Significance and Claims
The authors claim to extend the theoretical foundation for initialization and weight scaling in modern neural networks beyond the restrictive assumption of i.i.d. weights. By establishing a universality principle, the paper provides rigorous guarantees for spectral stability in broader classes of network models, including those with structured randomness (sparsity) and dependent weights.

The work highlights that:

Rescaling is Essential: Simply pruning a network or introducing mild correlations is insufficient; specific, method-dependent rescaling is required to maintain the "Edge of Chaos" stability regime.
Universality: Despite differences in initialization (pruning strategies, correlation structures), the spectral behavior of the Jacobian converges to the same limit as the critical i.i.d. case, provided the variance is matched and correlations/sparsity remain within derived bounds.

Limitations
The authors explicitly note that their current theoretical framework is limited to the Multi-Layer Perceptron (MLP) architecture. Furthermore, the analysis of sparse networks is restricted to Random Pruning and Magnitude-based Pruning performed at initialization. Extending these results to more complex architectures (e.g., Convolutional Neural Networks, Transformers) and other pruning algorithms (e.g., gradient-based pruning like SNIP or GraSP) remains an open question for future work. Additionally, the results are derived in the fixed-depth, infinite-width limit and do not address the joint limit where depth and width diverge simultaneously.