FANet: A Feedback Attention Network for Improved Biomedical Image Segmentation

In plain words

Segmentation means labelling every pixel of an image, for example marking which pixels belong to a polyp in a colonoscopy frame. FANet is a network for this in medical images. Its main trick is to reuse its own earlier answer. The mask it predicted in the previous training round is passed back in as a hint about where the target probably is, and the network uses that hint to decide which of its internal features to keep.

At test time the same loop runs repeatedly: the network's output becomes the hint for the next pass, so a rough first mask can be cleaned up over a few rounds. The authors report that this gives better overlap scores than several earlier medical segmentation networks on the datasets they tried. These are their own numbers, not independent results.

TL;DR

FANet is a U-Net-style encoder-decoder built from residual blocks with squeeze-and-excitation, plus a MixPool block in every encoder and decoder stage. MixPool uses the previous mask as hard attention on the feature maps. The mask is kept between training epochs and fed back at test time for up to 10 iterations. The authors report the best F1 on 4 of 7 datasets, but two of the paper's own text claims are contradicted by its tables (see Pitfalls).

Problem

Biomedical segmentation models are usually trained on small, varied datasets and make a single forward prediction. The paper argues that existing methods do not reuse information from earlier training epochs, and that a one-step prediction cannot correct its own mask §1. Its test-time refinement is intended to avoid heavy test-time augmentation or ensembles §1. The paper notes that the refinement adds inference time §5.

Method

FANet is a fully convolutional encoder-decoder with four encoder and four decoder blocks §3.3. Three building blocks matter:

• SE-residual block §3.1. Two 3×3 convolutions with batch norm, plus a shortcut. The paper describes the shortcut as an identity mapping. A squeeze-and-excitation layer re-weights the channels §3.1.
• MixPool block §3.2. Its input is the feature map F_l and the previous-epoch mask M_l. A sigmoid head on F_l gives a binary spatial attention map M_l' at threshold 0.5 eq. 1. The mask M_l is max-pooled to the feature size and combined with M_l' by union. That union multiplies F_l, and the result is concatenated with the plain features eq. 2.
• Feedback loop §3.3. Each training sample's predicted mask is run-length encoded and stored, then used as the input mask in the next epoch §3.3. At test time the same mask is fed back for up to 10 iterations §1, §5.

The encoder runs two SE-residual blocks, then MixPool, then pooling. The skip connection to the decoder is taken before MixPool §3.3. Each decoder starts with a 4×4 transposed convolution, concatenates the skip, applies two SE-residual blocks and MixPool §3.3. The last decoder output is concatenated with the previous mask, then a 1×1 convolution with sigmoid gives the output §3.3. The loss is binary cross-entropy plus Dice §3.3.

Initial mask: Otsu thresholding of each image gives the mask for the first epoch and for iteration 0 at test time §3.3, fig. 1.

The architecture is drawn in fig. 2; the initial mask and iteration effect are in fig. 1 and fig. 5.

# FANet forward pass  [§3.3] [fig. 2]  [code: model.py#L39-L71]
# Input:  image (B, 3, H, W); prev_mask (B, 1, H, W)
# Output: logits (B, 1, H, W); the sigmoid is applied in the loss

def fanet(image, prev_mask):
    # Encoder: SE-residual blocks, then MixPool, then 2x2 max-pool.
    # The skip is taken before MixPool  [§3.3] [code: model.py#L5-L20]
    skip_1, x = encoder(image, prev_mask, 32)
    skip_2, x = encoder(x, prev_mask, 64)
    skip_3, x = encoder(x, prev_mask, 128)
    skip_4, x = encoder(x, prev_mask, 256)
    # Decoder: 4x4 transposed conv (x2), concat skip, two SE-residual
    # blocks, then MixPool  [§3.3] [code: model.py#L22-L37]
    x = decoder(x, skip_4, prev_mask, 128)
    x = decoder(x, skip_3, prev_mask, 64)
    x = decoder(x, skip_2, prev_mask, 32)
    x = decoder(x, skip_1, prev_mask, 16)
    # Concatenate the previous mask, then a 1x1 conv  [§3.3]
    # [code: model.py#L68-L69]
    x = concat_channels(x, prev_mask)     # (B, 17, H, W)
    return conv1x1(x)                     # (B, 1, H, W)

# SE-residual block as implemented  [§3.1] [code: blocks.py#L23-L54]
# The paper says identity shortcut; the code uses a 1x1 projection.
# Input:  x (B, C_in, H, W); output channels C_out

def se_residual(x, out_channels):
    # Main branch: conv3x3-BN-ReLU, then conv3x3-BN  [§3.1]
    # [code: blocks.py#L40-L46]
    main = relu(batch_norm(conv3x3(x, out_channels)))
    main = batch_norm(conv3x3(main, out_channels))
    # Shortcut: 1x1 conv and BN, then squeeze-and-excitation
    # [code: blocks.py#L47-L49]
    shortcut = squeeze_excite(batch_norm(conv1x1(x, out_channels)))
    # Sum, then ReLU  [code: blocks.py#L51-L52]
    return relu(main + shortcut)

# MixPool block  [§3.2] [eq. 1] [eq. 2]  [code: blocks.py#L57-L88]
# Input:  features (B, C, h, w); prev_mask (B, 1, H, W)
# Output: (B, C, h, w), half the channels from each branch

def mix_pool(features, prev_mask):
    # Spatial attention: sigmoid head, thresholded  [eq. 1]
    # [code: blocks.py#L61-L67, L82]
    attention_map = sigmoid(attention_head(features))
    attention = attention_map > 0.5
    # Previous mask max-pooled to the feature size  [§3.2]
    # [code: blocks.py#L83]
    mask = max_pool(prev_mask, to_size=features.shape[-2:])
    # Union of the two binary maps gates the features  [§3.2] [eq. 2]
    # [code: blocks.py#L84]
    keep = attention or mask
    gated = features * keep
    # Two branches, each with half the channels, concatenated
    # [eq. 2] [code: blocks.py#L69-L87]
    gated = conv3x3_bn_relu(gated, C // 2)
    plain = conv3x3_bn_relu(features, C // 2)
    return concat_channels(gated, plain)

# Loss  [§3.3] [§4.1.2] [code: loss.py#L23-L41]
# Input:  logits (B, 1, H, W); targets (B, 1, H, W) in {0, 1}

def dice_bce_loss(logits, targets):
    probs = sigmoid(logits)
    # Soft Dice with smoothing term 1  [code: loss.py#L37]
    dice = (2 * sum(probs * targets) + 1)
    dice = dice / (sum(probs) + sum(targets) + 1)
    bce = mean(binary_cross_entropy(probs, targets))
    # Equal weights 0.5 and 0.5; the paper does not state weights
    # [code: loss.py#L39]
    return 0.5 * bce + 0.5 * (1 - dice)

# Training with feedback masks  [§3.3] [code: train.py#L226-L243]
# Input:  train and valid sets; the number of epochs
# The paper updates every stored mask after each epoch  [§3.3].
# The code updates them only when validation loss improves.

def train_with_feedback(model, train_set, valid_set, epochs):
    # Initial masks from Otsu thresholding  [§3.3] [code: utils.py#L90-L110]
    train_masks = otsu_masks(train_set.images)
    valid_masks = otsu_masks(valid_set.images)
    best_valid_loss = infinity
    for epoch in range(epochs):
        # Forward pass with the stored masks; new masks are taken from
        # the output before this batch's optimizer step
        # [code: train.py#L82-L96]
        new_train_masks, train_loss = train_epoch(model, train_set,
                                                  train_masks)
        new_valid_masks, valid_loss = evaluate(model, valid_set,
                                               valid_masks)
        # Stored masks change only when this epoch beats the best
        # validation loss so far  [code: train.py#L236-L243]
        improved = valid_loss < best_valid_loss
        if improved:
            best_valid_loss = min(best_valid_loss, valid_loss)
            save_checkpoint(model)
            train_masks = new_train_masks
            valid_masks = new_valid_masks

# Test-time refinement  [§1] [§5] [code: test.py#L101-L161]
# Input:  trained model; test images; number of iterations (10 in paper)
# Output: one prediction per iteration, per image

def refine(model, images, num_iter):
    # Iteration 0 starts from the Otsu mask  [§3.3] [fig. 1]
    masks = otsu_masks(images)
    predictions_per_iter = []
    for step in range(num_iter):
        # Each prediction becomes the input mask of the next pass
        probs = sigmoid(model(images, masks))
        masks = probs > 0.5
        predictions_per_iter.append(masks)
    return predictions_per_iter

Hyperparameters

Hardware: training on a Volta 100 and an NVIDIA DGX-2 with PyTorch 1.6; inference on a GTX 1050 Ti §4.1.2.

Results

Author-reported. Setup: seven public datasets, with the train and test splits of earlier work §4.1.1, tab. 1. Metrics: F1 (Dice), mIoU, recall and precision. Strongest other method is the best F1 in the same table of the paper.

Ablation on Kvasir-SEG, with the same F1 and mIoU at 512×512 tab. 9:

Efficiency tab. 10: FANet has 7.72 M parameters, 94.75 GMac, and 8.25 ms inference. U-Net has 31.04 M parameters and 3.14 ms. Attention U-Net has 34.88 M parameters and 4.47 ms. The paper says the inference time is higher than the baselines §4.5.

Extended ablation on Kvasir-SEG tab. 11: F1 0.8803 with the design used, against 0.8682 for four FANets run in series and 0.8570 for three SE blocks.

Comparators in the paper: U-Net, ResUNet, ResUNet++, FCN8, HRNet, DoubleU-Net, PSPNet, DeepLabv3+, U-Net++, Attention U-Net, R2U-Net, BCDU-Net, IterNet, DUNet and DenseBlock-UNet tab. 2–8.

Paper vs code

The official repository is linked from the abstract. Its owner is the first author. code_commit is the HEAD of the shallow clone. It contains only Kvasir-SEG training and testing, plus a README with the paper's abstract and citation README.md. The repository has no LICENSE file, so the code is cite-only: no verbatim snippets, and the pseudocode above is our own.

• Feedback update. The paper says the stored mask is updated after each epoch §3.3. The code stores new masks only when validation loss improves train.py.
• Epochs. The paper uses 100 epochs §4.1.2. The code runs 500 train.py. It keeps the checkpoint with the lowest validation loss train.py.
• Input size. The paper uses 512×512 for Kvasir-SEG §4.1.2. The code uses 256×256 train.py and also at test time test.py.
• Learning rate. The paper uses 1e-3 for DRIVE and CHASE-DB1 §4.1.2. The code has one rate, 1e-4 train.py.
• Augmentation. The paper lists random crop, elastic and grid and optical distortion, grayscale, brightness, contrast and channel changes §4.1.2. The code applies rotation, horizontal and vertical flips, and coarse dropout, each at probability 0.3 train.py.
• Residual shortcut. The paper says identity mapping §3.1. The code uses a 1×1 convolution with batch norm, and applies the squeeze-and-excite layer on that shortcut branch blocks.py.
• Second conv activation. The paper says each convolution is followed by BN and ReLU §3.1. In the code, the second conv has BN but no ReLU before the sum blocks.py.
• SE width. The paper gives no reduction ratio §3.1. The code passes the channel count as the reduction, so the hidden layer has one unit blocks.py.
• Threshold. The paper uses σ(·) ≥ 0.5 eq. 1. The code uses > 0.5 blocks.py.
• Loss weights. The paper says BCE plus Dice with no weights §4.1.2. The code uses 0.5 and 0.5 loss.py.
• Test split. The paper calls its evaluation set the test set tab. 1. test.py reads val.txt as the test set, because load_data returns the validation split second test.py. Split files are not in the repository.
• Iterations. The paper selects 10 as optimal §5. test.py runs 10 and reports every iteration; no rule selects one test.py.
• Models in the ablation. The repository has only the full model (B4). The B1, B2 and B3 configurations are not in the code.
• Metrics. The paper does not say how scores are averaged. The code averages per-image scores over the test set test.py.

No discrepancies found in the encoder and decoder layout, the channel counts used in the decoder, or the MixPool union and concatenation model.py.

Pitfalls

• Gradient does not reach the attention head. The attention map is thresholded and cast to float, so no gradient flows through it blocks.py. From the code alone, the attention head's weights do not learn, while its batch-norm statistics still update. I did not run this.
• CUDA is hard-coded. MixPool casts to torch.cuda.FloatTensor blocks.py. test.py falls back to CPU test.py, but MixPool would still need CUDA. I did not run this.
• Mask orientation. The RLE encoder flattens the transpose utils.py, and train and test each transpose on decode train.py. Keep both transposes if you port this.
• Otsu and the first iteration. The initial mask is an Otsu threshold after a 5×5 Gaussian blur utils.py. The paper does not state the blur. The code uses a binary threshold, so pixels above the Otsu value become 1 utils.py.
• Hidden cost of iterations. Each iteration is one full forward pass per image, so inference time scales with the iteration count §5.
• Paper text vs tables. Three claims disagree with the paper's own tables:
  • CVC-ClinicDB: the text says FANet "outperforms all the SOTA methods … by a large margin" §4.2.2, but U-Net++ has higher F1 (0.9377 vs 0.9355) tab. 3.
  • 2018 DSB: the text says FANet has "the highest F1 score (0.9176)" §4.2.3, but Attention U-Net has 0.9179 tab. 4.
  • 2018 DSB: the text gives a 2.02% F1 gain over U-Net++ §4.2.3. The table values (0.9176 vs 0.9117) give about 0.65% relative, or 0.59 points.
• Table values. In the DSB table, the DoubleU-Net row lists F1 0.7683 below its mIoU 0.8407 tab. 4. The paper gives no explanation.
• Timing. The efficiency table gives 8.25 ms per image tab. 10. The ablation table reports 66 to 71 FPS at 512 for the same model tab. 9. The two figures are not reconciled.
• "Far fewer epochs." The conclusion claims near-SOTA results with far fewer training epochs §6. The paper reports no baseline epoch counts.
• Hard-coded paths. train.py and test.py point at local dataset directories train.py.
• Global variables. size is set inside __main__ but used in train, evaluate and the dataset class train.py. Import these functions from another script and they fail.

Lineage

• Builds on: squeeze-and-excitation residual networks (Hu et al., cited as §3.1); the mask-guided attention of Song et al., cited as §2; U-Net as the encoder-decoder base §2.
• Compared against: the methods listed under Results tab. 2–8.
• Superseded by: not stated.