Shallow Attention Network for Polyp Segmentation

In plain words

Polyps are growths in the colon that have to be found and outlined in camera images. A neural network can mark where each polyp is in a frame. The trouble is that the same kind of polyp looks different when it is filmed in different clinics, under different light, and polyps are often tiny. Deep networks shrink the picture several times on the way through, and small polyps get lost in that process.

This paper, SANet, makes three changes. First, during training it swaps the colour statistics of one image with those of another image, so the network cannot rely on colour and has to learn the shape of the growth. Second, it lets the fine detail from early layers be selected by the more abstract information from later layers, so the background noise is suppressed and the small polyps keep their edges. Third, at test time it adjusts the network's raw scores so that the rare polyp pixels are not swamped by the many background pixels. The authors report that this beats earlier polyp models on five public datasets.

TL;DR

SANet is a polyp segmentation network built on a Res2Net encoder with a light decoder. It adds colour exchange during training, a shallow attention module that gates early features with deep ones, and a logit correction at inference. The paper reports the best mean Dice and IoU on five public polyp benchmarks against six baselines, and about 72 FPS on one GPU §4.2. Those numbers are author-reported. The released code omits the shallow attention module, so do not treat the code as a faithful implementation of the paper's architecture.

Problem

• Colour distributions differ across collection conditions, and the network can overfit to colour §1 fig. 1.
• Polyps are small. The text reports that most polyp areas are below 0.1 of the image area §1 fig. 1. Repeated downsampling loses them §1.
• Background pixels dominate, which biases the model towards low confidence on polyp pixels §1 §3.3.

The paper assumes a labelled dataset of polyp images with pixel masks, and it attributes the overfitting to the limited data size §1 §5.

Method

Overview. The encoder is Res2Net §3. The paper uses the features of the last three of the five blocks, written f_3, f_4 and f_5 §3. fig. 2 draws these as res-3 (44×44×512), res-4 (22×22×1024) and res-5 (11×11×2048) for a 352×352 input. The decoder applies a shallow attention module to res-3 and res-4, guided by the next deeper level. The results are concatenated, then a sigmoid and the probability correction produce the mask fig. 2. The text does not spell out the concatenation or the final prediction head §3. The description in this section follows fig. 2 and the equations.

Colour exchange (training only) §3.1 alg. 1

For each training image, another image is drawn from the dataset. The colour statistics of the two are exchanged in the LAB colour space, so the training input keeps the content of the first image but takes the colour of the second §3.1. The statistics are the per-channel mean and standard deviation alg. 1. Colour exchange is applied only during training. Inference uses the original images §3.1.

# Colour exchange for one training image  [alg. 1] [§3.1]
# Input:  img1, img2 as RGB arrays (H, W, 3)
# Output: out1, the new training input (H, W, 3)

def color_exchange(img1, img2):
    # Convert both images to the LAB colour space  [alg. 1]
    lab1 = to_lab(img1)                          # (H, W, 3)
    lab2 = to_lab(img2)                          # (H', W', 3)

    # Per-channel mean and standard deviation of each image  [§3.1] [alg. 1]
    mean1, std1 = channel_stats(lab1)            # (3,)
    mean2, std2 = channel_stats(lab2)            # (3,)

    # Give the content of img1 the colour statistics of img2  [§3.1] [alg. 1]
    new_lab1 = (lab1 - mean1) / std1 * std2 + mean2

    # Convert back to RGB; this is the new input image  [alg. 1] [§3.1]
    out1 = to_rgb(new_lab1)
    return out1

The algorithm also produces a second image by swapping the statistics the other way alg. 1. The method text only uses the first output as the new input §3.1, so the second is not needed for training.

Shallow attention module §3.2 eq. 1 eq. 2 fig. 2

The deep features f_d are upsampled to the size of the shallow features f_s. They are passed through ReLU to form an attention map, which is multiplied element-wise into f_s. The paper states this as Att = σ(Up(f_d)) and f_s = Att ⊗ f_s, with σ the ReLU function eq. 1 eq. 2 §3.2. The module is described as a gate that filters background noise from shallow features, not as self-attention §3.2. The paper also says it weights levels by their contribution, instead of a fixed weight for each level §3.2. The text does not say how the attention map's channel count is matched to f_s, which is required for the element-wise product.

# Shallow attention module  [§3.2] [eq. 1] [eq. 2]
# shallow: f_s (B, C, H, W); deep: f_d (B, C2, h, w), with h <= H

def shallow_attention(shallow, deep):
    # Upsample the deep features to the shallow resolution  [§3.2] [eq. 1]
    up = upsample(deep, size=(H, W))             # (B, C2, H, W)
    # ReLU turns them into a non-negative attention map  [§3.2] [eq. 1]
    att = relu(up)                               # (B, C2, H, W)
    # Gate the shallow features with the attention  [§3.2] [eq. 2]
    return att * shallow                         # (B, C, H, W)

Probability correction at inference §3.3 eq. 3 eq. 4

Before the sigmoid, the positive logits of an image are divided by the fraction of positive logits in that image, and the negative logits by the fraction of negative logits eq. 3 eq. 4 §3.3. The positive fraction is written rate^p and the negative fraction rate^n. The paper says this is applied only at inference, adds little compute, and does not need retraining §3.3. fig. 3 shows predictions with and without the correction.

# Probability correction at inference  [§3.3] [eq. 3] [eq. 4]
# logits: (H, W) for one image, before the sigmoid

def correct_logits(logits):
    # Fraction of positive and negative logits in this image  [§3.3]
    rate_pos = mean(logits > 0)                  # scalar
    rate_neg = mean(logits < 0)                  # scalar
    # Divide positive logits by the positive fraction  [§3.3] [eq. 3]
    corrected = logits.copy()
    corrected[logits > 0] = logits[logits > 0] / rate_pos
    # Divide negative logits by the negative fraction  [§3.3] [eq. 4]
    corrected[logits < 0] = logits[logits < 0] / rate_neg
    return sigmoid(corrected)                    # (H, W)

Logits that are exactly zero are not changed by the equations eq. 3 eq. 4. The paper does not say what happens when a rate is zero.

Loss §3.4 eq. 5

Training uses binary cross-entropy plus Dice loss, with weights λ1 and λ2 both set to 1 eq. 5 §3.4. The paper does not define the Dice formula §3.4.

# Training loss  [§3.4] [eq. 5]
# logits, mask: (B, H, W); mask holds 0 or 1

def training_loss(logits, mask):
    # Binary cross-entropy on the raw logits  [§3.4] [eq. 5]
    bce = binary_cross_entropy_with_logits(logits, mask)
    # Dice on the sigmoid probabilities; the formula is not stated here  [§3.4]
    probs = sigmoid(logits)                      # (B, H, W)
    inter = sum(probs * mask, over H and W)      # (B,)
    total = sum(probs, over H and W) + sum(mask, over H and W)
    dice = 1 - mean(2 * inter / (total + 1))     # scalar
    # Equal weights, lambda_1 = lambda_2 = 1  [§3.4] [eq. 5]
    return bce + dice

The "+1" smoothing term and the per-image averaging are from the code, not the paper train.py.

Backbone §3

Res2Net is the encoder §3. The paper uses the features of the last three of five blocks §3. The paper does not state which Res2Net variant or pretraining it uses.

Hyperparameters

Values not in the paper are marked and taken from the code, which is a separate source.

Results

Author-reported. Setup: five public polyp datasets, six baselines, a Res2Net-based model. Image resolution is 352×352 §4.1. Table 1 is reproduced in full tab. 1. A dash means the paper did not report that value; the text does not explain the dashes. The paper colours second-best entries blue; the table below does not reproduce colours.

Source: tab. 1. The paper states that SANet is best on both mDice and mIoU across the five datasets §4.2.

Ablation on ColonDB and Kvasir tab. 2:

Source: tab. 2. The paper's text says each component is needed §4.4. Each row adds one component to the one above it, but the paper does not describe how the backbone-only row is built.

Speed: about 72 FPS on an RTX 2080 Ti, against 64 FPS for PraNet, as the paper reports §4.2. The repository has no timing code, so this figure is not reproducible from the release.

Visual results are in fig. 5 and the Dice-vs-threshold curves on three datasets are in fig. 4. The text says SANet's curve is above the others across thresholds §4.2.

Paper vs code

Code: weijun-arc/SANet at commit 68da18e, MIT license, copyright line "Jun Wei" LICENSE. The paper does not link the code. The repository's README links the paper, and its clone URL names weijun88/SANet, which GitHub resolves to the same repository. The owner matches the first author, and the repository was created before the arXiv posting.

• Shallow attention module is absent. The paper defines the ReLU-gated multiply of §3.2 eq. 1 eq. 2, and fig. 2 draws it. The code has no such module. The decoder reduces the features of the last three stages to 64 channels with 1×1 convolution, batch norm and ReLU, then combines them as cat([out5, out4*out5, out3*out4*out5]) and predicts with a 1×1 convolution model.py model.py. Those products are a multiplicative cascade between levels, not the gate of eq. 1 eq. 2.
• Colour exchange partner. The paper says a partner is drawn at random from the dataset §3.1. The code draws it from a subset of files whose names are all digits, with probability 0.7, and from the other subset otherwise train.py train.py. This depends on file names, and the rule for the other subset is implicit.
• Colour exchange is always on. The paper describes CE as a training step §3.1. The code applies it to every training sample with no probability gate train.py.
• Colour exchange cast. The result is cast to uint8 without clipping train.py. Values outside 0 to 255 wrap around rather than saturate. The paper does not say how the result is converted.
• Learning rate. The paper gives 0.04 §4.1. The code uses 0.04 for the backbone group (0.1 × 0.4) and 0.4 for the rest, so the 0.04 figure applies only to the backbone train.py train.py.
• Learning-rate decay is in the code at epochs 64 and 96 train.py. The paper does not mention it.
• Augmentation. The paper says "random rotation" §4.1. The code rotates only by multiples of 90° train.py.
• Multi-scale is per batch in the code, with sizes 256 to 352 train.py.
• Pretraining. The code loads a pretrained Res2Net checkpoint from res/ and ignores missing keys (strict=False) res2net.py. The paper does not say the backbone was pretrained. A mismatched checkpoint would load without an error.
• Loss uses raw logits for BCE and sigmoid probabilities for Dice, with +1 smoothing in the denominator train.py. The paper defines neither the Dice formula nor the smoothing.
• Ablation switches do not exist. The model has no flag for SAM, colour exchange or PCS model.py train.py test.py. The rows of tab. 2 cannot be regenerated from the release without editing code.
• PCS at inference matches eq. 3 eq. 4 in form: positive logits are divided by the positive fraction and negatives by the negative fraction test.py. The code applies it after the logits are upsampled to the original image size test.py. The paper does not say at what resolution PCS is applied.
• Test set name. The paper lists EndoScene §4.1. The code's test folders are CVC-300, CVC-ClinicDB, CVC-ColonDB, ETIS-LaribPolypDB and Kvasir test.py. Neither source says CVC-300 and EndoScene are the same images.
• Metric. The paper reports mDice and mIoU tab. 1 without a threshold. The code averages over 256 thresholds main.m main.m. Compare numbers only after checking this.
• Speed. The repository has no timing code. The 72 FPS figure is not reproducible from it §4.2.

Pitfalls

• The main pitfall is that the released code is not the paper's architecture. Anyone implementing from the code will miss the shallow attention module and get a different network. Implement from eq. 1 and eq. 2, and treat the code as the source for hyperparameters only where the paper says nothing.
• The learning-rate split (0.04 backbone, 0.4 head in code) changes the training dynamics by a factor of ten on the head. Confirm the intended value before running.
• Channel matching in the attention is not stated: the upsampled deep features have a different channel count from f_s, and the paper does not say how they are matched. The element-wise product eq. 2 needs equal shapes.
• Shallow attention ReLU sets every negative attention value to zero, so those positions of the shallow features are removed entirely eq. 1 eq. 2. The paper does not say whether any normalisation follows.
• PCS is sensitive to rate. Dividing by a small positive fraction scales the logits by a large factor before the sigmoid, which can saturate it. The paper has no clamp. Zero rates need a guard, which the paper does not state eq. 3 eq. 4.
• PCS needs the logits, not probabilities. Apply it before the sigmoid §3.3 fig. 2.
• Colour exchange depends on the file-name convention in the code, and the cast to uint8 without clipping can wrap colours train.py.
• Test-time resolution. Inputs are 352 × 352, and logits are upsampled to the original size before PCS in the code test.py test.py. The paper does not say where PCS sits relative to upsampling.
• Table 1 dashes are not defined. Do not read them as zero.
• Cross-reference slip. The sentence in §2 that says "as shown in the Fig. 1" for the architecture points to the distribution figure §2. The architecture is fig. 2.
• Cost. The paper says colour exchange adds no test-time cost §3.1 and PCS adds very little computation §3.3. It does not give a cost for the shallow attention module. The 72 FPS figure has no timing code behind it.

Lineage

• Builds on: Res2Net as the encoder §3; PraNet's dataset split, which the paper follows §4.1. Section 2 describes PraNet's reverse attention as a boundary-refinement approach §2, which this paper does not use.
• Compared against: U-Net, U-Net++, ResUNet, ResUNet++, SFA and PraNet §4.1 tab. 1.
• Superseded by: not stated in the wiki.