🤖 machine learning

Deep Learning Method for Stationary Distribution of Reflected Brownian Motion

This paper proposes a deep learning framework that leverages the basic adjoint relationship to accurately and efficiently compute the Laplace transform and tail probabilities of high-dimensional reflected Brownian motion, overcoming the limitations of existing closed-form solutions.

Original authors: Jim Dai, Zhanhao Zhang

Published 2026-07-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jim Dai, Zhanhao Zhang

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

Imagine you are trying to predict the weather in a city that is constantly being hit by rain, but the rain bounces off the buildings in weird, complicated ways. In the world of math and computer science, this "bouncing rain" is called Reflected Brownian Motion (RBM). It's a fancy way of describing how things (like customers in a line or data packets in a network) move around when they get pushed against walls and bounce back.

For a long time, mathematicians have been able to write down exact formulas to describe the "steady state" of this bouncing—basically, what the system looks like after it's been running for a long time. But here's the catch: those formulas only work for a few special cases, regardless of how many dimensions the system has. While some high-dimensional systems might theoretically have a solution, for most practical high-dimensional setups (like a massive, multi-layered warehouse with 20 or 30 dimensions), no closed-form solution exists. It's like trying to solve a puzzle where the pieces keep changing shape the more you look at them.

The Big Idea: Teaching a Neural Network to "Feel" the Math

In this paper, Jim Dai and Zhanhao Zhang from Cornell University propose a clever workaround. Instead of trying to solve the impossible math equation directly, they teach a deep learning neural network to learn the answer. Think of the neural network as a super-smart student who is given a set of rules (called the "Basic Adjoint Relationship" or BAR) and asked to figure out the pattern of the bouncing rain.

The goal isn't just to guess the average behavior; they want to know the Laplace transform. If you imagine the system's behavior as a complex song, the Laplace transform is the sheet music that lets you play any part of the song, including the rare, extreme notes (like a sudden, massive traffic jam). Once the network learns this "sheet music," the authors can use a special math trick (called the Talbot method) to turn it back into real-world predictions, like the probability of a queue getting impossibly long.

Why the Old Way Failed (and Why This One Works)

The authors tried a "naive" approach first: just throwing random data points at a standard neural network and asking it to minimize the error. It was a disaster.

The "Corner" Problem: In high dimensions, random sampling is terrible at finding the "corners" of the space. It's like trying to find a specific rare candy in a giant box by blindly grabbing handfuls; you'll mostly grab the common stuff in the middle and miss the rare stuff in the corners. But those corners are exactly where the most extreme (and important) behaviors happen.
The "Stability" Problem: The numbers involved can get huge or tiny very fast, causing the computer to get confused (numerical instability).
The "Scalability" Problem: Standard networks get too big and slow as the dimensions increase.

To fix this, the authors built a custom toolkit:

A Special Loss Function: Instead of just checking if the answer is right, they added "penalties" to the training. If the network predicts something that violates the laws of physics (like a probability that isn't smooth or doesn't decrease correctly), it gets a "scolding" (a penalty). They also added a "pairwise consistency" rule to make sure the network understands how the walls interact with the bouncing.
Smart Sampling: Instead of grabbing random handfuls, they designed a two-step sampling strategy. First, they pick a "target zone," then they specifically look for data points near the tricky corners and edges where the action is. This ensures the network sees the rare, extreme events it needs to learn.
A Scalable Architecture: They built a neural network that doesn't get bigger just because the room gets bigger. Instead of having a unique neuron for every single dimension, they use a "shared encoder" that treats each dimension similarly, then adds up the results. It's like having one master chef who can cook for 2, 20, or 30 people using the same recipe, rather than hiring a new chef for every extra guest.

The Results: Near-Perfect Predictions

The authors tested their method on three scenarios:

A 2-dimensional case where they knew the answer but didn't have a simple formula for the Laplace transform.
A 20-dimensional case.
A 30-dimensional case.

In all three, the neural network's predictions for "tail probabilities" (the chance of extreme events) matched the ground truth almost perfectly. In the 2D case, they compared the network's output to a known density function. In the 20D and 30D cases, they compared it to a known product-form solution. The results showed that the network could capture complex structures and scale up without losing accuracy.

What They Didn't Solve (Yet)

While the results are impressive, the authors are careful not to claim they've solved everything.

Memory Hunger: The current method requires a lot of computer memory. For each update, they sample 16,384 data points. If they tried to scale this up to systems with hundreds or thousands of dimensions, the memory usage would become a massive bottleneck, or the training would take forever.
Moments vs. Tails: When they tried to calculate specific "moments" (like the average or variance) using the learned transform, the results were good for low dimensions but got a bit wobbly in higher dimensions. The authors suggest this is because calculating moments requires very precise local information near zero, which is harder to get right than the broader picture needed for tail probabilities.
Future Work: They explicitly state that extending this to systems with hundreds or thousands of dimensions is a challenge for the future, as is applying it to other types of stochastic systems beyond reflected Brownian motion.

The Bottom Line

This paper suggests that deep learning can be a powerful tool for understanding complex, high-dimensional systems where traditional math fails. By combining a smart way of sampling data, a custom-built neural network architecture, and a loss function that respects the mathematical rules of the game, they created a method that can accurately predict extreme behaviors in 20 and 30-dimensional systems. It's not a magic wand that solves every problem instantly, but it's a significant step forward in making the "unanalyzable" analyzable.

Technical Summary: Deep Learning Method for Stationary Distribution of Reflected Brownian Motion

Problem Statement

The stationary distribution of Reflected Brownian Motion (RBM) is a critical tool for analyzing high-dimensional stochastic systems, particularly multiclass queueing networks under heavy traffic. While RBMs provide diffusion approximations for these systems, closed-form solutions for their stationary distributions are known only for a few special cases. Consequently, computing essential performance metrics, such as tail probabilities (e.g., tail latency), remains analytically intractable for general high-dimensional settings.

Existing numerical methods often struggle with the "curse of dimensionality." While deep learning has been applied to low-dimensional stochastic systems and partial differential equations (PDEs), extending these approaches to high-dimensional RBMs presents unique challenges. Specifically, the goal of this work is to develop a scalable method to compute the Laplace transform of the stationary distribution of a $d$ -dimensional RBM. Unlike methods that estimate steady-state expectations directly, recovering the full Laplace transform allows for the numerical inversion to obtain the full stationary distribution and tail probabilities.

Methodology

The proposed framework approximates the Laplace transforms $\phi_k(\theta)$ (for $k=0, \dots, d$ ) using feedforward neural networks operating in the complex domain. The approach is built upon the Basic Adjoint Relationship (BAR), which characterizes the Laplace transforms via the equation:
$\gamma_0(\theta)\phi_0(\theta) = \sum_{k=1}^d \gamma_k(\theta)\phi_k(\theta)$
where $\gamma_0(\theta) = -\frac{1}{2}\langle\theta, \Sigma\theta\rangle + \langle\mu, \theta\rangle$ and $\gamma_k(\theta) = -\langle R^{(k)}, \theta\rangle$ .

To address the challenges of high-dimensional learning, the authors introduce three core methodological components:

1. Loss Function Design

Instead of minimizing the raw residual of the BAR equation, the authors parameterize the logarithms of the transforms ( $\phi_k = \exp(f_k)$ ) and construct a composite loss function $L(\theta)$ comprising five terms:

Normalized BAR Error ( $L_{BAR}$ ): To handle the exponential variation of Laplace transforms, the equation sides are normalized by a factor $\nu(\theta)$ before computing the squared relative error. This ensures numerical stability across different magnitudes of $\theta$ .
Pairwise Consistency Penalty ( $L_{pair}$ ): To accelerate convergence, the authors construct specific boundary points $\tilde{\theta}^{(k)}$ where only one reflection term is active. This creates explicit constraints linking the interior function $\phi_0$ to each boundary function $\phi_k$ , rather than relying solely on the aggregate residual.
Monotonicity Penalty ( $L_{mono}$ ): Enforces that the real part of the Laplace transform is non-increasing with respect to the real part of $\theta$ , and ensures real-valued outputs on the real axis.
Cauchy-Riemann Penalty ( $L_{CR}$ ): Enforces the analyticity of the learned functions by penalizing violations of the Cauchy-Riemann equations between the real and imaginary parts.
Zero-Anchoring Penalty ( $L_{zero}$ ): Fixes the normalization at the origin ( $\phi_0(0)=1$ ).

2. Training Data Sampling Strategy

Standard uniform sampling in high-dimensional boxes fails to capture "corner regions" where coordinates are simultaneously near their extremes. These regions are critical for accurate numerical inversion (e.g., via the Talbot method). The authors propose a two-stage conditional sampling procedure:

Sample scalar reference thresholds for real and imaginary parts.
Conditionally sample coordinates to increase the probability of drawing points near the lower corner of the real domain and the outer corners of the imaginary domain.

3. Scalable Neural Network Architecture

To ensure the number of trainable parameters does not scale with the dimension $d$ , the architecture employs shared encoders and additive aggregation:

Shared Coordinate Encoder: Each coordinate $\theta_j$ is processed by a shared network that combines Fourier features (to capture oscillatory patterns) with a coordinate index embedding.
Interior Network: Aggregates the outputs of the coordinate encoders additively to produce the interior transform $\phi_0$ .
Shared Boundary Network: Uses a separate shared network (with boundary index embeddings) to aggregate contributions for each boundary transform $\phi_k$ .
This design allows the model to handle high-dimensional inputs while maintaining a fixed parameter count regardless of $d$ .

Key Contributions

First Framework for High-Dimensional RBM Laplace Transforms: This work presents the first deep learning framework specifically designed to estimate the Laplace transform $\phi_k(\cdot)$ for high-dimensional RBMs.
Scalable Architecture: The proposed neural network architecture decouples the number of parameters from the system dimension $d$ , addressing a major bottleneck in applying deep learning to high-dimensional stochastic systems.
Structured Loss and Sampling: The integration of a normalized loss function, structural penalties (monotonicity, analyticity), and targeted sampling schemes overcomes issues of poor generalization, numerical instability, and imbalanced data distribution inherent in naive approaches.

Experimental Results

The method was evaluated on three RBM instances with known ground-truth tail probabilities:

2-Dimensional RBM: A case with a known stationary density but no closed-form Laplace transform. The neural network successfully captured complex structural behaviors.
20-Dimensional and 30-Dimensional RBMs: Cases with product-form Laplace transforms.
Performance: The predicted tail probabilities matched the ground truth "almost perfectly" across all dimensions. The results demonstrate that the framework is both expressive enough to represent complex transforms and scalable to high-dimensional settings (up to 30 dimensions in the experiments).
Moment Estimation: The paper also reports on moment estimation via numerical inversion. While first-order moments remained reasonably accurate in high dimensions, errors increased for higher-order moments and in higher dimensions, attributed to the sensitivity of contour-based derivative approximations near $\theta=0$ .

Significance and Claims

The authors claim that this framework provides a general tool for analyzing stochastic systems beyond analytically tractable regimes. By accurately computing the Laplace transform, the method enables the estimation of critical performance metrics like tail latencies in large-scale queueing networks where traditional analytical methods fail.

The paper maintains a modest stance regarding its limitations:

Computational Cost: Current gradient updates require large batch sizes (16,384 samples), leading to significant GPU memory usage or runtime as dimensionality increases. Scaling to hundreds or thousands of dimensions requires more efficient training strategies.
Scope: The current work is specific to RBMs. Extending the framework to broader classes of stochastic systems and other performance metrics (beyond tail probabilities) is identified as future work.

The authors emphasize that while the method is not yet ready for systems with hundreds or thousands of dimensions without further optimization, it represents a significant step toward scalable analysis of high-dimensional stochastic systems.