---
title: Conditional Spherical U-Net (CoS-UNet)
url: https://www.emergentmind.com/topics/conditional-spherical-u-net-cos-unet
type: topic
---

# Conditional Spherical U-Net (CoS-UNet)

Conditional Spherical U-Net (CoS-UNet) is a denoising architecture introduced as the core backbone of the Spherical Brownian Bridge Diffusion Model (SBDM) for individualized, vertex-wise cortical thickness forecasting on registered spherical cortical surfaces. It is designed for a setting in which the prediction target is a non-Euclidean cortical surface mesh rather than a Euclidean image, while the forecasting problem also requires conditioning on tabular clinical variables such as age, sex, diagnosis, and follow-up interval. Within SBDM, CoS-UNet serves as the denoiser \(f_\theta\), combining spherical convolutions with dense cross-attention so that geometry-aware cortical representations and longitudinal or clinical covariates are integrated at every resolution level [2509.08442].

## 1. Motivation and problem setting

CoS-UNet is introduced to address two limitations identified in prior cortical-thickness forecasting approaches. First, the cerebral cortex is a non-Euclidean manifold, so regular 2D or 1D convolutions are ill-suited for vertex-level cortical thickness maps. Second, the forecasting task requires the integration of heterogeneous information: a high-dimensional cortical surface map and low-dimensional tabular covariates. Previous diffusion approaches for cortical thickness used a 1D U-Net, which is described as too weak for vertex-wise cortical structure and as failing to exploit surface topology [2509.08442].

The architectural rationale is therefore explicitly twofold. On the one hand, a spherical U-Net is naturally suited to meshes on the sphere and preserves a geometry-aware inductive bias for cortical surfaces. On the other hand, a vanilla spherical U-Net does not natively accept the conditioning variables required for longitudinal forecasting. CoS-UNet is introduced to combine local geometric processing on the spherical mesh, multi-scale encoder-decoder feature extraction, and conditioning via cross-attention on tabular variables and time.

The reported ablations make this motivation operational rather than merely conceptual. Replacing CoS-UNet with an MLP, the original Spherical U-Net, or SiT within the same diffusion framework degrades performance, and CoS-UNet achieves the best MAE. This suggests that the model is not only a geometric reformulation of an existing U-Net, but a targeted response to the simultaneous demands of spherical structure and dense conditional control [2509.08442].

## 2. Function inside SBDM

SBDM formulates cortical-thickness forecasting as a Brownian bridge diffusion process between baseline cortical thickness and a future thickness-change map. The model works with cortical thickness change
\[
\tau_t = \tau_0 + \Delta \tau_t,
\]
where \(\tau_0\) is baseline thickness and \(\Delta \tau_t\) is the future relative change at time \(t\) [2509.08442].

The paper adopts a reverse-oriented Brownian bridge formulation. Instead of diffusing from baseline to future, it defines a stochastic bridge in which the baseline thickness is the starting point at inference and the model iteratively reconstructs the future change map \(\Delta \tau_t\). During training, both endpoints are known, so intermediate bridge states can be sampled; during inference, only baseline \(\tau_0\), time interval \(t\), and conditions \(c\) are available, and the model samples the bridge from \(\tau_0\) toward the predicted \(\Delta \tau_t\).

Within this process, CoS-UNet is the denoising network \(f_\theta\). Its role is to take an intermediate bridge state \(x_\beta\), the diffusion step \(\beta\), the time interval \(t\), and the conditioning variables \(c\), and to predict the bridge noise or residual signal required by the reverse Brownian bridge update. In other words, the architecture is not an auxiliary conditioning module attached to SBDM; it is the mechanism by which the reverse process is parameterized.

A plausible implication is that CoS-UNet inherits the generative flexibility of the Brownian bridge formulation while constraining denoising to respect cortical topology. The paper’s framing treats this as essential for forecasting dense vertex-wise cortical thickness changes rather than region-wise summaries [2509.08442].

## 3. Architectural design on spherical cortical surfaces

CoS-UNet is built on the Spherical U-Net framework and is adapted for cortical surfaces represented on icosphere or geodesic polyhedra obtained by recursive subdivision of an icosahedron and resampling to the FsAverage cortical template. Its macro-architecture follows a standard U-Net pattern: an encoder progressively downsamples the spherical surface, a bottleneck operates at the coarsest resolution, a decoder upsamples back to the original surface resolution, and skip connections link encoder and decoder stages at matching scales [2509.08442].

Each encoder stage contains two spherical residual blocks. Every residual block consists of one-hop spherical graph convolution, group normalization, and SiLU activation. The building block is explicitly described as using a one-hop spherical graph convolution as in Spherical U-Net, so computation is carried out directly on the spherical mesh and preserves neighborhood structure on the cortical surface. A generic spherical graph convolution is written as
\[
h_v' = \sigma\!\left( W_0 h_v + \sum_{u \in \mathcal{N}(v)} W_1 h_u \right),
\]
where \(\mathcal{N}(v)\) is the one-hop neighborhood on the icosphere, \(W_0\) and \(W_1\) are learnable weights, and \(\sigma\) is the nonlinearity. The exact implementation follows the Spherical U-Net operator rather than a Euclidean kernel [2509.08442].

Downsampling is performed by spherical pooling, which in the icosphere setting corresponds to inverting the subdivision hierarchy. The decoder mirrors this with transposed spherical convolutions or spherical unpooling to restore resolution. This is important because the target is a dense vertex-wise cortical thickness forecast rather than a region-level statistic.

Skip connections preserve fine-scale spatial detail by fusing encoder features into decoder stages at matching scales. The paper emphasizes their importance for cortical morphology, where local thinning patterns can be small and spatially localized. This suggests that CoS-UNet preserves the conventional U-Net argument for multiscale spatial fidelity, but transfers that argument to spherical cortical geometry rather than planar image grids [2509.08442].

## 4. Conditioning mechanism and mathematical formulation

In addition to the bridge state, CoS-UNet uses three forms of conditioning: the Brownian-bridge step index \(\beta\), the time interval \(t\), and covariates
\[
c = (a,s,d_0,d_t),
\]
including age, sex, baseline diagnosis, and optionally follow-up diagnosis. The design separates these pathways. The diffusion step \(\beta\) is embedded separately and added to latent feature maps, keeping it disentangled from clinical conditioning. By contrast, \(t\) and \(c\) are embedded jointly and injected into every residual stage via cross-attention, providing dense stage-wise guidance [2509.08442].

The paper identifies dense cross-attention as the key mechanism enabling conditioning. Conceptually, cortical surface features act as queries, while embedded tabular condition tokens provide keys and values:
\[
\mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V.
\]
Here, \(Q\) is derived from spherical feature maps, whereas \(K\) and \(V\) come from the conditioning embedding of \(t\) and \(c\). The result is that subject-specific and time-specific information modulates cortical features at each scale [2509.08442].

The Brownian bridge forward process is defined as
\[
q(x_\beta \mid x_0, x_B) = \mathcal{N}\left(x_\beta;\, (1-m_\beta)x_0 + m_\beta x_B,\; \delta_\beta I\right),
\]
with \(0 \le \beta \le B\), and the variance schedule
\[
\delta_\beta = 2(m_\beta - m_\beta^2).
\]
For cortical-thickness forecasting, the identifications are
\[
x_0 \equiv \tau_0,\qquad x_B \equiv \Delta \tau_t.
\]

At training time, intermediate bridge states are sampled as
\[
x_\beta = (1-m_\beta)\tau_0 + m_\beta \Delta\tau_t + \sqrt{\delta_\beta}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,I).
\]
The denoiser \(f_\theta\) is trained with the squared error objective
\[
\mathcal{L}_{\text{SBDM} = \left\| (1-m_\beta)(\tau_0-\Delta\tau_t) + \sqrt{\delta_\beta}\,\epsilon - f_\theta(x_\beta,\beta,t,c) \right\|^2.
\]
The paper distinguishes this from DDPM-style training because the bridge explicitly connects baseline and future change rather than starting from pure Gaussian noise [2509.08442].

## 5. Training, inference, and trajectory generation

Training uses longitudinal samples \((\tau_0, \Delta \tau_t, t, c)\). A bridge step \(\beta\) is sampled, noise \(\epsilon\) is drawn, and \(x_\beta\) is synthesized from the bridge equation. CoS-UNet then predicts the residual or noise, and the parameters are optimized with the squared error loss. A key implementation choice is that the model operates directly in cortical-thickness space rather than in a latent autoencoder space. The authors highlight this because latent diffusion can lose fine-grained anatomical detail and complicate decoding [2509.08442].

At inference, the process starts from
\[
x_0 = \tau_0,
\]
and the reverse bridge recursion is applied until \(x_B\) is reached, which is interpreted as the predicted \(\Delta \tau_t\). The forecasted cortical thickness is then
\[
\hat{\tau}_t = \tau_0 + \hat{\Delta \tau}_t.
\]
The reverse sampling distribution is written as
\[
p_\theta(x_{\beta+1}\mid x_\beta,\tau_0,t,c) = \mathcal{N}\!\left(x_{\beta+1}; \mu_\theta(x_\beta,\beta,t,c), \tilde{\delta}_\beta I\right),
\]
with recursive update
\[
x_{\beta+1} = \zeta_{1,\beta}x_\beta + \zeta_{2,\beta}\tau_0 - \zeta_{3,\beta}f_\theta(x_\beta,\beta,t,c) + \sqrt{\tilde{\delta}_\beta}\,\eta, \qquad \eta \sim \mathcal{N}(0,I).
\]
The coefficients \(\zeta_{1,\beta}, \zeta_{2,\beta}, \zeta_{3,\beta}, \tilde{\delta}_\beta\) are specified by the bridge derivation. The paper also uses a non-Markovian sampling strategy similar to DDIM/BBDM, so only a subset of bridge steps can be sampled for speed [2509.08442].

The same architecture supports both factual and counterfactual trajectory generation. For standard forecasting, \(d_t\) is not part of the condition, and the model predicts the likely future trajectory from baseline state and observed covariates. For counterfactual trajectories, the paper trains a variant that includes follow-up diagnosis \(d_t\) in \(c\). Changing \(d_t\) at inference allows the same baseline subject to be simulated under different diagnostic outcomes, such as a stable MCI trajectory versus a hypothetical progression to AD. The reported qualitative result is that switching the target diagnosis produces a marked drop in predicted cortical thickness. This suggests that CoS-UNet’s dense conditioning pathway is sufficiently expressive to alter the generative trajectory in diagnostically meaningful ways [2509.08442].

## 6. Data, implementation, and empirical performance

The training and evaluation protocol is centered on longitudinal cortical-thickness forecasting using ADNI for training and internal testing, with OASIS used exclusively for external generalization testing. Cortical thickness was extracted using longitudinal FreeSurfer v7.2, resampled or registered to the FsAverage icosphere template, and the medial wall was ignored unless otherwise noted [2509.08442].

### Dataset configuration

| Dataset | Subjects / scans | Role |
|---|---:|---|
| ADNI train | 921 subjects / 4,112 scans | Training |
| ADNI validation | 306 subjects / 1,412 scans | Validation |
| ADNI test | 306 subjects / 1,387 scans | Testing |
| OASIS | 590 subjects / 1,750 scans | Testing only |

For ADNI, the age range is 55–97 years, the cohort includes CN, MCI, and AD subjects, and longitudinal follow-up extends up to 168 months. For OASIS, the age range is 42–95 years and the dataset includes CN and AD cases.

The implementation details associated with CoS-UNet are specific. It is implemented in PyTorch 2.0.1 and uses the Sphericalunet package. The channel width is set to \(C=64\), training proceeds for up to 2,000 epochs, the Brownian bridge horizon is \(B=1000\), and only 200 intermediate bridge stages are sampled at inference. Optimization uses AdamW with learning rate \(10^{-4}\), the learning rate is reduced on plateau, and exponential moving average is used for final weights [2509.08442].

The main empirical claim is that SBDM with CoS-UNet achieves the best prediction accuracy among all compared methods.

| Setting | Method | MAE |
|---|---|---:|
| ADNI test | SBDM | 0.097 mm |
| ADNI test | SiT | 0.108 mm |
| ADNI test | Spherical U-Net | 0.111 mm |
| ADNI test | DDPM/CoS-UNet | 0.119 mm |
| ADNI test | original CTh-DDPM | 0.208 mm |
| OASIS test | SBDM | 0.100 mm |

On OASIS, SBDM is reported to outperform all baselines and to generalize without fine-tuning. The comparison between DDPM/CoS-UNet and the original CTh-DDPM is identified as especially informative: replacing the original 1D U-Net in CTh-DDPM with CoS-UNet substantially improves the diffusion baseline, indicating that the poor performance of the original CTh-DDPM was largely attributable to its denoiser rather than to the diffusion paradigm itself [2509.08442].

The denoising-model ablation on the ADNI validation set further isolates the contribution of CoS-UNet:

| Denoising model | MAE |
|---|---:|
| Spherical U-Net-based | 0.281 mm |
| MLP-based | 0.269 mm |
| SiT-based | 0.097 mm |
| CoS-UNet | 0.095 mm |

The paper interprets these results as showing that naive MLPs and even plain spherical U-Nets are insufficient within the bridge diffusion framework, that attention-based conditioning is important, and that CoS-UNet slightly but consistently outperforms the SiT-based variant. Vertex-wise error maps are described as showing that CoS-UNet-based SBDM reduces errors consistently across the cortical sheet and avoids the spatial heterogeneity seen in some alternative methods. Statistical significance is reported using paired two-sided Wilcoxon signed-rank tests with \(p < 10^{-3}\) for SBDM versus all four reference methods on both datasets [2509.08442].

Taken together, these results position CoS-UNet as the principal architectural contribution that enables SBDM to operate effectively on vertex-wise cortical thickness forecasting: it adapts the U-Net denoising paradigm to spherical cortical surfaces, preserves multiscale topology-aware processing, and injects tabular clinical and temporal conditioning through dense cross-attention.

Source: https://www.emergentmind.com/topics/conditional-spherical-u-net-cos-unet