module

Attention U-Net

3 members

lucid.models.vision.attention_unet

Attention U-Net (Oktay et al., MIDL 2018).

Paper: "Attention U-Net: Learning Where to Look for the Pancreas"

Semantic Segmentation

Paper

Oktay, Ozan, et al. "Attention U-Net: Learning Where to Look for the Pancreas." Medical Imaging with Deep Learning, 2018.

Overview

Attention U-Net augments the standard U-Net skip pathways with soft spatial attention gates. In a plain U-Net every encoder feature is concatenated unconditionally into the decoder, so irrelevant background activations still propagate to the output. For small structures like the pancreas — where the target occupies a tiny fraction of the field of view — this leaks gradient signal into the wrong regions and hurts both convergence and final Dice.

Each Attention Gate (AG) at decoder level $\ell$ takes two inputs: the skip feature $x^{\ell} \in \mathbb{R}^{C \times H \times W}$ from the encoder and a gating signal $g^{\ell+1} \in \mathbb{R}^{C_g \times H/2 \times W/2}$ from the deeper decoder stage. Both are projected to a common intermediate space with $1 \times 1$ convolutions and additively combined:

q^{\ell} = \psi^{\top}\!\bigl( \sigma_1\!\bigl( W_x^{\top} x^{\ell} + W_g^{\top} g^{\ell+1} + b_g \bigr) \bigr) + b_\psi, \qquad \alpha^{\ell} = \sigma_2(q^{\ell}),

where $\sigma_1$ is ReLU and $\sigma_2$ is sigmoid, so $\alpha^{\ell}(h, w) \in [0, 1]$ is a per-pixel attention coefficient. The gated skip $\tilde{x}^{\ell} = \alpha^{\ell} \odot x^{\ell}$ then replaces the raw skip in the U-Net concatenation.

Because the gate is differentiable and trained jointly with the rest of the network, no extra supervision (bounding boxes, landmarks, …) is needed: the model learns where to look purely from the segmentation loss. The added cost is small (two $1 \times 1$ convs per gate) but consistently improves Dice on small-organ targets and tightens the response near tissue boundaries — a now-standard add-on in medical-imaging U-Net variants.