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

In plain words

Segmenting a small organ in a scan, such as the pancreas on a CT image, is hard because the organ is small, its shape varies from patient to patient, and it looks much like the tissue around it. A common workaround is to chain two networks: the first finds a rough box around the organ, and the second segments inside that box. This works, but it costs extra models and extra compute.

The paper's idea is to let one network learn where to look. A small "attention gate" sits on each of the links that pass detail from the encoder to the decoder. It uses a coarser, higher-level view of the image to decide which parts of the detail are relevant, and it turns down the rest. The gates are trained together with the rest of the network, so no separate localisation step is needed.

The gate is added to a standard U-Net and adds only a small amount of compute. The authors test it on two abdominal CT data sets and report that it helps across training-set sizes and data sets.

TL;DR

• An additive attention gate weights each skip connection with sigmoid coefficients, computed from the skip features and a coarser gating signal §2 eq. 1 eq. 2.
• Gates are inserted into a 3D U-Net with deep supervision and need no separate localisation model §2 §3 fig. 1.
• Author-reported: pancreas Dice 0.840 ± 0.087 for Attention U-Net vs 0.814 ± 0.116 for U-Net, trained on 120 and tested on 30 CT-150 scans tab. 1.
• Author-reported: on the public TCIA Pancreas-CT set the gains are smaller (0.712 vs 0.690 before fine-tuning) tab. 3.
• The cost is modest but not zero: 6.40 M vs 5.88 M parameters and 0.179 s vs 0.167 s per forward pass in the authors' setup tab. 1.

Problem

Standard fully convolutional networks such as U-Net handle large organs well. For small organs with large shape variation, such as the pancreas, the state of the art at the time used cascades: a first model finds a region of interest (ROI), and a second model segments inside it §1. The authors argue that cascades repeat similar low-level feature extraction across models and need more parameters and compute §1. Earlier attention work for image classification pools features over regions, which the authors note does not suit dense prediction §1.

The paper asks whether attention can replace the localisation model inside one network, while keeping dense predictions and standard back-propagation training §1 §2.

Method

Base network. The gate is built on a standard U-Net: convolutional blocks at several scales, a downsampling path, an upsampling path, and skip connections that merge coarse and fine features §2 fig. 1. Coarse features carry the category and location of foreground objects, and fine features carry detail §2.

Attention gate. For each skip connection, the gate takes the skip features x^l and a gating signal g from a coarser scale §2 fig. 2. It computes one attention coefficient per voxel, or a vector of coefficients when there are several classes §2. The output is the skip features multiplied element-wise by the coefficients §2. Gating signals come from the coarser scale and aggregate information from several scales §2 fig. 1. Low-level skip connections, the first ones, are not gated §2.

Additive attention, formulas. The authors use additive attention §2. The attention logit and coefficient are given by eq. 1 and eq. 2:

q^l_{att} = \psi^T \left(\, \sigma_1\,(\,W_x^T x_i^l + W_g^T g_i + b_g\,)\,\right) + b_{\psi} eq. 1

\alpha_i^l = \sigma_2 (\, q^l_{att}(x_{i}^l\,,\, g_i \,;\, \Theta_{att}) \,) eq. 2

Here σ1 is a ReLU and σ2 is the sigmoid §2. Parameters Θ_att are W_x (F_l × F_int), W_g (F_g × F_int), ψ (F_int × 1), and the biases b_g and b_ψ. Linear maps are 1×1×1 convolutions §2. The authors choose the sigmoid over softmax, because sequential softmax gives sparser outputs, and they report better convergence with sigmoid §2.

Grid-based gating. Unlike the global-vector gating in §1 (Jetley et al.), the gating signal here is a grid conditioned on image position, so each voxel gets its own gating value §1 §2. The skip features are downsampled to the gating resolution before the linear map, similar to non-local blocks §2.

Gradients. The gate is trained by ordinary back-propagation. The gradient reaching the layer below has two terms: the usual term scaled by α, and a term through the coefficient itself §2 eq. 3:

\frac{\partial (\hat{x}_i^l)}{\partial \, (\Phi^{l-1})}= \alpha_i^l \, \frac{\partial(f(x_i^{l-1}; \Phi^{l-1}))}{\partial \, (\Phi^{l-1})} + \frac{\partial (\alpha_i^l)}{\partial \, (\Phi^{l-1})} \, x_i^l eq. 3

Deep supervision. Each decoder scale gets an auxiliary output. The authors use this to make the intermediate features discriminative at every scale, so that attention units at all scales can influence the output §2.

Multiple gates per scale. Several gates are trained per scale, and each one learns to focus on a different subset of organs §2 §3. The paper does not state how many gates are used.

Training loop, sketch. The pseudocode below follows the text. The released code differs in several places; see "Paper vs code".

# Additive attention gate on one skip connection  [§2] [eq. 1] [eq. 2]
# Input:  skip features x (B, F_l, H, W, D); gating signal g (B, F_g, h, w, d)
#         where (h, w, d) is the coarser grid of g
# Output: gated features (B, F_l, H, W, D); coefficients (B, 1, H, W, D)

def attention_gate(x, g):
    # Linear map of the skip features, down to the gating grid  [§2] [eq. 1]
    x_proj = conv1x1(downsample(x, to_grid_of=g))   # (B, F_int, h, w, d)
    # Linear map of the gating signal, bias b_g  [§2] [eq. 1]
    g_proj = conv1x1(g)                             # (B, F_int, h, w, d)
    # Additive combination with ReLU  [eq. 1]
    hidden = relu(x_proj + g_proj)
    # Project to one score per voxel, bias b_psi  [eq. 1]
    logits = conv1x1(hidden)                        # (B, 1, h, w, d)
    # Sigmoid gives coefficients in [0, 1]  [eq. 2]
    alpha_coarse = sigmoid(logits)
    # Resample the coefficients to the skip grid (trilinear)  [fig. 2] [§2]
    alpha = upsample_trilinear(alpha_coarse, to_grid_of=x)
    # Element-wise scaling, so background activations are suppressed  [§2]
    gated = alpha * x
    return gated, alpha

# Soft Dice loss over all classes  [§3] [code: models/layers/loss.py#L29-L48]
# Input:  logits (B, C, N) per voxel; one_hot (B, C, N) targets
# Output: scalar loss

def soft_dice_loss(logits, one_hot):
    # Softmax over the class axis  [code: models/layers/loss.py#L39]
    probs = softmax(logits, dim=1)
    # Per-sample, per-class overlap and sizes, with smoothing 0.01
    # [code: models/layers/loss.py#L36-L43]
    overlap = sum(probs * one_hot, over N) + 0.01
    total = sum(probs, over N) + sum(one_hot, over N) + 0.01
    # Mean Dice over samples and classes; loss is one minus Dice
    # [code: models/layers/loss.py#L45-L46]
    dice = mean(2 * overlap / total)
    return 1 - dice

Hyperparameters

Values come from the paper unless the "Where" column names code. "not stated" means the paper and this card's cited code do not give the value.

Results

Author-reported. The paper uses 3D models, 2 mm isotropic resolution and the datasets named below. "±" is shown as printed; the paper does not say what it measures §3. Statistical tests are named only as p-values, without the test §3.

Datasets: CT-150 is 150 abdominal CT scans of patients with gastric cancer, with pancreas, liver and spleen labels §3. CT-82 is the public NIH TCIA Pancreas-CT set, with 82 contrast-enhanced scans and slice-by-slice pancreas annotation §3 tab. 3.

For context, the state-of-the-art table lists single- and multi-model CNN methods on CT-82. The best-listed numbers are 82.3 ± 6.7 (2D FCN with recurrent network) and 82.2 ± 5.7 (multi-model 2D FCN), using 62/20 splits and 4-fold CV tab. 4. The paper describes its 3D model as performing "similar to the state-of-the-art" §3. No Attention U-Net row appears in tab. 4.

Reported significance: attention gates increase recall on CT-150 at p = .005 (120/30) and improve DSC at p = .01 for 30 training images §3. Capacity-matched U-Net comparison gives p = .007 §3. On CT-82 from scratch, S2S distance has p = .03 and recall p = .09 §3.

Paper vs code

Official code: ozan-oktay/Attention-Gated-Networks (linked in the paper's footnote; the owner is the first author) at commit eee4881fdc31920efd873773e0b744df8dacbfb6, MIT license. Checked: gate layer, the multi-gate 3D U-Net, the CT configs, the optimiser and loss code, and the CT augmentation pipeline.

• Optimiser. Paper: Adam §3. Code: get_optimizer builds SGD with momentum 0.9 and Nesterov unless the config sets optim to adam utils.py. The released CT configs do not set optim config_unet_ct_multi_att_dsv.json, so they train with SGD.
• Gate initialisation. Paper: gates initialised to pass all features at every location §3. Code: convolution weights get Kaiming init; the attention biases are left at the framework default grid_attention_layer.py networks_other.py. A pass-through bias of 3.0 appears only in a separate gate variant, not in the class the CT networks use grid_attention_layer.py.
• Output projection. The paper gives the output as the gated features §2. Code adds a 1×1×1 convolution and batch normalisation after the scaling grid_attention_layer.py grid_attention_layer.py.
• Gate fusion. The paper says gate outputs are "fused" §2. Code concatenates the two gates and applies 1×1×1 convolution, batch normalisation and ReLU unet_CT_multi_att_dsv_3D.py unet_CT_multi_att_dsv_3D.py.
• Flips. Paper: "axial flips" §3. Code: the CT training transform flips along both axes, with probability 0.5 transforms.py config_unet_ct_multi_att_dsv.json.
• Pancreas-only fine-tuning. Paper: spleen and kidney output nodes are removed from the softmax and only background and pancreas get gradients during fine-tuning §3. Code: a CustomSoftDiceLoss with class indices [0, 2] exists loss.py and is selected by dice_loss_pancreas_only utils.py. The released configs use dice_loss config_unet_ct_multi_att_dsv.json. The mapping from indices to organs is not stated in the code.
• Dice smoothing. Paper: no smoothing constant §3. Code: 0.01 in numerator and denominator, with the background channel included loss.py.
• Learning rate and schedule. Paper: not stated §3. Code: 1e-4, step decay ×0.5 every 250 epochs, 1000 epochs config_unet_ct_multi_att_dsv.json config_unet_ct_multi_att_dsv.json networks_other.py.
• Resampling and intensity. Paper: 2 mm isotropic resampling and N(0,1) intensities §3. Code: intensity normalisation is called transforms.py but the function is not defined in this repository, and no resampling code was found. Neither claim can be checked against the code.
• Model variant. The paper does not name its class. The text (multiple gates per scale, §3) matches the multi-gate class; the single-gate class also exists unet_CT_single_att_dsv_3D.py.

Pitfalls

• Missing config keys. The shipped multi-gate config does not set nonlocal_mode or attention_dsample config_unet_ct_multi_att_dsv.json, but the model builder reads opts.nonlocal_mode and opts.attention_dsample feedforward_seg_model.py. Set them explicitly, with the values the paper used, or the run may fail. This was read, not run.
• Table 2 labels. Both rows of tab. 2 are labelled "U-Net (120/30)". The only way to tell them apart is the parameter count (6.44 M and 10.40 M).
• Table 4 has no attention row. The attention model's CT-82 cross-validation number (81.48 ± 6.23) appears only in the text §3. Table 4 compares methods trained on different splits and fold counts, so it is not a like-for-like ranking.
• Undefined spread. The paper does not say what "±" is computed over.
• Statistical tests. The tests behind the p-values are not named §3.
• Dataset access. CT-82 is public. The paper does not say whether CT-150 is available §3.
• Hidden cost. The paper says the gates add "minimal computational overhead" [abstract]. Table 1 shows 6.40 M vs 5.88 M parameters, and 0.179 s vs 0.167 s per forward pass, which is about 7% more time tab. 1. The overhead is therefore small but measurable.
• Sigmoid, not softmax. The paper reports better convergence with sigmoid §2 and has no table for the softmax comparison.
• Gate scale. The gating signal comes from the coarser scale §2 fig. 1. The paper does not say how the gate's grid should be set if the encoder's downsampling changes.
• Memory. The model is 3D on 160×160×96 patches tab. 1. The paper says gradient averaging over several passes is used for larger networks §3 but does not give the number of passes.
• Sub-sampling. The skip features are downsampled to the gating grid before the linear map §2. The released code uses a strided convolution with kernel = stride = (2, 2, 2) grid_attention_layer.py.
• Scale of the attention map. The attention coefficients are in [0, 1] eq. 2. The code's output projection (see "Paper vs code") can change the output scale, so compare against the released code, not the equations alone.

Lineage

• Builds on: the U-Net §1 (cited, no card here; not used beyond naming). Additive attention, as in Bahdanau et al. 2014 §2. Attention for classification and captioning, as in Jetley et al. and Anderson et al. §1.
• Compared against: the standard 3D U-Net, with and without extra capacity tab. 1 tab. 2; cascaded and multi-model CT pancreas methods on CT-82 tab. 4.