---
title: 'PS-U-Net: Shared Encoder & Expert Decoders'
url: https://www.emergentmind.com/topics/partially-shared-u-net-ps-u-net
type: topic
---

# PS-U-Net: Shared Encoder & Expert Decoders

Partially Shared U-Net (PS-U-Net) refers to a family of deep neural architectures that blend universal and task- or modality-specific processing streams in a U-Net framework. PS-U-Nets allocate a common encoder or intermediate block for information shared across tasks or modalities, while dedicating private decoder or auxiliary branches for sources of signal divergence, such as expert label variability or distinct input modalities (e.g., image and text). This design addresses challenges in multitask learning, data annotation from multiple experts, and multimodal conditional generation by mitigating interference and capturing fine-grained specifics.

## 1. Architectural Principles of PS-U-Net

The PS-U-Net principle is characterized by the architectural split between a shared feature extraction pathway and expert- or modality-specific decoding. In biomedical segmentation settings, the encoder is implemented as a deep CNN (e.g., ResNet-50 up to conv5, omitting the final fully connected layer), operated on input \(x\in\mathbb{R}^{H\times W\times C}\) to produce a compact latent \(h=E(x;\theta)\in\mathbb{R}^{h\times w\times d}\) where typically \(h=w=H/32\) and \(d=2048\) [2108.04658]. Each decoder \(D_i\) (for the \(i\)th expert) mirrors the encoder’s depth, performing up-sampling via transpose convolutions, 3×3 convolutions, batch normalization, and ReLU activations. Skip connections link corresponding encoder and decoder stages. The model fuses predictions by summing or averaging decoder outputs. Only encoder parameters \(\theta\) are shared, while decoder parameters \(\phi_i\) are separated per expert.

In multimodal diffusion models, PS-U-Net organizes computation into three parallel towers: a modality-specific image path, a modality-specific text path, and a central shared path. Each encoder/decoder branch comprises Transformer (U-ViT) or ViT blocks, enabling separate treatment of visual and linguistic features. Cross-modal fusion is achieved at configurable depths by concatenating outputs from modality-specific branches into the shared stream, followed by joint processing blocks [2311.16488]. Dedicated skip connections preserve modality-specific detail and reduce interference, supporting joint or conditional generation while maintaining pathway-specific granularity.

## 2. Mathematical Formulation

For expert-driven biomedical segmentation, the encoder mapping is formalized as
\[
h = E(x; \theta)
\]
where \(E\) is a shared feature extractor, and expert decoders yield
\[
O_i = D_i(h; \phi_i)\,.
\]
The final prediction is obtained by
\[
\hat{p}(x) = \sum_{i=1}^N O_i(x)
\]
where \(N\) is the number of expert decoders. The hybrid loss combines cross-entropy components weighted by expert agreement (Dice/IoU), e.g.
\[
L_{\mathrm{hybrid}}(\hat{p}, y_1, y_2) = w \, \mathrm{CE}(\hat{p}, y_1) + (1-w) \, \mathrm{CE}(\hat{p}, y_2)
\]
with adaptive weight \(w\).

For multimodal joint diffusion, let \(x = (x^{\mathrm{image}}, x^{\mathrm{text}})\) and the forward process be
\[
q(x_t \mid x_{t-1}) = \mathcal{N}\bigl(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I\bigr)
\]
with joint reverse model
\[
p_\theta(x_{t-1} \mid x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(t)).
\]
Training is performed using the standard L2 denoising objective:
\[
L(\theta) = \mathbb{E}_{t, x_0, \epsilon}\left\|\epsilon - \epsilon_\theta(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilon, t)\right\|^2.
\]

Classifier-free guidance for masked infilling uses two evaluations of the score network:
\[
\hat\epsilon_t = (1+w)\epsilon_\theta^{\mathrm{cond}} - w\epsilon_\theta^{\mathrm{uncond}}
\]
with \(w\) controlling the trade-off between conditional and unconditional signals.

## 3. Empirical Results and Performance Metrics

Biomedical segmentation performance was benchmarked for “U-Net-and-a-half” (dual-decoder PS-U-Net) against single-expert U-Nets on two datasets: pathologist-annotated glomerular segmentation (whole slide images, 10 subjects) and radiologist-annotated arteriovenous fistula cross-sections (IVUS, 10 subjects) [2108.04658].

Representative metrics (mean ± SD):

| Model              | Dice (1)        | Core-Dice (1)  | IoU (1)         | Dice (2)        | Core-Dice (2)     | IoU (2)         |
|--------------------|-----------------|---------------|-----------------|-----------------|-------------------|-----------------|
| UNet 1→1           | 0.9867±0.005    | 0.9410±0.02   | 0.7768±0.01     | –               | –                 | –               |
| UNet 2→2           | 0.9886±0.004    | 0.9435±0.01   | 0.7846±0.02     | –               | –                 | –               |
| UN aah→1           | 0.9874±0.005    | 0.9397±0.02   | 0.7573±0.02     | –               | –                 | –               |
| UN aah→2           | 0.9869±0.003    | 0.9316±0.01   | 0.7758±0.02     | –               | –                 | –               |

On cross-expert evaluation, UN aah (PS-U-Net) outperformed single-expert U-Nets when tested against the alternate annotation, indicating increased generalizability.

In multimodal joint diffusion, PS-U-Net attained the following Fréchet Inception Distance (FID) comparisons (MS-COCO, zero-shot) [2311.16488]:

| Model                  | Params   | FID   |
|------------------------|----------|-------|
| U-ViT-multi+Infilling  | 214M     | 13.90 |
| PS-U-Net+Infilling     | 230M     | 9.40  |
| Unidiffuser            | 1.1B     | 9.71  |

PS-U-Net reduced FID by approximately 4.5 points compared to the closest-size U-ViT baseline, converged three times faster, and demonstrated improved robustness at higher guidance scales.

## 4. Loss Functions and Training Protocols

For the segmentation scenario, the hybrid loss is central. When two expert annotations are available per image, per-pixel cross-entropy is computed separately and combined using a weight \(w\) that dynamically reflects expert agreement (as measured by Dice or IoU). For whole slide image segmentation, a focal variant is employed to better handle class imbalance. The loss generalizes to \(N\) decoders as:
\[
L_{\mathrm{total}} = \sum_{i=1}^N \alpha_i \mathrm{CE}(D_i(E(x;\theta); \phi_i), y_i), \quad \sum_i \alpha_i = 1.
\]

For multimodal diffusion, a joint denoising score-matching loss suffices, requiring no explicit multimodal regularizer. Data are embedded using Stable Diffusion autoencoder latents for images and Word2Vec for captions. Optimizer settings include AdamW with \(lr=2\times10^{-4}\), weight decay 0.03, and \(\beta=(0.9,0.9)\).

Inference employs masked classifier-free guidance with guidance scale \(w\) (optimal \(w\approx3\)), and arbitrary sub-blocks of either modality can be masked for infilling.

## 5. Extensions and Architectural Variants

In the segmentation context, potential PS-U-Net extensions include:
- Increasing the number of decoders for multi-expert learning beyond two annotators (general \(N\)-decoder form).
- Learning the fusion weights (\(w_i\)) through an auxiliary gating network.
- Sharing only a subset of decoder layers to model shared post-processing while retaining some expert specificity.
- Incorporating weakly supervised objectives to alleviate annotation burden.
- Extending convolutional kernels to 3D for volumetric data.

For multimodal diffusion, the architecture admits flexible fusion depth (parameterizing at which stage information is combined) and supports masking or activating individual towers for efficient conditional generation.

*This suggests an ongoing research interest in dynamic and learnable partial sharing mechanisms, as well as the integration of PS-U-Net paradigms into a broader range of generative and discriminative tasks.*

## 6. Significance, Challenges, and Future Directions

PS-U-Net architectures formalize the principle of sharing “what is universally informative” (through encoders or central fusion blocks) and specializing “what is expert- or modality-dependent” (through dedicated decoders or branches). This approach generalizes standard U-Nets, multitask segmentation models, and diffusion U-Nets by explicitly decoupling invariant and variant task structure [2108.04658, 2311.16488].

Key empirical results establish substantial gains in generalization, label fusion, multimodal sample quality (measured by FID), inference speed, and convergence rate compared to monolithic or fully shared approaches.

A plausible implication is that PS-U-Nets are especially pertinent when models must reconcile annotation ambiguity or interact with heterogeneous input spaces, whether arising from human expert disagreement or cross-modal data translation. Future research may explore scalable variants, learnable fusion strategies, and broader application domains including natural language understanding, audio-visual learning, and federated annotation settings.

Source: https://www.emergentmind.com/topics/partially-shared-u-net-ps-u-net