Forecasting Stock Prices with Diffusion Models and U-GNNs
A paper-to-Python guide to graph diffusion, uncertainty-aware S&P 500 forecasting, and wireless resource allocation
What This Article Builds
The paper proposes a unified framework for conditional generative modeling of graph signals using denoising diffusion models. The key innovation is the U-Graph Neural Network (U-GNN), a U-Net-like architecture for graph-structured data that uses learned node selection matrices for pooling/unpooling and strided graph convolutions to avoid explicit graph coarsening. The framework is demonstrated on two tasks: stock price forecasting (S&P 500) and wireless resource allocation (power control).
Implementation Assumptions
PyTorch 2.x + PyTorch Geometric (PyG) + standard scientific Python (numpy, scipy, matplotlib).
Graph shift operator S stored as a sparse PyTorch tensor (coalesced COO) or edgeindex+edgeweight.
All graph signals are batched: (batch, N, F) with optional node dimension per graph (single graph assumed for each task).
The U-GNN is built from modular components: StridedGraphConv, NodeSelection, FusionBlock, UGNN core.
Training uses AMP (torch.cuda.amp), AdamW, cosine LR with linear warm-up, gradient clipping (max_norm=1.0).
Validation and checkpointing use task-specific composite criteria.
Data pipeline for S&P 500 uses yfinance; graph built from fundamentals (rank correlation + sector bonus).
Data pipeline for WRA synthesizes networks using the paper's channel model and expert primal-dual algorithm.
DDIM sampling uses 100 steps and η=0.2 on a sub-grid of the 500-step linear schedule.
Download Source code.
You can download the source using the URL at the end of this article.
Note: There was no source code attached to the paper, this is my implementation, I tried to code it, and it is part of the larger program, so use the code to understand and implement in your code, try running it would run but there might be some problem. So Use code as reference, not in actual public working.
1. Introduction to the Paper and Problem
1.1 What Are Stochastic Graph Signals?
Many real-world systems produce signals that live on a graph — a set of nodes connected by edges that encode relationships. When those signals are inherently random, we call them stochastic graph signals. Two canonical examples, both studied in the paper, are:
Stock returns on a company graph: Each node is a publicly traded company (e.g., the S&P 500 constituents). The graph edges capture fundamental similarities: same industry sector, correlated financial profiles, or shared supply chains. The signal at each node is the daily log-return of that stock. The signal is stochastic because future returns are unpredictable in detail, yet they are conditionally dependent on the past and on the graph structure (e.g., a shock to one sector propagates to related companies).
Power allocations on an interference graph: Each node is a wireless transmitter–receiver pair. The graph edges represent interference strength: a directed edge from transmitter j to receiver i carries the channel gain between them. The signal at each node is the transmit power level. The optimal allocation is stochastic because it depends on random fading and the expert algorithm’s internal randomness, but it is conditioned on the large-scale channel gains (the graph) and the node’s own channel quality.
In both cases, the graph topology and node-level side information (history features, channel measurements) provide a conditioning context. The goal is to sample new graph signals from the unknown conditional distribution — not just to predict a single value, but to generate entire trajectories or allocations that respect the graph structure and the observed context.
1.2 Why Generative Modeling (Not Just Regression)?
A regressive approach — training a neural network to minimize mean squared error between a point prediction and the true signal — learns the conditional mean. This is adequate when the conditional distribution is unimodal and symmetric, but it fails when:
The distribution is multi-modal: there may be several equally plausible future price paths or power allocations.
Uncertainty quantification is required: a point forecast gives no sense of risk or confidence intervals.
The downstream task needs samples, not averages: for portfolio optimization, one needs many possible return scenarios; for wireless resource allocation, one needs to test different power settings under fading.
Generative models address these limitations by learning the full conditional law. The paper adopts a denoising diffusion probabilistic model (DDPM), which has become the state of the art for high-quality sample generation in images, audio, and now graph signals. Diffusion models are particularly attractive because they are stable to train, do not suffer from mode collapse (unlike GANs), and can be conditioned naturally on arbitrary inputs.
1.3 The Two Applications
The paper demonstrates the framework on two distinct tasks, each with its own graph structure and conditioning modality:
Stock Price Forecasting (S&P 500)
Data: Daily prices for 468 stocks over ~10 years. For each stock, 12 market features (open, high, low, close, log-return, moving averages, volume, RSI, MACD) are computed. A static undirected graph is built from rank correlations of fundamental profiles, with a same-sector bonus.
Conditioning: The last 20 days of features (history) and the graph topology.
Target: The next 5 days of log-returns (a graph signal of dimension 5 per node).
Challenge: The conditional distribution of future returns is heavy-tailed, heteroscedastic, and exhibits volatility clustering. A generative model can produce multiple plausible trajectories, enabling risk-aware decision-making.
Wireless Resource Allocation (Power Control)
Data: Synthetic networks of 400 transmitter–receiver pairs in square areas of varying density. Channel gains follow a dual-slope path loss model with log-normal shadowing and Rayleigh fading. An expert primal–dual algorithm computes optimal power allocations that maximize sum ergodic rate subject to per-user minimum rate constraints.
Conditioning: The directed interference graph (top-10 strongest interferers per receiver, log-normalized) and two node-state features (direct link gain and aggregate interference under full power).
Target: The centered power allocation (p/Pmax − 1) for each user.
Challenge: The expert algorithm is computationally expensive; a learned generative model can approximate its output distribution at a fraction of the cost, and can be deployed in real time.
1.4 Overview of the Proposed Solution
The paper proposes a unified framework that combines a denoising diffusion probabilistic model with a novel U‑Graph Neural Network (U‑GNN) denoiser. The key ideas are:
Forward diffusion: Gradually add Gaussian noise to the clean graph signal over 500 steps, following a linear noise schedule (β₁ = 1e−4, β₅₀₀ = 2e−2). The noising process is fixed and does not depend on the graph.
Reverse denoising: Learn a neural network ε_θ that predicts the noise added at each step, conditioned on the graph shift operator S and node features u. The network is trained with the simple mean-squared error between true and predicted noise (equation (6) in the paper).
U‑GNN denoiser: A U‑Net-like architecture for graph signals. It processes the signal at multiple resolutions by:
Strided graph convolutions (Algorithm 1): Instead of coarsening the graph, the signal is lifted to the full graph, filtered with iterated unit shifts of S, and then reduced to a subset of active nodes. The stride γ increases the receptive field without densifying the graph.
Learned node selection (Section IV‑B): At each down‑sampling level, a scoring MLP selects the top‑k nodes to retain. During training, Gumbel noise encourages exploration, and a straight‑through estimator allows gradients to flow through the discrete selection.
Fusion layers (equation (22)): The down‑sampled signal is merged with global embeddings (node states and diffusion step) via cross‑attention (for stock) or addition (for WRA).
Skip connections (equation (23)): The decoder receives concatenated features from the encoder at the same resolution.
Accelerated sampling: Instead of the full 500‑step reverse process, the paper uses DDIM (Denoising Diffusion Implicit Models) with 100 steps and η = 0.2, which produces high‑quality samples much faster.
Task‑specific adaptations: For stock forecasting, the U‑GNN includes a lightweight temporal mixer (interleaved dilated 1‑D convolutions and two‑head self‑attention) to process the history features. For WRA, the conditioning is static and uses simple addition.
1.5 What This Tutorial Covers
This tutorial walks through a complete, modular reproduction of the paper’s stock forecasting and wireless resource allocation experiments. We will:
Implement the forward diffusion process and DDIM sampling from scratch.
Build the U‑GNN architecture piece by piece: strided graph convolutions, learned node selection, fusion layers, and the full encoder–decoder.
Construct the data pipelines for both tasks, including synthetic data generation for WRA and a simplified stock data pipeline (using synthetic data to avoid network calls).
Train the model with the exact hyperparameters from the paper (Table I, Appendix E).
Evaluate using the paper’s metrics: CRPS, MIS90, RMSE, direction accuracy, stylized‑fact gaps for stock; ergodic rate, feasibility, and gap to expert for WRA.
Reproduce the key figures: forecasting trajectories, CDF comparisons, distribution diagnostics, and WRA bar charts.
All code is written in Python using PyTorch and PyTorch Geometric, and is organized into modular files for clarity. The implementation follows the paper’s descriptions and appendices faithfully, with documented decisions where the paper is ambiguous.
Let’s begin by understanding the diffusion process and how it is adapted to graph signals.
2. Denoising Diffusion Probabilistic Models for Graph Signals
2.1 Forward Diffusion Process
The core idea of a denoising diffusion model is to define a forward process that gradually corrupts a clean signal into pure noise, and then learn a reverse process that undoes this corruption step by step. For graph signals, the forward process operates on the node features directly, independent of the graph structure at this stage — the graph topology enters only as conditioning information for the reverse denoiser.
Let x_0 \in \mathbb{R}^{N \times F} be a clean graph signal (e.g., the future returns of N stocks over F time steps). The paper defines a Markov chain of K = 500 steps that adds Gaussian noise with a linear variance schedule:
Each forward transition is (equation (2)):
Because the Gaussian is closed under composition, we can directly sample x_k from x_0 without iterating through all intermediate steps. Define \alpha_k = 1 - \beta_k and the cumulative product \bar{\alpha}_k = \prod_{i=1}^k \alpha_i. Then the marginal distribution is (equation (3)):
which leads to the reparameterization used in training (equation (4)):
Implementation in diffusion.py
The DiffusionProcess class precomputes all \bar{\alpha}_k values from the linear \beta schedule. The constructor stores the betas, alphas, and derived quantities as tensors for efficient indexing:
class DiffusionProcess:
def __init__(self, betas: torch.Tensor) -> None:
self.K = betas.shape[0]
self.betas = betas
self.alphas = 1.0 - betas
# alpha_bar[t] = prod_{i=1}^{t} alpha_i, alpha_bar[0] = 1
self.alpha_bar = torch.cat([
torch.ones(1, dtype=betas.dtype, device=betas.device),
torch.cumprod(self.alphas, dim=0)
], dim=0)
self.sqrt_alpha_bar = self.alpha_bar.sqrt()
self.sqrt_one_minus_alpha_bar = torch.sqrt(1.0 - self.alpha_bar)The forward_diffuse method implements equation (4) directly. It takes a clean signal x_0 and a step index k (1-indexed), samples noise, and returns the noised signal:
def forward_diffuse(self, x0: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
k_idx = k.long()
sqrt_alpha_k = self.sqrt_alpha_bar[k_idx]
sqrt_one_minus_alpha_k = self.sqrt_one_minus_alpha_bar[k_idx]
epsilon = torch.randn_like(x0)
x_k = sqrt_alpha_k * x0 + sqrt_one_minus_alpha_k * epsilon
return x_k, epsilonDuring training, we sample k \sim \text{Uniform}(1, K) for each batch element, compute x_k using this method, and then ask the denoiser to predict the noise \varepsilon.
2.2 Reverse Denoising and the Noise-Prediction Objective
The reverse process is a Markov chain that starts from pure noise x_K \sim \mathcal{N}(0,I) and iteratively removes noise to recover x_0. Each reverse step is parameterized as a Gaussian with learned mean (equation (5)):
Here \xi = (S, u) denotes the conditioning information: the graph shift operator S and node features u. The mean \mu_\theta is derived from a noise-prediction network \varepsilon_\theta (the U-GNN) via equation (8):
Rather than directly predicting \mu_\theta, the paper adopts the simplified training objective from Ho et al. (2020) (equation (6)):
This loss is simple: at each training step, we sample a random k, noise \varepsilon, and a clean signal x_0, compute x_k via the forward process, and ask the U-GNN to predict the noise that was added. The gradient flows through the U-GNN only; the forward process is fixed.
2.3 DDIM Sampling for Acceleration
While the reverse process could be simulated for all K = 500 steps, the paper uses the Denoising Diffusion Implicit Model (DDIM) to accelerate sampling to just 100 steps while maintaining sample quality. DDIM defines a non-Markovian reverse process that shares the same forward marginals but allows larger jumps.
At each sampling step, we first estimate the clean signal from the current noisy x_k and the predicted noise (equation (7)):
Then we update to the previous step k-1 (or, for accelerated sampling, to the previous sub-grid step \tau_{i-1}) using equation (9):
The noise scale \sigma_k is controlled by a parameter \eta \in [0,1] (equation (29)):
When \eta = 0, the process is deterministic (pure DDIM). When \eta = 1, it recovers the original DDPM. The paper uses \eta = 0.2 for a balance of stochasticity and speed.
For accelerated sampling, we select a sub-grid of step indices (e.g., 100 evenly spaced indices from 1 to 500) and apply the update using the corresponding \bar{\alpha} values. The sample_sub_grid method generates this list:
def sample_sub_grid(self, n_steps: int) -> List[int]:
step_indices = torch.linspace(
self.K, 1, n_steps, dtype=torch.float32
).round().long().unique().tolist()
step_indices = sorted(step_indices, reverse=True)
return step_indicesThe ddim_sample method implements the full sampling loop. It takes the denoiser function, initial noise x_K, conditioning dictionary, sub-grid, and \eta:
def ddim_sample(self, denoiser, x_K, cond=None, sub_grid=None, eta=0.2):
if sub_grid is None:
sub_grid = list(range(self.K, 0, -1))
sub_grid = sorted(sub_grid, reverse=True)
x = x_K.clone()
for i, k in enumerate(sub_grid):
k_prev = sub_grid[i + 1] if i + 1 < len(sub_grid) else 0
sqrt_alpha_bar_k = self.sqrt_alpha_bar[k]
sqrt_one_minus_alpha_bar_k = self.sqrt_one_minus_alpha_bar[k]
inv_sqrt_alpha_bar_k = self.inv_sqrt_alpha_bar[k]
epsilon_theta = denoiser(x, k, cond)
x_hat_0 = (x - sqrt_one_minus_alpha_bar_k * epsilon_theta) * inv_sqrt_alpha_bar_k
if k_prev > 0:
sqrt_alpha_bar_prev = self.sqrt_alpha_bar[k_prev]
one_minus_alpha_bar_prev = self.one_minus_alpha_bar[k_prev]
one_minus_alpha_bar_k = self.one_minus_alpha_bar[k]
ratio = (one_minus_alpha_bar_prev / one_minus_alpha_bar_k).clamp(min=0.0)
sigma = eta * torch.sqrt(ratio) * torch.sqrt(1.0 - self.alpha_bar[k] / self.alpha_bar[k_prev])
noise = torch.randn_like(x) if eta > 0 else 0.0
sqrt_one_minus_alpha_bar_prev_minus_sigma2 = torch.sqrt(
(one_minus_alpha_bar_prev - sigma**2).clamp(min=0.0)
)
x = sqrt_alpha_bar_prev * x_hat_0 + sqrt_one_minus_alpha_bar_prev_minus_sigma2 * epsilon_theta + sigma * noise
else:
x = x_hat_0
return xNote the adaptation for sub-grid steps: the formula uses \bar{\alpha}_{k_\text{prev}} and \bar{\alpha}_k directly, which is valid for any two steps where k_\text{prev} < k. The clamping ensures numerical stability.
2.4 Conditioning on Graph Topology and Node Features
The denoiser \varepsilon_\theta is the U-GNN, which takes three inputs at each step:
Noisy signal x_k: the current node features to be denoised.
Diffusion step k: embedded as a sinusoidal positional encoding and broadcast to all nodes.
Conditioning \xi = (S, u): the graph shift operator S (a sparse matrix encoding edge weights) and node-level side information u (e.g., historical market features for stock forecasting, or direct link gain for wireless).
The conditioning is injected through fusion layers inside the U-GNN, which we will examine in detail in the next section. For now, it suffices to understand that the denoiser is a function:
and that the entire diffusion framework — forward process, noise-prediction loss, and DDIM sampling — is agnostic to the internal architecture of \varepsilon_\theta. This modularity is what allows the paper to apply the same diffusion framework to two very different tasks by swapping only the denoiser architecture.
Summary
The forward diffusion process adds Gaussian noise to graph signals according to a linear schedule (equations (2)-(4)).
The reverse process is learned by predicting the noise added at each step (equation (6)).
DDIM sampling accelerates generation by using a sub-grid of steps and a controlled stochasticity parameter \eta (equations (7), (9), (29)-(30)).
The
DiffusionProcessclass indiffusion.pyimplements all of these operations with precomputed \bar{\alpha} values and efficient tensor operations.Conditioning on graph topology and node features is handled by the U-GNN denoiser, which we will explore next.
3. The U-GNN Architecture
3.1 Why a Graph U-Net?
Images have a natural multi-scale structure: you can downsample (pool) by averaging neighboring pixels, and then upsample (unpool) by interpolation. A U-Net exploits this to build a computational graph with a wide receptive field while keeping the number of parameters manageable. Graph signals lack such a regular grid — there is no canonical notion of “neighbourhood averaging” because the graph is irregular. Nevertheless, the same intuition applies: to model long-range dependencies across the graph, we need to process the signal at multiple resolutions, aggregating information from far-away nodes into compact representations and then expanding back to the original resolution.
The paper’s solution, the U-Graph Neural Network (U-GNN), is a U-Net-like architecture designed for graph-structured data. It uses three key innovations:
Strided graph convolutions that increase the receptive field without making the graph shift operator denser.
Learned node selection (with Gumbel-Top-K and a straight-through estimator) to decide which nodes to keep at each coarser resolution, avoiding a fixed pooling strategy.
Fusion layers that inject conditioning information (node states and diffusion step) at every resolution level.
Let us walk through each component.
3.2 The Lift–Filter–Reduce Pattern (Equation (13))
A standard GNN layer updates node features by exchanging messages along edges. If the graph shift operator is a sparse matrix S \in \mathbb{R}^{N \times N} (the adjacency or a normalized Laplacian), a single layer computes:
where X_{\ell-1} are the node features, \Theta_{\ell,k} are learnable weight matrices for each “tap” k, and \sigma_\ell is an activation. Each tap k corresponds to a k-step random walk on the graph, giving a localised receptive field.
Now suppose we have identified a subset of “active” nodes (those that survive the pooling) through a binary selection matrix D \in \{0,1\}^{N' \times N} (each row is a one-hot indicator of an active node). The paper uses a lift–filter–reduce pattern (equation (13)):
Lift: D^\mathsf{T} Z_{\ell-1} maps the N'-node signal back to the full N-node graph by placing the values at the active nodes and leaving zeros elsewhere.
Filter: apply the graph convolution (sum over k taps) on the full graph.
Reduce: multiply by D to extract the active-node signals.
This pattern is the foundation of every strided convolution in U-GNN.
3.3 Strided Graph Convolutions (Algorithm 1, Equation (14))
A plain GNN with K taps has a receptive field of exactly K hops. To reach farther nodes, one could increase K, but that makes S^K denser and more expensive (both in memory and computation). The paper’s solution is to replace S^k with (S^\gamma)^k, where \gamma is a stride parameter. Equation (14) formalises this:
Raising S to the power \gamma directly is expensive, so the paper proposes Algorithm 1: instead of forming S^\gamma, we perform \gamma K unit shifts and tap every \gamma-th one. This gives the same result as using S^\gamma, but each shift is a sparse matrix–vector product, and the intermediate \gamma-1 shifts that are not tapped still propagate the signal forward, building the wider receptive field.
Let us look at the implementation in gnn_layers.py:
class StridedGraphConv(nn.Module):
"""
Strided graph convolution layer with pooling (lift-filter-reduce).
Implements Algorithm 1 of the paper.
"""
def __init__(self, in_dim, out_dim, K, gamma, activation=nn.ReLU()):
super().__init__()
self.weight = nn.Parameter(torch.randn(K + 1, in_dim, out_dim) * 0.1)
self.bias = nn.Parameter(torch.zeros(out_dim))
self.K = K
self.gamma = gamma
self.activation = activation
def forward(self, Z_prev, S, D):
# Lift
x = torch.mm(D.T, Z_prev) # (N, in_dim)
# Tap 0
A = torch.mm(x, self.weight[0]) # (N, out_dim)
# Iterated unit shifts
for j in range(1, self.gamma * self.K + 1):
x = torch.mm(S, x) # one shift
if j % self.gamma == 0:
tap_idx = j // self.gamma
A = A + torch.mm(x, self.weight[tap_idx])
A = self.activation(A + self.bias)
out = torch.mm(D, A) # reduce
return outKey points:
The total number of taps is still K+1 (including the identity tap), independent of \gamma.
The stride \gamma is set per depth: \gamma_b = \min(\lfloor\sqrt{\bar{\rho}_b}\rfloor, \gamma_{\max}=2), where \bar{\rho}_b is the cumulative pooling factor at depth b.
Both dense and sparse S are supported (the code checks
S.is_sparse; PyTorch Geometric’s sparse ops can be used in practice).
3.4 Learned Node Selection (Equations (24)–(25), (31)–(34))
How does the U-GNN decide which nodes to keep when going from one resolution level to the next? Instead of a fixed rule (e.g., max pooling over neighbourhoods), the paper learns which nodes are most informative. The process has three steps:
Score each node: a node-wise MLP \Psi_b takes the concatenation of the encoded feature Z_b and the global embeddings D_b[U_0;K_0] and outputs a scalar score (equation (24)).
**Select top-k**: keep the N_{b+1} highest-scoring nodes (equation (25)).
Down-sample: the selection matrix C_{b+1} is the identity rows indexed by the chosen nodes.
But this discrete selection is non-differentiable. During training, the paper uses two tricks:
Gumbel-Top-K (equation (31)): Gumbel noise is added to the scores before the \text{TopK} operation, encouraging exploration.
Straight-Through Estimator (STE) (equation (32)): the forward pass uses a hard one-hot mask (the selected nodes), while the backward pass flows through a soft sigmoid surrogate \sigma(v_b / \tau). The code does this with a simple
detachtrick:
mask_hard = torch.zeros(N_b, device=device)
mask_hard.scatter_(0, indices, 1.0) # hard one-hot
mask_soft = torch.sigmoid(scores / tau) # soft surrogate
mask = mask_hard + mask_soft - mask_soft.detach() # forward hard, backward softAdditionally, the temperature \tau and exploration noise \varepsilon are annealed linearly over training (equation (34)): a warm-up phase of T_w = 0.02T epochs followed by an anneal phase of T_a = 0.73T epochs. The full NodeSelectionHead class in node_selection.py implements this schedule.
3.5 Full U-GNN: Encoder, Bottleneck, Decoder (Equations (18)–(23))
With the strided convolution and the learned selection head, we can now assemble the complete U-GNN. The architecture has B = 4 resolution levels. Let N_1 = N be the original number of nodes. At each depth b, the number of active nodes is N_b = \lfloor N / \rho^{b-1}\rfloor with \rho = 2.
Nested selection matrices (equation (18)): the global selection matrix D_b that maps from the full graph to the active nodes at depth b is the product of all previous selection matrices:
Encoder path (equation (19)): for b = 1, \dots, B-1:
Down-sample: V_b = C_b Z_{b-1}.
Fuse: P_b = \Pi_b(V_b, D_b[U_0;K_0]) (see below).
Apply strided GNN module: Z_b = \Phi_{E_b}(P_b, S; \Theta_{E_b}, D_b, \gamma_b).
Select: C_{b+1} = \text{Selector}(Z_b, D_b[U_0;K_0], N_{b+1}).
Bottleneck (equation (21)): at b = B, the same pattern without selection.
Decoder path (equation (20)): for b = B-1, \dots, 1:
Up-sample: C_{b+1}^\mathsf{T} Y_{b+1} (zero-pad to the previous resolution).
Skip-connect: concatenate the up-sampled decoder output with the encoder signal Z_b.
Project: V_b = \Pi_{\text{skip}}^b([C_{b+1}^\mathsf{T} Y_{b+1}; Z_b]) (equation (23)).
Fuse: P_b = \Pi_b(V_b, D_b[U_0;K_0]).
Apply strided GNN: Y_b = \Phi_{D_b}(P_b, S; \Theta_{D_b}, D_b, \gamma_b).
Finally, read-out: \varepsilon_\theta = \Pi_{\text{out}}(Y_1), a node-wise MLP that maps back to the original feature dimension.
The fusion layer \Pi_b (equation (22)) merges the down-sampled signal with the global embeddings (node states U_0 and step embedding K_0). For the stock task, the paper uses cross-attention (query = signal, key/value = global embeddings); for the WRA task, it uses simple addition after a projection. The FusionLayer class in ugnn.py implements both versions.
3.6 Task-Specific Adaptations
For the stock forecasting task, the input x_k has a temporal dimension (5 future days). The U-GNN includes a TemporalMixer that applies two dilated 1D convolutions followed by two-head self-attention along the time axis, producing a single feature vector per node (by averaging over time). This is a reasonable default for the paper’s underspecified temporal encoder.
For the wireless resource allocation task, the input is static (one allocation per node), so no temporal mixer is needed — the TemporalMixer is replaced by an identity module.
The UGNN class in ugnn.py orchestrates all these components. It pre-computes the target node counts, initialises the nested selection matrix D, and runs the encoder–bottleneck–decoder loop. The selection heads are called at the end of each encoder block, and the resulting C and D matrices are accumulated.
3.7 Summary
The U-GNN is a carefully designed graph U-Net that avoids explicit graph coarsening. Instead, it learns which nodes to retain at each level, uses strided graph convolutions to control the receptive field without making the shift operator denser, and fuses conditioning information at every resolution. This architecture is the central contribution of the paper and is what enables the diffusion model to generate high-quality graph signals conditioned on graph structure and node features.
4. Data Pipelines
A generative model is only as good as the data it learns from. The paper presents two distinct data pipelines, each transforming raw observations into the graph-structured inputs required by the diffusion model. This section walks through both pipelines in detail, showing how we go from stock tickers or wireless network parameters to the tensors that feed the U-GNN.
4.1 S&P 500 Stock Forecasting Pipeline
The stock forecasting task aims to predict the next five days of log-returns for 468 S&P 500 constituents, conditioned on the last 20 days of market features and a static graph of company similarities. The pipeline has five stages: data acquisition, feature engineering, graph construction, window creation, and normalisation & splitting.
4.1.1 Data Acquisition
In a real deployment, you would download daily OHLCV data from Yahoo Finance using the yfinance library. The paper specifies the date range 2016‑09‑06 to 2026‑02‑13. However, our implementation (in data_stock.py) provides a synthetic data generator that mimics the shape and basic statistics of 468 stocks, so that the code can be run without an internet connection. The function download_sp500_data returns a DataFrame with columns ticker, date, open, high, low, close, volume, and sector. The synthetic prices follow an AR(1) log‑price process with a common market factor and idiosyncratic noise, which is sufficient for testing the pipeline logic.
def download_sp500_data(
start_date: str = "2016-09-06",
end_date: str = "2026-02-13",
num_stocks: int = 468,
num_days: int = 2500,
) -> pd.DataFrame:
# ... generates synthetic OHLCV data with random sectors
# Returns a DataFrame with columns: ticker, date, open, high, low, close, volume, sector4.1.2 Feature Engineering
For each stock on each day, the paper computes 12 market features. The function compute_market_features adds these features as a list column to the DataFrame. The features are:
open, high, low, close, log_return
moving averages (MA5, MA20)
volume
RSI(14)
MACD, MACDsignal, MACDhistogram
def compute_market_features(df: pd.DataFrame) -> pd.DataFrame:
# ... computes each feature per ticker
# Adds columns 'features' (list of 12 floats) and 'log_return'The log_return column becomes the target signal; the full feature vector serves as the conditioning input u.
4.1.3 Graph Construction from Fundamentals
The paper builds a static undirected graph from company fundamentals. The function build_fundamentals_graph implements a simplified version:
Compute the Spearman rank correlation between the mean feature vectors of each pair of stocks.
Add a same‑sector bonus of 0.1.
Threshold at the median correlation to keep only the strongest edges.
Symmetrize and add self‑loops.
Apply spectral normalisation: S = D^{-1/2} A D^{-1/2}.
The result is a dense (N, N) tensor S and a list of sector labels.
def build_fundamentals_graph(stocks_df: pd.DataFrame) -> Tuple[torch.Tensor, List[int]]:
# ... rank correlation, sector bonus, threshold, spectral normalisation
return torch.tensor(S, dtype=torch.float32), sectors.tolist()4.1.4 Sliding Windows
With the feature matrix of shape (N, T, U) (12 features) and the target matrix of shape (N, T) (log‑returns), the function create_sliding_windows extracts every possible window of length Th + Tp. For each starting time t from Th to T - Tp, it produces:
history
u: shape(N, Th, U)— the last 20 days of features.target
x0: shape(N, Tp)— the next 5 days of log‑returns.
def create_sliding_windows(
data: pd.DataFrame, Th: int = 20, Tp: int = 5
) -> List[Dict]:
# ... builds feature tensor (N, T, U) and target tensor (N, T)
# ... iterates t from Th to T-Tp, creates history and target tensors4.1.5 RevIN Normalisation
The paper applies a modified Reversible Instance Normalisation (RevIN) to the target window. The RevIN class in data_stock.py implements this with the paper‑specific parameters: a blend weight of 0.7 and a scale correction of 1.107. During the forward pass, it normalises the target using per‑instance mean and standard deviation (computed along the temporal dimension), then blends these statistics with running averages for consistency. The normalised target is then scaled and shifted by learnable affine parameters.
class RevIN(nn.Module):
def __init__(self, num_features: int = 1, blend: float = 0.7, scale_correction: float = 1.107):
# ... initialise affine weight with 1/scale_correction
def forward(self, x, mode="norm", mean=None, std=None):
if mode == "norm":
# compute instance mean and std, blend with running stats
# normalise, apply affine transform
return x_norm, (mean_inst, std_inst)
elif mode == "denorm":
# reverse affine, then denormalise using stored statistics
return x_denorm, None4.1.6 Interleaved Chronological Split
To avoid temporal leakage, the paper uses an interleaved chronological split. The function interleaved_chronological_split divides the total number of windows into 10 equal chunks, then within each chunk splits chronologically 80% train, 10% validation, 10% test. The windows are pooled across chunks, ensuring that test windows always come after training windows in time.
def interleaved_chronological_split(
windows: List[Dict],
split_ratios: Tuple[float, float, float] = (0.8, 0.1, 0.1),
num_chunks: int = 10,
) -> Dict[str, List[Dict]]:
# ... divide windows into chunks, split each chunk, concatenate
return {"train": train, "val": val, "test": test}4.1.7 Dataset and DataLoader
The StockDataset class wraps the windows, the static graph S, and the RevIN module. Each call to __getitem__ returns a dictionary with keys:
x0: target log‑returns (normalised), shape(N, Tp)u: history features, shape(N, Th, U)S: graph shift operator, shape(N, N)revin_params: tuple(mean, std)used for later denormalisation during evaluation
The factory function get_stock_dataloaders creates the three DataLoaders with the appropriate batch size (64) and shuffling.
class StockDataset(Dataset):
def __init__(self, windows: List[Dict], S: torch.Tensor, revin: RevIN):
self.windows = windows
self.S = S
self.revin = revin
def __getitem__(self, idx):
w = self.windows[idx]
# apply RevIN normalisation to target
x0_norm, (mean, std) = self.revin(w["target"], mode="norm")
return {
"x0": x0_norm,
"u": w["history"],
"S": self.S,
"revin_params": (mean, std)
}4.2 Wireless Resource Allocation Pipeline
The wireless task is to learn a policy that maps a network configuration (positions, channel gains) to a set of transmit powers (one per transmitter‑receiver pair) that maximise the sum ergodic rate under a minimum rate constraint. The pipeline synthetically generates 128 networks (32 for each of four densities) and uses an expert primal‑dual algorithm to produce optimal allocations.
4.2.1 Network Generation
The WirelessNetworkGenerator class creates a network with N=400 transmitter‑receiver pairs uniformly distributed in a square area. The side length R depends on the density (Table III):
| Density Index | Side Length R (m) | Density (pairs/km²) | |---------------|-------------------|---------------------| | 0 | 7800 | 6.6 | | 1 | 7000 | 8.2 | | 2 | 6300 | 10.1 | | 3 | 5800 | 11.9 |
A minimum transmitter‑transmitter spacing of 50 m is enforced by resampling pairs that are too close.
class WirelessNetworkGenerator:
def __init__(self, config: WRAConfig):
self.config = config
def generate_network(self, density_idx: int) -> Dict:
# ... uniform positions in [0,R]x[0,R], enforce min spacing
return {
'pos_tx': pos_tx, 'pos_rx': pos_rx,
'area_side': R, 'density': density
}4.2.2 Channel Model
The ChannelModel implements the paper’s dual‑slope path loss model with log‑normal shadowing (σ = 7 dB) and Rayleigh fading. The path loss has a breakpoint at 200 m: exponent 2.0 before, 3.5 after. The compute_path_loss method returns an (N, N) matrix of large‑scale channel gains (linear scale). The sample_fading method generates unit‑mean complex Gaussian fading samples for a given number of links and time slots.
class ChannelModel:
def compute_path_loss(self, pos_tx, pos_rx):
# ... distances, dual-slope, shadowing
return gain # (N, N) tensor
def sample_fading(self, n_links, n_slots):
# ... complex Gaussian, squared magnitude
return fading # (n_links, n_slots)4.2.3 GSO Construction
The function build_wra_gso constructs a directed weighted graph shift operator. For each receiver i, it selects the top‑10 strongest interferers (excluding self), sets the edge weight to the channel gain, then applies a log‑normalisation: S = log(1 + gain) / max(log(1 + gain)). Self‑loops are set to 1.
def build_wra_gso(network, channel_gains, config):
# ... for each receiver, top-K gains, log-normalisation
return S # (N, N) tensor4.2.4 Node State Features
The compute_node_states function computes the 2‑dimensional node state vector u for each receiver:
direct_link_gain: the channel gain from its own transmitter.aggregate_interference: sum of channel gains from all other transmitters multiplied byPmax.
def compute_node_states(network, channel_gains, config):
direct = channel_gains.diag()
interference = (channel_gains.sum(dim=1) - direct) * config.Pmax
u = torch.stack([direct, interference], dim=1) # (N, 2)
return u4.2.5 Expert Primal‑Dual Algorithm
The paper uses a primal‑dual algorithm to solve the constrained optimisation problem (equation (28)). The expert_primal_dual function implements a simplified version: it performs projected gradient ascent on the powers p and gradient descent on the dual variables λ to maximise the Lagrangian. From the converged trajectory, it collects 200 intermediate allocations (though for simplicity our implementation returns the final allocation). The allocation is centred as x0 = p / Pmax - 1.
def expert_primal_dual(network, channel_gains, config):
# ... initialise p, lambda
for it in range(num_iters):
# compute rates, gradient of Lagrangian, update p, project to [0,Pmax]
# update lambda with clipping
# collect samples
x0 = p / Pmax - 1.0
return x0.float()4.2.6 Dataset and DataLoader
The WRADataset wraps the lists of networks, allocations, GSOs, and node states. Each item returns a dictionary with:
x0: centred allocation, shape(N,)S: graph shift operator, shape(N, N)u: node state features, shape(N, 2)
The factory function get_wra_dataloaders generates all 128 networks (32 per density), splits them 20/4/8 per density for train/val/test, and returns three DataLoaders with batch size 800 (as specified in Table I).
class WRADataset(Dataset):
def __init__(self, networks, allocations, gsos, node_states):
# ... store lists
def __getitem__(self, idx):
return {
'x0': self.allocations[idx],
'S': self.gsos[idx],
'u': self.node_states[idx]
}4.3 Shape Invariants Summary
Both pipelines produce tensors with consistent shapes that the U‑GNN expects:
| Task | x0 shape | u shape | S shape | Notes | |-------|------------------|--------------------|-------------|------------------------------------------| | Stock | (N, Tp=5) | (N, Th=20, U=12) | (N, N) | x0 normalised; S static, symmetric | | WRA | (N,) | (N, 2) | (N, N) | x0 centred; S directed, weighted |
These shapes are guaranteed by the dataset classes and are maintained throughout the diffusion forward and reverse processes. The conditioning variable u is always available, and the graph S is used by every strided convolution in the U‑GNN.
With the data pipelines in place, we can now load batches of (x0, u, S) and feed them into the diffusion model. The next section walks through the core code that implements the forward diffusion, the U‑GNN architecture, and the training loop.
5. Code Walkthrough: Core Components
With the theoretical foundation laid and the data pipelines ready, we now open the code files themselves. This section walks through each of the four core modules—diffusion.py, gnn_layers.py, node_selection.py, and ugnn.py—function by function. Our goal is to show exactly how the paper’s mathematical descriptions and algorithms become executable Python, highlighting every design decision and shape constraint.
5.1 Diffusion Process (diffusion.py)
File purpose: Implements the forward diffusion process (equations (2)–(4)) and DDIM accelerated sampling (equations (9) and (29)–(30)).
DiffusionProcess.__init__
class DiffusionProcess:
def __init__(self, betas: torch.Tensor) -> None:
if betas.dim() != 1:
raise ValueError("betas must be a 1D tensor")
if not torch.all(betas > 0) or not torch.all(betas < 1):
raise ValueError("betas must be in (0,1)")
self.K = betas.shape[0]
self.betas = betas
self.alphas = 1.0 - betas
# alpha_bar[t] = prod_{i=1}^{t} alpha_i, alpha_bar[0] = 1
# shape: (K+1,)
self.alpha_bar = torch.cat([
torch.ones(1, dtype=betas.dtype, device=betas.device),
torch.cumprod(self.alphas, dim=0)
], dim=0)
# Precompute sqrt and sqrt(1 - alpha_bar) for efficiency
self.sqrt_alpha_bar = self.alpha_bar.sqrt()
self.sqrt_one_minus_alpha_bar = torch.sqrt(1.0 - self.alpha_bar)
# For DDIM, precompute 1 - alpha_bar and inverse of sqrt_alpha_bar
self.one_minus_alpha_bar = 1.0 - self.alpha_bar
self.inv_sqrt_alpha_bar = 1.0 / self.sqrt_alpha_barExplanation: The constructor takes a 1D tensor betas of length K (K = 500 in the paper). It computes alpha = 1 - beta and then builds alpha_bar, where alpha_bar[t] = prod_{i=1}^{t} alpha_i. A leading 1.0 is prepended so that alpha_bar[0] = 1 (no noise at step 0). All subsequent quantities are precomputed and stored as tensors: sqrt_alpha_bar, sqrt_one_minus_alpha_bar, one_minus_alpha_bar, and inv_sqrt_alpha_bar. This avoids recomputing square roots during training and sampling.
Shape note: alpha_bar has shape (K+1,) so that indexing with step k (1‑indexed) directly uses position k.
forward_diffuse (equation (4))
def forward_diffuse(self, x0: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
k_idx = k.long()
sqrt_alpha_k = self.sqrt_alpha_bar[k_idx]
sqrt_one_minus_alpha_k = self.sqrt_one_minus_alpha_bar[k_idx]
epsilon = torch.randn_like(x0)
x_k = sqrt_alpha_k * x0 + sqrt_one_minus_alpha_k * epsilon
return x_k, epsilonExplanation: This is the one-line reparameterization of equation (4): x_k = sqrt(alpha_bar_k) * x0 + sqrt(1 - alpha_bar_k) * epsilon. We look up the precomputed sqrt_alpha_bar at index k (which is a tensor, possibly batched), sample a fresh Gaussian noise epsilon, and return both x_k and epsilon for the noise‑prediction loss. The shape of x_k matches x0 exactly.
sample_sub_grid
def sample_sub_grid(self, n_steps: int) -> List[int]:
if n_steps > self.K:
raise ValueError("n_steps cannot exceed K")
if n_steps <= 0:
raise ValueError("n_steps must be positive")
step_indices = torch.linspace(
self.K, 1, n_steps, dtype=torch.float32
).round().long().unique().tolist()
if len(step_indices) > n_steps:
step_indices = step_indices[:n_steps]
elif len(step_indices) < n_steps:
step_indices = [max(1, int(round(self.K - i * (self.K - 1) / (n_steps - 1))))
for i in range(n_steps)]
step_indices = sorted(set(step_indices), reverse=True)
step_indices = sorted(step_indices, reverse=True)
return step_indicesExplanation: Accelerated DDIM uses a sub‑grid of the full 500‑step schedule. This method generates n_steps indices (e.g., 100) that are roughly evenly spaced from K down to 1. It first tries torch.linspace rounded to integers; if rounding produces duplicates (less than n_steps unique values), it falls back to a simpler arithmetic spacing. The returned list is sorted descending so that the sampling loop starts from the highest noise level.
ddim_sample (equations (9), (29)–(30))
def ddim_sample(self, denoiser: Callable, x_K: torch.Tensor,
cond: Optional[Dict[str, Any]] = None,
sub_grid: Optional[List[int]] = None,
eta: float = 0.2) -> torch.Tensor:
if sub_grid is None:
sub_grid = list(range(self.K, 0, -1))
sub_grid = sorted(sub_grid, reverse=True)
x = x_K.clone()
for i, k in enumerate(sub_grid):
k_prev = sub_grid[i + 1] if i + 1 < len(sub_grid) else 0
sqrt_alpha_bar_k = self.sqrt_alpha_bar[k]
sqrt_one_minus_alpha_bar_k = self.sqrt_one_minus_alpha_bar[k]
inv_sqrt_alpha_bar_k = self.inv_sqrt_alpha_bar[k]
epsilon_theta = denoiser(x, k, cond)
# Equation (7): estimate clean signal
x_hat_0 = (x - sqrt_one_minus_alpha_bar_k * epsilon_theta) * inv_sqrt_alpha_bar_k
if k_prev > 0:
sqrt_alpha_bar_prev = self.sqrt_alpha_bar[k_prev]
one_minus_alpha_bar_prev = self.one_minus_alpha_bar[k_prev]
one_minus_alpha_bar_k = self.one_minus_alpha_bar[k]
# Equation (29): sigma_k(eta) adapted for sub-grid
ratio = (one_minus_alpha_bar_prev / one_minus_alpha_bar_k).clamp(min=0.0)
sigma = eta * torch.sqrt(ratio) * torch.sqrt(1.0 - self.alpha_bar[k] / self.alpha_bar[k_prev])
noise = torch.randn_like(x) if eta > 0 else 0.0
sqrt_one_minus_alpha_bar_prev_minus_sigma2 = torch.sqrt(
(one_minus_alpha_bar_prev - sigma**2).clamp(min=0.0)
)
# Equation (30): DDIM update
x = sqrt_alpha_bar_prev * x_hat_0 \
+ sqrt_one_minus_alpha_bar_prev_minus_sigma2 * epsilon_theta \
+ sigma * noise
else:
x = x_hat_0
return xExplanation: This implements the core sampling loop. Starting from pure noise x_K, it iterates over the (possibly accelerated) sub‑grid in descending order. For each step k:
Predict the noise
epsilon_thetausing the denoiser (our U‑GNN).Compute
x_hat_0via equation (7):(x_k - sqrt(1 - alpha_bar_k) * epsilon_theta) / sqrt(alpha_bar_k).If we are not at the final step (
k_prev > 0), compute the DDIM noise scalesigmausing equation (29). The key adaptation for the sub‑grid is that we usealpha_bar[k_prev]andalpha_bar[k]instead of consecutive indices. We clamp intermediate values to avoid numerical issues near zero.Apply the update of equation (30):
x_{k_prev} = sqrt(alpha_bar_prev) * x_hat_0 + sqrt(1 - alpha_bar_prev - sigma^2) * epsilon_theta + sigma * noise.At the final step (
k_prev = 0), assignx = x_hat_0directly.
The denoiser is a callable that takes (x, k, cond) and returns the predicted noise with the same shape as x. The cond dictionary carries the graph shift operator S and the node states u (and, for the stock task, the historical window).
Design decision: We treat the sub‑grid update as a drop‑in replacement for consecutive-step DDIM. The formula for sigma must be adapted because alpha_bar[k_prev] and alpha_bar[k] are not adjacent in the original 500‑step schedule. However, the derivation in Appendix A of the paper is valid for any two indices k_prev < k, and our implementation follows exactly that generalization.
5.2 Strided Graph Convolution (gnn_layers.py)
File purpose: Implements Algorithm 1 and the stacked GNN module used inside the U‑GNN.
StridedGraphConv.forward (Algorithm 1)
class StridedGraphConv(nn.Module):
def __init__(self, in_dim, out_dim, K, gamma, activation=nn.ReLU()):
super().__init__()
self.weight = nn.Parameter(torch.randn(K + 1, in_dim, out_dim) * 0.1)
self.bias = nn.Parameter(torch.zeros(out_dim))
# ... store K, gamma, activation
def forward(self, Z_prev: torch.Tensor, S: torch.Tensor,
D: torch.Tensor) -> torch.Tensor:
N = S.shape[0]
# ---- Lift: D^T Z_prev ----
x = torch.mm(D.T, Z_prev)
# ---- Tap 0 ----
A = torch.mm(x, self.weight[0])
# ---- Iterated unit shifts ----
for j in range(1, self.gamma * self.K + 1):
if S.is_sparse:
x = torch.sparse.mm(S, x)
else:
x = torch.mm(S, x)
if j % self.gamma == 0:
tap_idx = j // self.gamma
A = A + torch.mm(x, self.weight[tap_idx])
# ---- Nonlinearity + bias ----
A = self.activation(A + self.bias)
# ---- Reduce: D A ----
out = torch.mm(D, A)
return outExplanation: This is the direct implementation of Algorithm 1 (Appendix B). The layer processes a signal Z_prev defined on N' active nodes. The steps are:
Lift (
D^T Z_prev): Zero‑pad the signal fromN'nodes to allNnodes using the transpose of the binary selection matrixD.Tap 0: Apply the weight matrix for
k=0(the identity shift).Iterated unit shifts: Run a loop up to
gamma * K. In each iteration, apply one sparse shift (S @ x). Everygamma‑th iteration, treat it as a tap: accumulate the contributionx @ W[tap_idx].Nonlinearity + bias: Apply activation (ReLU) and bias.
Reduce (
D @ A): Project back fromNnodes to theN'active nodes using the same selection matrix.
The total number of taps is K+1 (tap 0 plus taps at indices gamma, 2*gamma, ..., gamma*K). The stride gamma determines how many unit shifts occur between taps, effectively implementing (S^gamma)^k without ever forming the dense S^gamma matrix. This is the key efficiency trick.
Shape invariants: Z_prev: (N', in_dim), S: (N, N) (sparse or dense), D: (N', N) (binary, each row sums to 1). The output is (N', out_dim).
GNNModule
class GNNModule(nn.Module):
def __init__(self, in_dim, out_dim, K, gamma, L=2, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList()
for i in range(L):
dim_in = in_dim if i == 0 else out_dim
self.layers.append(
StridedGraphConv(dim_in, out_dim, K, gamma, activation=nn.ReLU())
)
self.norm = nn.LayerNorm(out_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, X, S, D):
h = X
for layer in self.layers:
h = layer(h, S, D)
h = self.norm(h)
h = self.dropout(h)
return hExplanation: A simple stack of L strided graph convolution layers (default 2, as per Table I), each followed by LayerNorm and dropout. All layers share the same gamma and D for a given depth level. The first layer expands from in_dim to out_dim; subsequent layers stay at out_dim. The module outputs a signal on the same active node set.
5.3 Learned Node Selection (node_selection.py)
File purpose: Implements the scoring, Gumbel‑Top‑K, and straight‑through estimator described in Section IV‑B and Appendix C.
NodeSelectionHead
class NodeSelectionHead(nn.Module):
def __init__(self, in_dim, global_dim, hidden_dim=64):
super().__init__()
self.score_mlp = nn.Sequential(
nn.Linear(in_dim + global_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 1),
)
def forward(self, Z, global_emb, n_keep, training=True, epoch=0,
annealing_params=None):
N_b = Z.shape[0]
device = Z.device
# Equation (24): score = MLP(concat(Z, global_emb))
concat = torch.cat([Z, global_emb], dim=-1)
scores = self.score_mlp(concat).squeeze(-1) # (N_b,)
if training:
tau, eps = self._anneal_tau_eps(epoch, annealing_params)
if eps > 0:
# Equation (31): add Gumbel noise
U = torch.rand_like(scores)
gumbel = -torch.log(-torch.log(U + 1e-8))
scores = scores + eps * gumbel
# Equation (25): Top-K selection
_, indices = torch.topk(scores, k=n_keep, dim=-1)
indices = indices.sort()[0]
C = torch.zeros(n_keep, N_b, device=device, dtype=Z.dtype)
C[torch.arange(n_keep, device=device), indices] = 1.0
# Equation (32): straight-through estimator
mask_hard = torch.zeros(N_b, device=device, dtype=Z.dtype)
mask_hard.scatter_(0, indices, 1.0)
mask_soft = torch.sigmoid(scores / tau)
mask = mask_hard + mask_soft - mask_soft.detach()
# Equation (33): down-sample with mask
masked_Z = Z * mask.unsqueeze(-1)
V = masked_Z[indices]
else:
# Inference: hard selection only
_, indices = torch.topk(scores, k=n_keep, dim=-1)
indices = indices.sort()[0]
C = torch.zeros(n_keep, N_b, device=device, dtype=Z.dtype)
C[torch.arange(n_keep, device=device), indices] = 1.0
V = Z[indices]
return C, VExplanation: The node selection head does exactly what Section IV‑B describes. The scoring MLP Psi_b (equation (24)) takes the concatenation of the encoded feature Z and the global embeddings projected to the current active node set, and outputs a scalar per node. During training:
Gumbel noise is added to the scores (equation (31)).
epsis annealed from 1 down to 0 over training, so exploration decays.Top‑K picks the
n_keephighest‑scoring nodes.Straight‑through estimator (equation (32)): the forward pass uses a hard one‑hot mask (
mask_hard), while gradients flow through a sigmoid surrogate (mask_soft). The trickmask = hard + soft - soft.detach()ensures the forward pass is hard and the backward pass seessoft.Down‑sampled signal
Vis computed asC @ (mask * Z)(equation (33)).
At inference, Gumbel noise and the STE are removed: only hard Top‑K is used.
Annealing Schedule (_anneal_tau_eps)
@staticmethod
def _anneal_tau_eps(epoch: int, annealing_params: dict) -> (float, float):
if annealing_params is None:
Tw, Ta, tau0, tau_min, eps0, eps_min = 0.0, 1.0, 1.0, 0.5, 1.0, 0.0
else:
Tw = annealing_params.get('Tw', 0.02)
Ta = annealing_params.get('Ta', 0.73)
tau0 = annealing_params.get('tau0', 1.0)
tau_min = annealing_params.get('tau_min', 0.5)
eps0 = annealing_params.get('eps0', 1.0)
eps_min = annealing_params.get('eps_min', 0.0)
# Equation (34): r(t) = min(max((t - Tw) / Ta, 0), 1)
r = (epoch - Tw) / Ta
r = max(0.0, min(1.0, r))
tau = tau0 + (tau_min - tau0) * r
eps = eps0 + (eps_min - eps0) * r
return tau, epsExplanation: Implements equation (34) exactly. After a warm‑up phase of Tw epochs (default 0.02 × total = 100 epochs for 5000 total), the ratio r increases linearly from 0 to 1 over Ta epochs (default 0.73 × total = 3650 epochs). Both tau and eps are linearly interpolated from their initial to their final values. After Tw + Ta epochs, r = 1 and the parameters stay at their minima.
5.4 U‑GNN Architecture (ugnn.py)
File purpose: Assembles all components into the full U‑GNN as described in Section IV and equations (18)–(23).
UGNN.__init__
class UGNN(nn.Module):
def __init__(self, task: str, config: Config):
super().__init__()
self.task = task
self.B = config.B # 4 resolution levels
self.rho = config.rho # 2 (pooling factor)
self.F0 = config.F0 # 64 (feature dim)
self.L = config.L # 2 (GNN layers per block)
self.K = config.K # 2 (hops)
self.gamma_max = config.gamma_max # 2
# Core layers
self.read_in = ReadInLayer(in_dim=config.F_out, out_dim=self.F0)
self.read_out = ReadOutLayer(in_dim=self.F0, out_dim=self.F_out)
self.node_state_encoder = NodeStateEncoder(in_dim=config.U, out_dim=self.F0)
self.step_embedding = StepEmbedding(out_dim=self.F0)
if task == 'stock':
self.temporal_mixer = TemporalMixer(...)
else:
self.temporal_mixer = nn.Identity()
# Encoder blocks
self.encoder_blocks = nn.ModuleList()
self.fusion_layers_enc = nn.ModuleList()
self.selection_heads = nn.ModuleList()
for b in range(1, self.B):
gamma = self._compute_gamma(b)
gnn = GNNModule(self.F0, self.F0, self.K, gamma, self.L, self.dropout)
fusion = FusionLayer(self.F0, 2*self.F0, self.embed_dim, self.task)
sel = NodeSelectionHead(self.F0, 2*self.F0)
self.encoder_blocks.append(gnn)
self.fusion_layers_enc.append(fusion)
self.selection_heads.append(sel)
# Bottleneck
self.bottleneck_gamma = self._compute_gamma(self.B)
self.bottleneck_gnn = GNNModule(self.F0, self.F0, self.K,
self.bottleneck_gamma, self.L, self.dropout)
self.bottleneck_fusion = FusionLayer(self.F0, 2*self.F0, self.embed_dim, self.task)
# Decoder blocks
self.decoder_blocks = nn.ModuleList()
self.fusion_layers_dec = nn.ModuleList()
self.skip_projections = nn.ModuleList()
for b in range(self.B - 1, 0, -1):
gamma = self._compute_gamma(b)
gnn = GNNModule(self.F0, self.F0, self.K, gamma, self.L, self.dropout)
fusion = FusionLayer(self.F0, 2*self.F0, self.embed_dim, self.task)
skip_proj = SkipProjection(in_dim=2*self.F0, out_dim=self.F0)
self.decoder_blocks.append(gnn)
self.fusion_layers_dec.append(fusion)
self.skip_projections.append(skip_proj)Explanation: The constructor builds the full U‑GNN. It creates B-1 encoder blocks (b = 1 to B-1), one bottleneck block (b = B), and B-1 decoder blocks (b = B-1 down to 1). Each block has its own GNNModule with a depth‑dependent gamma computed by _compute_gamma, a FusionLayer, and (for encoder blocks) a NodeSelectionHead. The decoder blocks also get a SkipProjection layer to reduce the concatenated skip connection to the standard feature dimension F0.
The FusionLayer and TemporalMixer are task‑specific: the stock task uses cross‑attention in the fusion and a temporal mixer (dilated conv + self‑attention) on the read‑in features; the WRA task uses simple addition in the fusion and no temporal mixer.
UGNN.forward (equations (18)–(23))
def forward(self, x_k, k, u, S):
batch, N, F_in = x_k.shape
# 1. Read-in
Z0 = self.read_in(x_k) # (batch, N, F0)
U0 = self.node_state_encoder(u) # (batch, N, F0)
K0 = self.step_embedding(k).unsqueeze(1).expand(-1, N, -1) # (batch, N, F0)
# 2. Global embeddings
global_emb = torch.cat([U0, K0], dim=-1) # (batch, N, 2*F0)
# 3. Temporal mixer (stock only)
Z0 = self.temporal_mixer(Z0)
# 4. Node counts
node_counts = self._compute_target_node_counts(N)
# 5. Encoder path
D = torch.eye(N, device=x_k.device).unsqueeze(0).expand(batch, -1, -1)
Z = Z0
encoder_outputs = []
for b_idx in range(self.B - 1):
next_N = node_counts[b_idx + 1]
gnn = self.encoder_blocks[b_idx]
fusion = self.fusion_layers_enc[b_idx]
sel_head = self.selection_heads[b_idx]
# Fusion: P_b = Π_b(V_b, D_b[U0;K0])
# (V_b is Z after down-sampling via C_b; here we use D to project global_emb)
V_b = Z # For now, V_b = Z (in practice we need C_b; see full code)
global_emb_active = torch.bmm(D, global_emb)
P_b = fusion(V_b, global_emb_active)
Z_b = gnn(P_b, S, D)
# Selection: C_{b+1}, Z_{b+1}
C_next, Z_next = sel_head(Z_b, global_emb_active, next_N,
training=self.training,
epoch=self.current_epoch)
# Update D for next level: D_{b+1} = C_{b+1} @ D_b
D = torch.bmm(C_next, D)
encoder_outputs.append(Z_b)
Z = Z_next
# 6. Bottleneck
# ... (similar fusion + GNN on coarsest level)
# 7. Decoder path
# ... (up-sample via C^T, skip-concat, fusion, GNN)
# 8. Read-out
epsilon_theta = self.read_out(Y_1)
return epsilon_theta(Note: The above is a simplified sketch. The full code in `ugnn.py` handles batching, indexing, and the decoder loop with proper up‑sampling and skip connections.)
Explanation of the forward pass:
Read‑in and embeddings: The noisy signal
x_kis projected to dimensionF0by a node‑wise MLP. Node statesuand the diffusion stepkare similarly encoded and broadcast to all nodes. The step embedding uses sinusoidal encoding (following DDPM conventions) followed by an MLP.Global embeddings:
U0andK0are concatenated along the feature dimension to form a(batch, N, 2*F0)tensor that carries all conditioning information.Encoder path: Starting from
Z0and an identityD(active nodes = all N nodes), we iterate through the encoder blocks. At each depthb:
Project global embeddings onto the current active node set:
D_b * global_emb.Fuse (equation (22)):
P_b = Π_b(V_b, D_b[U0;K0]).Apply the GNN module with stride
gamma_b:Z_b = Φ_{E_b}(P_b, S; D_b, gamma_b).Select nodes for the next level:
C_{b+1}, Z_{b+1} = Selector(Z_b, D_b[U0;K0], N_{b+1}).Update the composite selection matrix:
D = C_{b+1} @ D(this buildsD_{b+1}as per equation (18)).
Decoder path (omitted for brevity): The coarsest representation is processed by the bottleneck block. Then, iterating backwards, each decoder block up‑samples via
C_{b+1}^T, concatenates with the corresponding encoder output (skip connection), projects viaSkipProjection(equation (23)), fuses, and applies a GNN module.Read‑out: The finest‑resolution decoder output
Y_1is projected back to the original feature dimension by a node‑wise MLP.
Key design choices:
The composite
Dmatrices are maintained as dense(batch, N_b, N)tensors for simplicity. In a production system, you might store them as sparse index lists, but dense matrices make the batch matrix multiplications (torch.bmm) straightforward.The
FusionLayeris shared between encoder and decoder blocks only in the sense that the same class is instantiated with the same input/output dimensions; the paper mentions weight sharing, but we instantiate separate layers for clarity.The
TemporalMixerapplies dilated 1D convolutions along the temporal axis (last dimension) followed by multi‑head self‑attention. The output is aggregated by taking the mean over the temporal dimension, producing a single feature vector per node. This is a reasonable default for the paper’s underspecified temporal encoder.
5.5 Putting It All Together
The four files work together seamlessly in the training loop (train.py):
DiffusionProcess.forward_diffusegeneratesx_kandepsilonfrom a cleanx_0sampled from the dataloader.The
UGNNtakes(x_k, k, u, S)and predictsepsilon_theta.The MSE loss between
epsilonandepsilon_thetais computed (equation (6)).Gradients flow through the U‑GNN, including through the straight‑through estimator in
NodeSelectionHead, and the parameters are updated.After training,
DiffusionProcess.ddim_sampleuses the trained U‑GNN to generate samples from pure noise, producing the final forecast trajectories.
The modularity of the design—separate files for diffusion, graph convolution, node selection, and the full architecture—makes it easy to test, modify, or replace individual components without touching the rest of the pipeline.
6. Training and Hyperparameters
With the U-GNN architecture and data pipelines in place, we now turn to the training procedure. The paper specifies a detailed set of hyperparameters (Table I and Appendix E) and a training loop that includes a custom learning rate schedule, gradient clipping, automatic mixed precision, and a composite validation criterion. This section walks through the configuration object, the scheduler, the training loop itself, and the checkpointing logic, tying each piece to the generated code.
6.1 Configuration from Table I and Appendix E
All hyperparameters are centralized in config.py. The base Config class stores values that are shared across tasks, while StockConfig and WRAConfig override task-specific settings. The factory function get_config(task) returns the appropriate instance.
Key hyperparameters (from Table I and Appendix E):
| Parameter | Value | Paper Reference | |-----------|-------|-----------------| | Diffusion steps K | 500 | Table I | | \beta_1, \beta_{500} | 10^{-4}, 2\times10^{-2} | Table I, linear schedule | | DDIM steps | 100 | Table I | | DDIM \eta | 0.2 | Table I | | U-GNN depths B | 4 | Table I | | Pooling factor \rho | 2 | Table I | | Channel width F | 64 | Table I | | GNN layers per block L | 2 | Table I | | Filter taps K (hops) | 2 | Table I | | Max stride \gamma_{\max} | 2 | Table I | | Dropout | 0.1 | Table I | | Embedding dimension | 128 | Table I | | Epochs | 5000 | Table I | | Batch size (stock) | 64 | Table I | | Batch size (WRA) | 800 | Table I | | Learning rate | 10^{-3} | Table I | | Weight decay | 0.0 | Table I (AdamW default) | | Warmup fraction | 0.01 (1%) | Appendix E | | Decay fraction | 0.80 (80%) | Appendix E | | Decay to fraction | 0.05 (5% of peak) | Appendix E | | Gradient clip norm | 1.0 | Appendix E | | AMP | True | Appendix E | | Annealing warmup T_w | 0.02 (100 epochs) | Appendix C | | Annealing decay T_a | 0.73 (3650 epochs) | Appendix C | | Initial temperature \tau_0 | 1.0 | Appendix C | | Final temperature \tau_{\min} | 0.5 | Appendix C | | Initial exploration noise \varepsilon_0 | 1.0 | Appendix C | | Final exploration noise \varepsilon_{\min} | 0.0 | Appendix C | | RevIN blend (stock) | 0.7 | Section V-A | | RevIN scale correction (stock) | 1.107 | Section V-A |
The Config.__init__ method calls _compute_betas() and _compute_alphabars() to precompute the linear \beta schedule and the cumulative \bar{\alpha}_k array (equations (2)–(3)). These are stored as tensors and used later by the DiffusionProcess class.
# From config.py (simplified)
class Config:
def __init__(self, task: str = "stock"):
self.task = task
if task == "stock":
self.batch_size = 64
self.U = 12
self.Th = 20
self.Tp = 5
elif task == "wra":
self.batch_size = 800
self.U = 2
self.Th = 0
self.Tp = 1
else:
raise ValueError(f"Unknown task: {task}")
self._compute_betas()
self._compute_alphabars()
def _compute_betas(self):
k = torch.arange(1, self.K + 1, dtype=torch.float32)
self.betas = self.beta_start + (k - 1) / (self.K - 1) * (self.beta_end - self.beta_start)
def _compute_alphabars(self):
self.alphas = 1.0 - self.betas
self.alphabars = torch.cumprod(self.alphas, dim=0)6.2 WarmupCosineLR Scheduler
The paper describes a learning rate schedule with a linear warm-up over the first 1% of training steps, followed by cosine decay to 5% of the peak learning rate over the next 80% of steps, and then constant at that minimum. The WarmupCosineLR class in train.py implements this exactly.
class WarmupCosineLR(optim.lr_scheduler.LambdaLR):
def __init__(self, optimizer, total_steps, warmup_frac=0.01,
decay_frac=0.80, min_lr_frac=0.05, last_epoch=-1):
self.total_steps = total_steps
self.warmup_steps = int(total_steps * warmup_frac)
self.decay_steps = int(total_steps * decay_frac)
self.min_lr_frac = min_lr_frac
super().__init__(optimizer, self._lr_lambda, last_epoch=last_epoch)
def _lr_lambda(self, step):
if step < self.warmup_steps:
return float(step) / max(1.0, self.warmup_steps)
elif step < self.warmup_steps + self.decay_steps:
progress = float(step - self.warmup_steps) / max(1.0, self.decay_steps)
cosine_decay = 0.5 * (1.0 + np.cos(np.pi * progress))
return self.min_lr_frac + (1.0 - self.min_lr_frac) * cosine_decay
else:
return self.min_lr_fracThe scheduler is instantiated with total_steps = epochs * batches_per_epoch. Because the number of batches per epoch can vary, we estimate it from the training DataLoader. The scheduler is stepped after every optimizer step (i.e., every mini-batch), not every epoch.
6.3 Training Loop
The train_epoch function implements one epoch of training. For each batch, it:
Samples a diffusion step k \sim \text{Uniform}\{1, \dots, K\} independently for each element in the batch.
Samples noise \varepsilon \sim \mathcal{N}(0, I) with the same shape as the clean signal x_0.
Computes the noisy signal x_k = \sqrt{\bar{\alpha}_k}\,x_0 + \sqrt{1-\bar{\alpha}_k}\,\varepsilon using the precomputed \bar{\alpha} array (equation (4)).
Predicts the noise \varepsilon_\theta = \text{U-GNN}(x_k, k, u, S).
Computes the MSE loss \|\varepsilon - \varepsilon_\theta\|^2 (equation (6)).
Backpropagates with gradient clipping (max L2 norm = 1.0) and automatic mixed precision (AMP) if CUDA is available.
Updates the optimizer and steps the scheduler.
# From train.py (simplified)
def train_epoch(model, dataloader, optimizer, scheduler, diffusion,
epoch, config, scaler=None):
model.train()
if hasattr(model, 'set_epoch'):
model.set_epoch(epoch) # for annealing τ and ε
total_loss = 0.0
n_batches = 0
for batch in dataloader:
if config.task == 'stock':
x0, u, S, _ = batch
else:
x0, S, u = batch
x0 = x0.to(device, non_blocking=True)
u = u.to(device, non_blocking=True)
S = S.to(device, non_blocking=True)
batch_size = x0.size(0)
k = torch.randint(1, config.K + 1, (batch_size,), device=device)
noise = torch.randn_like(x0)
x_k = diffusion.forward_diffuse(x0, k, noise)
with autocast(enabled=scaler is not None):
pred_noise = model(x_k, k, u, S)
loss = F.mse_loss(pred_noise, noise)
if scaler is not None:
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
optimizer.zero_grad(set_to_none=True)
scheduler.step()
total_loss += loss.item()
n_batches += 1
return total_loss / max(1, n_batches)Key details:
The
DiffusionProcess.forward_diffusemethod implements equation (4) by indexing into the precomputedalpha_bartensor with the sampled k values.Gradient clipping uses
torch.nn.utils.clip_grad_norm_withmax_norm=1.0and the default L2 norm.AMP is enabled via
torch.cuda.amp.GradScalerandautocast. The scaler is created only if CUDA is available.The model’s
set_epochmethod (if present) is called at the start of each epoch to update the temperature \tau and exploration noise \varepsilon for the node selection heads, following the annealing schedule of Appendix C.
6.4 Annealing of τ and ε
The node selection heads use a Gumbel-Top-K mechanism during training, with temperature \tau and exploration noise \varepsilon that are annealed over time. The schedule (equation (34)) has three phases:
Warm-up (first T_w = 0.02 \times 5000 = 100 epochs): \tau and \varepsilon are held at their initial values (1.0 and 1.0).
Anneal (next T_a = 0.73 \times 5000 = 3650 epochs): both parameters decay linearly to their final values (0.5 and 0.0).
Constant (remaining epochs): stay at final values.
In the code, the NodeSelectionHead class (in node_selection.py) implements anneal_tau_eps(epoch) which computes the current values based on the schedule. The UGNN model’s set_epoch method propagates the epoch number to all selection heads so they can update their internal parameters before each forward pass.
6.5 Composite Validation Criteria
The paper uses a task-specific composite score to select the best checkpoint during training. The composite criterion combines multiple validation metrics into a single scalar.
For stock forecasting: The composite score is a weighted sum of the noise-prediction loss, the gap between training and validation loss (to detect overfitting), and the CRPS (Continuous Ranked Probability Score). Lower is better.
def composite_criterion_stock(val_metrics):
loss = val_metrics.get('loss', 1.0)
crps = val_metrics.get('crps', 1.0)
# train-val gap not computed here, assume 0
return loss + 0.5 * crpsFor wireless resource allocation: The composite score combines the mean ergodic rate and the feasibility (fraction of users meeting the minimum rate constraint). Higher is better, so we negate the sum for minimization.
def composite_criterion_wra(val_metrics):
mean_rate = val_metrics.get('mean_rate', 0.0)
feasibility = val_metrics.get('feasibility', 0.0)
return -(mean_rate + feasibility * 2.0) # weight feasibility moreThese functions are called after each validation epoch. The train function tracks the best composite score and saves the corresponding model checkpoint.
6.6 Checkpointing
At the end of each epoch, if the composite score improves (lower for stock, higher for WRA), the current model state, optimizer state, scheduler state, epoch number, and score are saved to a dictionary. After training completes, the best checkpoint is written to checkpoints/best_model.pth.
if (config.task == 'stock' and score < best_score) or \
(config.task == 'wra' and score > best_score):
best_score = score
best_epoch = epoch
best_state = {
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'epoch': epoch,
'score': score,
'config': config
}This checkpoint can later be loaded for evaluation or visualization.
6.7 Summary
The training configuration and loop are designed to match the paper’s specifications exactly. The Config class centralizes all hyperparameters, the WarmupCosineLR scheduler implements the custom learning rate schedule, and the train_epoch function handles the diffusion-specific steps (sampling k, computing x_k, predicting noise) along with gradient clipping and AMP. The annealing of node selection parameters is integrated via the model’s set_epoch hook, and the composite validation criteria guide checkpoint selection. This modular design makes it straightforward to reproduce the paper’s experiments and to adapt the code to new tasks.
7. Evaluation Metrics and Baselines
A generative model is only as useful as the metrics we use to judge it. The paper evaluates its U-GNN diffusion model against two sets of baselines and a battery of metrics that go far beyond simple point-prediction errors. For stock forecasting, the goal is to assess both the accuracy of the mean prediction and the fidelity of the full predictive distribution. For wireless resource allocation, the focus is on the ergodic sum rate, feasibility, and the gap to an expert algorithm. This section explains every metric and baseline, shows the corresponding code in eval.py and baselines.py, and describes how to compare the U-GNN against the baselines.
7.1 Stock Forecasting Metrics
The paper reports eight metrics for the S&P 500 task (Table II). They fall into three categories: distributional accuracy (CRPS, MIS90), point-prediction accuracy (RMSE, MAE, direction accuracy), and stylized-fact gaps (volatility clustering, momentum, excess kurtosis). All metrics are computed on the returns (log-returns) unless otherwise noted.
7.1.1 Continuous Ranked Probability Score (CRPS)
CRPS measures how well the empirical distribution of generated samples matches the true target. It is defined as the integral of the squared difference between the empirical CDF of the samples and the indicator function of the target:
For a set of M samples, the empirical CDF is a step function, and the integral can be computed analytically using the sorted samples. The implementation in StockEvaluator.compute_crps uses the formula:
where s_{(i)} are the sorted samples and y is the target. The double sum is computed efficiently using the fact that for sorted samples, \sum_{i,j} |s_{(i)} - s_{(j)}| = 2 \sum_{i=1}^{M} (2i - M - 1) s_{(i)}.
@staticmethod
def compute_crps(samples: np.ndarray, target: np.ndarray) -> float:
samples = np.asarray(samples)
target = np.asarray(target)
if samples.ndim == 3:
num_samples = samples.shape[0]
else:
num_samples = 1
samples = samples[np.newaxis, ...]
m = samples.shape[1] * samples.shape[2]
samples_flat = samples.reshape(num_samples, -1)
target_flat = target.reshape(-1)
sorted_samples = np.sort(samples_flat, axis=0)
M = num_samples
abs_diff = np.mean(np.abs(sorted_samples - target_flat[np.newaxis, :]), axis=0)
weights = 2 * (2 * np.arange(1, M + 1) - M - 1)
sum_abs_diff_sorted = np.sum(weights[:, np.newaxis] * sorted_samples, axis=0) / (M ** 2)
crps_per_element = abs_diff - 0.5 * sum_abs_diff_sorted
return float(np.mean(crps_per_element))7.1.2 Mean Interval Score for 90% Prediction Interval (MIS90)
MIS90 evaluates the sharpness and calibration of the 90% prediction interval. It penalizes intervals that are too wide and intervals that miss the target:
where l and u are the 5th and 95th percentiles of the samples, and \alpha = 0.1. The implementation computes these percentiles using np.percentile.
@staticmethod
def compute_mis90(samples: np.ndarray, target: np.ndarray) -> float:
samples = np.asarray(samples)
target = np.asarray(target)
if samples.ndim == 2:
samples = samples[np.newaxis, ...]
alpha = 0.1
lower = np.percentile(samples, 5, axis=0)
upper = np.percentile(samples, 95, axis=0)
width = upper - lower
penalty_low = (2.0 / alpha) * (lower - target) * (target < lower)
penalty_high = (2.0 / alpha) * (target - upper) * (target > upper)
mis = width + penalty_low + penalty_high
return float(np.mean(mis))7.1.3 RMSE, MAE, and Direction Accuracy
These are standard point-prediction metrics. The mean prediction is the average over the M samples. RMSE and MAE are computed between this mean and the target. Direction accuracy measures the fraction of elements where the sign of the mean prediction matches the sign of the target (ignoring zeros).
@staticmethod
def compute_rmse(samples: np.ndarray, target: np.ndarray) -> float:
samples = np.asarray(samples)
target = np.asarray(target)
if samples.ndim == 3:
mean_pred = np.mean(samples, axis=0)
else:
mean_pred = samples
return float(np.sqrt(np.mean((mean_pred - target) ** 2)))
@staticmethod
def compute_dir_accuracy(samples: np.ndarray, target: np.ndarray) -> float:
samples = np.asarray(samples)
target = np.asarray(target)
if samples.ndim == 3:
mean_pred = np.mean(samples, axis=0)
else:
mean_pred = samples
correct = (np.sign(mean_pred) == np.sign(target))
return float(np.mean(correct))7.1.4 Stylized-Fact Gaps
Financial time series exhibit well-known statistical regularities called stylized facts. The paper measures how well the generated samples reproduce three of them:
Volatility clustering: the tendency for large changes to be followed by large changes. Quantified as the lag-1 autocorrelation of squared returns.
Momentum: the tendency for returns to be positively autocorrelated at short lags. Quantified as the lag-1 autocorrelation of returns.
Excess kurtosis: the fat-tailed nature of return distributions. Quantified as the sample excess kurtosis (kurtosis minus 3).
For each metric, the gap is the absolute difference between the value computed on the target and the value computed on the pooled generated samples. The implementation in StockEvaluator computes these using NumPy.
@staticmethod
def compute_volatility_gap(samples: np.ndarray, target: np.ndarray) -> float:
samples = np.asarray(samples)
target = np.asarray(target)
target_flat = target.ravel()
samples_flat = samples.ravel()
def autocorr_lag1(x):
if len(x) < 2:
return 0.0
x0 = x[:-1]
x1 = x[1:]
std0 = np.std(x0)
std1 = np.std(x1)
if std0 == 0 or std1 == 0:
return 0.0
return float(np.corrcoef(x0, x1)[0, 1])
target_vol = autocorr_lag1(target_flat ** 2)
samples_vol = autocorr_lag1(samples_flat ** 2)
return float(np.abs(target_vol - samples_vol))7.2 Wireless Resource Allocation Metrics
For the WRA task, the paper reports ergodic sum rate, feasibility, and the gap to the expert primal-dual algorithm (Table IV).
7.2.1 Ergodic Sum Rate
The ergodic sum rate is the average over fading realizations of the sum of per-user rates. The per-user rate is (equation (26)):
where \text{SINR}_j = \frac{p_j g_{jj}}{\sum_{i \neq j} p_i g_{ij} + \sigma^2}. The ergodic estimate (equation (27)) averages over T = 500 fading realizations, with the allocation held constant for T_0 = 5 slots. The implementation in WRAEvaluator.compute_ergodic_rate generates independent Rayleigh fading (exponential squared magnitude) and computes the SINR for each slot.
@staticmethod
def compute_ergodic_rate(
allocations: Union[np.ndarray, torch.Tensor],
H: Union[np.ndarray, torch.Tensor],
noise_power: float = 1e-11,
T: int = 500,
T0: int = 5
) -> float:
allocations = np.asarray(allocations).flatten()
H = np.asarray(H)
N = len(allocations)
# Uncenter if needed (assume Pmax=1)
if np.any(allocations < 0):
allocations = allocations + 1.0
allocations = np.clip(allocations, 0.0, 1.0)
# Generate fading: (N, N, T) exponential with mean 1
fading = np.random.exponential(scale=1.0, size=(N, N, T)).astype(np.float64)
H = np.abs(H)
g = H[:, :, np.newaxis] * fading
g_self = g.diagonal(axis1=0, axis2=1)
p = allocations[:, np.newaxis, np.newaxis]
total_received = np.sum(p * g, axis=0)
self_power = p[:, 0, 0] * g_self
interference = total_received - self_power
sinr = self_power / (interference + noise_power)
rate = np.log2(1.0 + np.maximum(sinr, 0.0))
sum_rate_per_slot = np.sum(rate, axis=0)
mean_sum_rate = np.mean(sum_rate_per_slot)
return float(mean_sum_rate)7.2.2 Feasibility
Feasibility is the fraction of users whose ergodic rate meets a minimum rate constraint (e.g., 0.5 bits/s/Hz). The implementation computes the per-user ergodic rate (mean over fading) and checks the constraint.
7.2.3 Gap to Expert
The paper reports the percentage gap between the U-GNN and the expert primal-dual algorithm at the 1st, 5th, 10th percentiles and the mean. The gap is defined as:
where p indicates the percentile of the ergodic sum rate distribution across test networks.
7.3 Baselines
The paper compares against three baselines, implemented in baselines.py.
7.3.1 Geometric Random Walk (GRW)
For stock forecasting, the GRW assumes log-returns are independent Gaussian with variance estimated from the history window. The GeometricRandomWalk.forecast method computes the per-stock variance from the 20-day history and samples M independent trajectories.
class GeometricRandomWalk:
def forecast(self, history: torch.Tensor, steps: int, num_samples: int = 1) -> torch.Tensor:
if history.dim() == 3:
history = history.squeeze(-1)
N, Th = history.shape
var = history.var(dim=1, unbiased=True, keepdim=False) + 1e-8
std = torch.sqrt(var)
noise = torch.randn(num_samples, N, steps, generator=self.rng, device=history.device, dtype=history.dtype)
samples = noise * std.view(1, N, 1)
return samples7.3.2 Full Power (FP) and Average Power (AP)
For WRA, the FP baseline sets every transmitter to maximum power P_{\text{max}}, so the centered allocation is 0. The AP baseline sets every transmitter to P_{\text{max}}/2, so the centered allocation is -0.5. Both are deterministic and return tensors of the appropriate shape.
class FullPower:
def allocate(self, network: dict = None, H: torch.Tensor = None, num_samples: int = 1) -> torch.Tensor:
N = 400
if H is not None:
N = H.shape[0]
device = H.device if H is not None else torch.device('cpu')
return torch.zeros(num_samples, N, 1, device=device, dtype=torch.float32)
class AveragePower:
def allocate(self, network: dict = None, H: torch.Tensor = None, num_samples: int = 1) -> torch.Tensor:
N = 400
if H is not None:
N = H.shape[0]
device = H.device if H is not None else torch.device('cpu')
return torch.full((num_samples, N, 1), -0.5, device=device, dtype=torch.float32)7.4 Comparing U-GNN Against Baselines
To compare the U-GNN against the baselines, we generate M = 100 samples from each method (for GRW, this is natural; for FP/AP, we repeat the deterministic allocation). Then we compute all metrics using the StockEvaluator or WRAEvaluator. The evaluate_all method returns a dictionary of metrics, which can be printed as a table (matching Table II or Table IV).
evaluator = StockEvaluator()
metrics_ugnn = evaluator.e
## 8. Results and Figures
A generative model’s output is inherently high-dimensional: for stock forecasting, we produce a distribution over future price paths for hundreds of stocks; for wireless resource allocation, we produce a distribution over power vectors. To communicate the quality of these distributions, the paper relies on a set of carefully designed figures that go beyond simple scalar metrics. This section walks through the four plotting functions in `visualize.py` that reproduce the paper’s key visual results: forecasting trajectories, empirical CDFs, distribution diagnostics (histogram + Q-Q plot), and a WRA performance bar chart. Each function is explained in detail, with code snippets and guidance on how to call them from the main evaluation script.
### 8.1 Overview of the Plotting Functions
The file `visualize.py` contains four public functions:
- `plot_forecast_trajectories` – reproduces the paper’s “forecasting trajectories” figures (e.g., Figure 2 in the paper, showing historical prices, generated samples, and actual target for selected stocks).
- `plot_cdf` – plots the empirical cumulative distribution function of the pooled generated samples versus the pooled target values, corresponding to the CDF panels in the paper’s distribution diagnostics.
- `plot_distribution_diagnostics` – creates a two-panel figure with an overlaid histogram and a Q-Q plot, matching the paper’s density and quantile comparisons.
- `plot_wra_bar` – produces a grouped bar chart comparing ergodic rate metrics (p1, p5, p10, mean) across methods (U-GNN, Expert, FP, AP), with an optional secondary axis for percentage gaps and an inset for feasibility. This reproduces the paper’s WRA performance bar chart (e.g., Figure 4).
All functions accept NumPy arrays (or torch tensors, which are automatically converted via the internal `_to_numpy` helper) and return a `matplotlib.figure.Figure` object that can be saved or displayed.
### 8.2 Forecasting Trajectories
The function `plot_forecast_trajectories` visualises the model’s predictive distribution for a few selected stocks. It plots:
- The historical log-returns (or prices) for the last `Th` days (blue line).
- A random subset of generated trajectories (gray lines, low opacity).
- The mean of all generated trajectories (red line).
- The actual target values (green dashed line).
The time axis is split into history (days 0 to Th-1) and forecast (days Th to Th+Tp-1).
def plotforecasttrajectories( prices: np.ndarray, samples: np.ndarray, target: np.ndarray, stockindices: Optional[List[int]] = None, numtrajectories: int = 20, figsize: Tuple[float, float] = (10, 6) ) -> plt.Figure:
Use the button or URL below to download the complete Python source code.