🤖 machine learning

Learning in PINNs: Phase transition, diffusion equilibrium, and generalization

This paper investigates the learning dynamics of neural networks by identifying a "diffusion equilibrium" phase characterized by aligned sample-wise gradients and homogeneous residuals, which drives faster convergence and improved generalization, leading to a proposed re-weighting scheme that enhances performance in physics-informed neural networks (PINNs).

Original authors: Sokratis J. Anagnostopoulos, Juan Diego Toscano, Nikolaos Stergiopulos, George Em Karniadakis

Published 2026-08-25

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Sokratis J. Anagnostopoulos, Juan Diego Toscano, Nikolaos Stergiopulos, George Em Karniadakis

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

In the vast landscape of modern artificial intelligence, a specific type of computer program has emerged that acts as a bridge between raw data and the fundamental laws of physics. These programs, known as physics-informed neural networks, are designed to solve complex equations that describe how the world works, from the flow of blood in the human body to the movement of air over an airplane wing. Unlike traditional methods that rely on rigid, step-by-step calculations, these networks learn by trial and error, constantly adjusting their internal settings to minimize the difference between their predictions and the known rules of nature. However, this learning process is often messy and unpredictable. The path to a correct solution is rarely a straight line; instead, it is a winding journey through a rugged terrain of possibilities, where the computer can easily get stuck in local traps or fail to learn the full picture. Understanding exactly how these digital minds navigate this difficult terrain, and what happens when they finally find their way, is crucial for making them faster, more reliable, and better at solving real-world problems.

A team of researchers has now mapped out the hidden stages of this learning journey, revealing that the process is not a single, continuous flow but rather a sequence of distinct phases. By observing how the computer's internal signals change over time, they discovered that the training process moves through an initial period of rapid adjustment, followed by a slower, more chaotic phase of exploration. But their most significant finding is the identification of a third, previously overlooked stage that acts as the true turning point for success. They call this stage "diffusion equilibrium." It is a moment of sudden clarity where the computer's internal signals, which were previously noisy and conflicting, suddenly align and agree on a single direction. This alignment is not just a minor improvement; it is the key that unlocks the ability to generalize, allowing the model to perform well on new, unseen data rather than just memorizing the training examples.

The researchers arrived at this conclusion by closely watching the behavior of these networks as they solved problems involving fluid dynamics and wave equations. They measured the ratio of the useful signal in the computer's learning process against the background noise. In the beginning, the signal is strong and the computer learns quickly. Then, as it enters the exploration phase, the noise takes over, and the learning becomes slow and uncertain. For a long time, scientists believed that this noisy phase was where the real learning happened. However, the team observed that for these physics-based models, the process does not end there. After a long period of noise, something dramatic happens: the noise suddenly drops, and the signals from different parts of the problem space snap into perfect agreement. This is the moment of diffusion equilibrium. It is a stable state where the computer has found a consistent path forward, and it is only after reaching this state that the error in its predictions begins to plummet.

What makes this discovery particularly powerful is the realization that this alignment of signals is not enough on its own. The researchers found that for the computer to truly succeed, the errors it makes across the entire problem space must also be uniform. If the computer is very accurate in some areas but makes huge mistakes in others, it will fail to generalize, no matter how well the signals align. To fix this, they developed a simple but effective technique that acts like a spotlight, focusing the computer's attention on the specific areas where it is struggling the most. By giving extra weight to the difficult parts of the problem, they forced the computer to balance its learning across the entire domain. This approach, which they call residual-based attention, helped the models reach the diffusion equilibrium stage much faster and with a much more even distribution of errors. In their tests, this method allowed the computer to solve problems ten times faster than the standard approach, while also producing results that were far more accurate and reliable.

The study also shed light on what happens inside the computer's "brain" during this critical transition. As the model moves into this stable equilibrium, the internal values that the computer uses to make decisions begin to saturate, meaning they hit their maximum or minimum limits. This saturation causes the computer's internal representation of the problem to become highly compressed, almost like turning a complex, colorful image into a simple black-and-white sketch. Surprisingly, this compression does not mean the computer is losing important information. Instead, it suggests that the model has found the most efficient way to store the essential details needed to solve the problem. The researchers noted that this compression happens most intensely in the middle layers of the network, creating a structure that resembles a system that first encodes information and then decodes it to find the solution. This finding challenges the old idea that compression is always a sign of information loss; here, it appears to be a sign of the model having finally understood the problem deeply enough to represent it efficiently.

These insights offer a new way to think about how artificial intelligence learns. The researchers suggest that the key to better performance is not just about finding the right settings or adding more data, but about guiding the learning process through these specific phases until it reaches that moment of equilibrium. By ensuring that the computer pays attention to all parts of the problem equally and waits for the signals to align, we can help it avoid getting stuck in suboptimal solutions. While this work focused on problems involving the laws of physics, the principles discovered here could apply to a much wider range of artificial intelligence tasks. The ability to recognize when a model has truly learned, and to steer it toward that state of balance, could lead to smarter, more efficient algorithms for everything from medical diagnosis to climate modeling. The journey from confusion to clarity is no longer a mystery; it is a process that can be understood, measured, and improved.

Technical Summary: Learning in PINNs: Phase Transition, Diffusion Equilibrium, and Generalization

Problem Statement

The optimization of Physics-Informed Neural Networks (PINNs) and deep learning models generally involves navigating non-convex loss landscapes characterized by complex, discrete phases rather than smooth convergence. While the Information Bottleneck (IB) theory posits a two-phase learning regime (fitting and diffusion) separated by a phase transition, the dynamics of full-batch optimizers like Adam in PINNs remain under-explored. Specifically, current methods often suffer from:

Gradient Stochasticity: Even with full-batch training, sample-wise gradients may disagree, leading to instabilities or convergence to sub-optimal local minima where gradients cancel out.
Residual Heterogeneity: In PINNs, samples (collocation points) are interdependent via Partial Differential Equations (PDEs). Standard quadratic loss averaging can mask "outlier" samples with large residuals but vanishing gradients, causing the optimizer to stall while other regions of the domain remain poorly resolved.
Generalization Gaps: Models may converge to stationary solutions with high test errors because the optimization process fails to achieve uniform sensitivity across all training samples.

Methodology

The authors analyze training dynamics through the lens of the Neural Gradient Signal-to-Noise Ratio (SNR) and propose a re-weighting mechanism to enforce residual homogeneity.

1. Gradient Signal-to-Noise Ratio (SNR)

The study defines SNR as the ratio of the batch-wise mean gradient (signal) to the batch-wise standard deviation (noise) across network parameters:
$\text{SNR} = \frac{\|\mathbb{E}_B[\nabla_\theta L_B]\|_2}{\|\text{std}_B[\nabla_\theta L_B]\|_2}$

High SNR: Indicates deterministic, aligned gradient directions.
Low SNR: Indicates stochastic, conflicting gradient directions.
Order Parameter ( $\phi$ ): A bounded metric derived from SNR to measure the degree of gradient alignment and typicality among samples.

2. Identification of Three Phases

The authors identify three distinct phases in PINN training with Adam:

Fitting Phase: High SNR, rapid deterministic descent.
Diffusion Phase: Low SNR, stochastic exploration where the model generalizes but convergence is slow.
Diffusion Equilibrium (DE): A newly identified third phase occurring after prolonged training. It is characterized by an abrupt transition where SNR jumps to a stable, high-value region. In this phase, sample-wise gradients align, and the optimizer reaches a stable equilibrium.

3. Residual-Based Attention (RBA)

To address the issue of heterogeneous residuals (where some samples have large errors but small gradients), the authors propose a Residual-Based Attention (RBA) re-weighting scheme.

Mechanism: The loss function is modified to $L^* = \frac{1}{n} \sum (\lambda_i R_i)^2$ , where $\lambda_i$ is a weight calculated as a moving average of the relative residuals.
Goal: Dynamically rescale gradients to emphasize samples with large, unresolved residuals. This prevents the optimizer from ignoring "stiff" samples where gradients vanish, effectively enforcing residual homogeneity across the sample space.

4. Information Compression Analysis

The study investigates the relationship between phase transitions and information compression (IB theory). Using a relative-binning approach to discretize activations, the authors measure mutual information $I(X; T)$ and the saturation of neuron activations (specifically tanh) to determine if compression correlates with the DE phase.

Key Contributions

Discovery of Diffusion Equilibrium (DE): The paper identifies a third phase in PINN training where stable convergence occurs. This phase is marked by a sudden increase in SNR and the alignment of sample-wise gradients.
Gradient and Residual Order: The authors demonstrate that both high gradient order (alignment) and high residual order (homogeneity) are required for optimal generalization. Gradient alignment alone is insufficient if residuals remain heterogeneous.
RBA Algorithm: A generic, low-cost re-weighting technique is proposed that accelerates the onset of the DE phase by targeting problematic samples, leading to faster convergence and better generalization.
Saturation-Induced Compression: The study observes that the transition to DE coincides with a saturation of neuron activations ("binarization"), particularly in middle layers, suggesting an encoder-decoder-like hierarchy. However, it clarifies that this geometric compression does not necessarily imply significant theoretical information loss.

Results

The methodology was validated on four benchmarks:

Allen-Cahn Equation
Helmholtz Equation
Lid-Driven Cavity Flow (Re=1000)
Transient 3D Navier-Stokes (Cerebrospinal Fluid flow)

Key Findings:

Convergence Speed: RBA models reached the DE phase significantly faster than vanilla models (e.g., 10k iterations vs. 100k for Allen-Cahn to reach 10% test error).
Generalization: Vanilla models often converged to sub-optimal solutions with heterogeneous residuals (accumulated errors in specific regions), leading to high test errors (e.g., ~~60% for cavity flow). RBA models achieved uniform residual distribution and significantly lower test errors (~~6%).
SNR Dynamics: The DE phase corresponds to a stable equilibrium where $\frac{\partial}{\partial t}\text{SNR} \approx 0$ . The correlation between the batch-wise SNR and the Adam-specific SNR ( $\text{SNR}_{Adam}$ ) confirms that Adam stabilizes when sample gradients are homogeneous.
Information Compression: During the DE transition, activation saturation increases sharply. While deeper layers show negligible information loss, middle layers exhibit the most pronounced saturation, consistent with an encoding-decoding process.

Significance and Claims

The paper claims that identifying the Diffusion Equilibrium phase provides a critical insight into deep learning optimization, particularly for PINNs where sample interdependence is inherent.

Optimization Strategy: The findings suggest that optimal learning occurs not just during the stochastic diffusion phase, but specifically when the system reaches a state of gradient agreement and residual homogeneity.
Generalization: The paper posits that generalization is enhanced when the optimization steps are equally sensitive to every sample, a condition enforced by the RBA scheme.
Theoretical Connection: The work bridges Information Bottleneck theory with scientific machine learning, offering a consistent description of training dynamics where "binarization" of activations coincides with convergence.
Practical Impact: The RBA method is presented as a generic algorithm applicable to any quadratic loss function suffering from slow-converging outliers, potentially enhancing performance beyond PINNs without significant computational overhead.

The authors maintain a modest stance, noting that while the DE phase is observed across benchmarks, the theoretical bounds of SNR and the universality of this transition across different architectures (e.g., CNNs) require further investigation. They do not claim that compression is a prerequisite for generalization, but rather that the observed saturation is a consequence of gradient alignment during the DE phase.