TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

In plain words

Medical image segmentation means outlining organs, or the chambers of the heart, pixel by pixel in a scan. The standard tool for this is U-Net, a convolutional network that is good at fine local detail but sees the image through small windows, so it struggles to relate parts that are far apart. Transformers relate every part of the input to every other part, but on their own they work on coarse chunks of the image and lose the sharp edges that an exact outline needs.

TransUNet combines the two. A convolutional stem first turns the scan into a feature map. A Transformer then mixes information across the whole map. A U-shaped decoder upsamples the result back to full resolution and borrows fine detail from the convolutional layers along the way. The paper reports that this beats earlier convolutional methods on abdominal CT and on cardiac MRI.

TL;DR

• Hybrid segmentation network: a ResNet-50 stem, a ViT-B/16 Transformer encoder, and a U-shaped decoder (called CUP) with skip connections §3.2.
• Reported on Synapse multi-organ CT (8 organs) with 77.48% average DSC and 31.69 mm HD; the best CNN baseline in the same table is R50-AttnUNet at 75.57% DSC tab. 1.
• Reported on ACDC cardiac MRI with 89.71% average DSC tab. 5.
• The paper leaves the loss, the learning-rate schedule and the checkpoint choice unstated; the released code fills these in (see Paper vs code).

Problem

Convolutional networks (FCN, U-Net) are the default for medical segmentation, but their locality makes long-range relations hard to model, which hurts targets whose texture, shape and size vary a lot between patients §1. Self-attention added on top of CNN features is one fix §1, §2. Applying a pure Transformer directly to image patches and upsampling the output is a naive alternative, and the paper reports it gives a coarse result because the low-resolution token grid lacks localization detail §1, §3.2.

Method

The model has three parts: a CNN stem that produces a low-resolution feature map, a Transformer encoder over that map, and a cascaded upsampler (CUP) that decodes back to full resolution and receives skip connections from the CNN §3.2. The authors call the combined design a U-shaped architecture §3.2.

1. Transformer encoder. Images are tokenized, embedded linearly, given learned position embeddings, and passed through L layers of multi-head self-attention (MSA) and MLP blocks with pre-layer-norm residuals §3.1 eq. 1 eq. 2 eq. 3.
2. Hybrid stem. The Transformer does not see raw pixels. A ResNet-50 stem produces the feature map, and each cell of that map is one token, so a 1×1 patch on the CNN map replaces the usual 16×16 pixel patch §3.2. The paper compares this hybrid against a pure ViT encoder tab. 1.
3. Cascaded upsampler. The token sequence is reshaped to its 2-D grid and decoded by repeated upsampling blocks. Each block doubles the spatial size, then applies a 3×3 convolution followed by a ReLU §3.2. Four such blocks reach full resolution §4.2.
4. Skip connections. Feature maps from the CNN at 1/8, 1/4 and 1/2 of the input resolution are concatenated into the first three decoder blocks. The best configuration in the paper uses three skips §4.4 fig. 1 fig. 2.
5. Segmentation head. A 3×3 convolution maps the last decoder features to one channel per class at full resolution vit_seg_modeling.py vit_seg_modeling.py. The text does not describe this head; softmax is applied afterwards at inference utils.py.

The released code also includes a 3×3 convolution to 512 channels in front of the decoder, and two convolutions per decoder block, each with batch normalisation; the text describes a single convolution per block §3.2. See Paper vs code.

# TransUNet forward pass  [§3.2] [eq. 1-3]
# Input:  image (B, 3, 224, 224)
# Output: logits (B, n_classes, 224, 224)

def transunet(image):
    # 1. ResNet-50 stem and stages 1-3; keep skips at 1/8, 1/4, 1/2  [§4.4]
    # cnn_map: (B, 1024, 14, 14); skips: 512 ch at 28, 256 at 56, 64 at 112
    cnn_map, skips = resnet50_features(image)
    # 2. Each cnn_map cell becomes one token (1x1 patch)  [§3.2]
    # (B, 196, 768)
    tokens = flatten(linear_1x1(cnn_map))
    # 3. Add learned position embeddings  [eq. 1]
    tokens = tokens + position_embedding
    # 4. Transformer encoder: 12 layers of MSA and MLP  [eq. 2] [eq. 3]
    for layer in range(12):
        # pre-LN multi-head self-attention, with residual  [eq. 2]
        tokens = tokens + attention(layer_norm(tokens))
        # pre-LN MLP, with residual  [eq. 3]
        tokens = tokens + mlp(layer_norm(tokens))
    # 5. Reshape the token sequence back to a 2-D grid  [§3.2]
    # (B, 768, 14, 14)
    grid = reshape_to_grid(tokens, 14, 14)
    # 6. 3x3 conv to 512 channels, batch norm, ReLU
    # [code: vit_seg_modeling.py#L331-L337]
    x = relu(batch_norm(conv3x3(grid, 512)))
    # 7. Four upsampling blocks; first three receive a skip  [§3.2] [§4.4]
    # Output widths: 256, 128, 64, 16
    for width, skip in zip((256, 128, 64, 16), skips + [None]):
        # bilinear 2x upsampling  [code: vit_seg_modeling.py#L284-L313]
        x = upsample_bilinear_2x(x)
        if skip is not None:
            x = concat(x, skip)
        x = conv3x3_bn_relu(x, width)
        x = conv3x3_bn_relu(x, width)
    # 8. 3x3 conv to the class count at 224 x 224
    # [code: vit_seg_modeling.py#L378-L390]
    logits = conv3x3(x, n_classes)
    return logits

Training objective, as released in the code (the paper does not state it):

# Training loss  [code: trainer.py#L56-L58] [code: utils.py#L22-L45]
# Input:  logits (B, C, H, W); labels (B, H, W) with C = 9 for Synapse

def training_loss(logits, labels):
    # Cross-entropy over all pixels  [code: trainer.py#L56]
    loss_ce = cross_entropy(logits, labels)
    # Soft Dice, averaged over all C classes including background
    # [code: utils.py#L32-L45]
    loss_dice = soft_dice(softmax(logits), labels)
    # Equal weighting of the two terms  [code: trainer.py#L58]
    return 0.5 * loss_ce + 0.5 * loss_dice

Hyperparameters

Results

Author-reported. Setup: Synapse multi-organ CT (8 abdominal organs, 18 training and 12 evaluation cases), input 224 × 224, default TransUNet = R50-ViT-B/16 with CUP decoder and three skips §4.1 §4.2. Baselines are V-Net, DARR, U-Net and AttnUNet, with U-Net and AttnUNet given an ImageNet-pretrained ResNet-50 encoder §4.3. The 12 cases the paper calls validation are the ones in the code's test_vol list §4.1 test_vol.txt. Metric: average DSC (%) and average Hausdorff distance (HD, mm) over organs; the paper does not say whether HD is HD95 §4.1.

Synapse, all rows tab. 1. Values as printed in the paper.

Note: the R50-AttnUNet aorta value (55.92) does not reproduce its average; see Pitfalls.

Input resolution tab. 2. Patch size 16 throughout; the 512 input has a sequence about 5× longer §4.4.

Patch size and sequence length tab. 3.

Model scale tab. 4.

Skip connections fig. 2. The figure plots average DSC against the number of skips (0, 1, 3). The text reports that more skips generally help and that three skips give the best average DSC and HD §4.4. No numeric values are given in the text for the 0- and 1-skip cases. An additive-attention gate in the 1/8 skip adds 1.4% DSC in the paper's text, with no table §4.4.

ACDC cardiac MRI tab. 5. Average DSC over RV, Myo and LV, on the ACDC challenge data. The paper's split is 70 / 10 / 20 cases §4.1. The text only says the trend matches Synapse, and it names the ACDC baselines §4.6.

Qualitative results fig. 3. Visual comparison of ground truth, TransUNet, R50-ViT-CUP, R50-AttnUNet and R50-U-Net on Synapse. The text reads it as TransUNet producing fewer false positives and finer boundaries §4.5. No numbers are given with the figure.

Paper vs code

Official code: https://github.com/Beckschen/TransUNet at code_commit 02ef0010b36eb8328b5e689eadaf613602edf9b8, Apache-2.0. Code cited below is from that commit.

• Loss. Paper: not stated. Code: 0.5 cross-entropy + 0.5 soft Dice over all classes including background trainer.py utils.py.
• Learning-rate schedule. Paper: learning rate 0.01 §4.2. Code: polynomial decay with power 0.9 applied per iteration trainer.py.
• Training length. Paper: 14k iterations for Synapse and 20k for ACDC §4.2. Code: the default is 150 epochs train.py; the trainer runs max_epochs × len(trainloader) iterations trainer.py. The --max_iterations flag is parsed but not used train.py. Derived, not stated: 2211 slices at batch 24 give 93 iterations per epoch, so 150 epochs is 13,950 iterations, close to the paper's 14k.
• Skip count for n_skip = 1. Paper: the single skip is added "at the 1/4 resolution scale" §4.4. Code: the first decoder block receives the first feature in the list, which is the 1/8-scale map (28 × 28, 512 channels) vit_seg_modeling.py vit_seg_modeling.py vit_seg_modeling_resnet_skip.py. The 3-skip layout matches the paper's 1/8, 1/4, 1/2 description.
• Decoder structure. Paper: one 3×3 convolution per upsampling block §3.2. Code: a 3×3 convolution to 512 channels with batch norm and ReLU at the decoder input, then two 3×3 convolutions with batch norm and ReLU per block vit_seg_modeling.py vit_seg_modeling.py.
• Hybrid stem. Paper: "a CNN" with no details §3.2. Code: a weight-standardised 7×7 stride-2 convolution with GroupNorm and ReLU, a max-pool applied in the forward pass, then ResNet-50 units (3, 4, 9) vit_seg_modeling_resnet_skip.py vit_seg_modeling_resnet_skip.py vit_seg_configs.py.
• Pretraining. Paper: "pretrained on ImageNet" §4.2. Code: config paths name ImageNet-21k checkpoints for ViT-B/16 and R50+ViT-B/16 vit_seg_configs.py vit_seg_configs.py. The checkpoints are not in the repository.
• Evaluation metric. Paper: "Hausdorff Distance (HD)" §4.1. Code: HD95 per organ, with mean HD95 logged utils.py test.py.
• Checkpoint used for evaluation. Paper: not stated. Code: the trainer never writes best_model.pth. The test script loads epoch_<max_epochs-1>.pth, so the last epoch is evaluated, with no selection on the held-out cases trainer.py test.py. The test script's default --max_epochs is 30, so a default-trained model must be evaluated with --max_epochs 150 test.py.
• Inference resampling. Paper: "slice-by-slice" §4.2. Code: each 2D slice is cubic-resampled to 224 × 224, predicted, and nearest-neighbour resampled back to its original size utils.py.
• Training augmentation. Paper: random rotation and flipping §4.2. Code: with probability 0.5, a random 90° rotation and a random flip; otherwise with probability 0.5, an integer rotation in [-20, 19] degrees dataset_synapse.py. Training slices are resized to 224 with cubic interpolation on the image and nearest on the label dataset_synapse.py.
• ACDC. Paper reports ACDC results tab. 5. Code: train.py and test.py configure only Synapse; no ACDC loader is present in the repository.
• Synapse training slices. Paper: 2212 axial slices §4.1. The file lists/lists_Synapse/train.txt has 2211 lines.

No further discrepancies were found in the encoder configuration: hidden size, layer count, MLP width and head count match the paper's Base and Large values §4.4 vit_seg_configs.py.

Pitfalls

• Table 1, R50-AttnUNet row. The per-organ DSC values in that row average to 72.10, not the 75.57 printed as the average (checked by hand). The aorta value 55.92 is the likely outlier, since about 83.7 would be needed to match the average. The other rows checked reproduce their averages. Treat that row with care when comparing against it.
• Loss, schedule and checkpoint are not in the paper. An implementation that follows only the text will use a constant learning rate and no Dice term, and will not match the reported numbers. Use the code's values trainer.py.
• Base model sentence. The model-scaling paragraph gives four numbers for Base, and the labels it lists for them are in a different order from the numbers. Read literally, it would put the value 12 on the hidden size. The code and the other parts of the paper give hidden size 768 and 12 layers §4.4 vit_seg_configs.py.
• Skip scale for one skip. The text's 1/4 for the single skip §4.4 conflicts with the code's 1/8 (see Paper vs code). The ablation numbers in the text are for 0, 1 and 3 skips, so the 1-skip result depends on which layout is implemented.
• Cost of higher resolution. The 512 input gives a 5× longer sequence and a 6.88-point DSC gain; the paper calls the compute cost "much larger" §4.4 tab. 2. Patch size 8 gives the best average DSC (77.83), at a sequence length of 784 tab. 3. The default of 224 input and patch 16 is the paper's choice for cost reasons §4.4.
• Memory limit. The additive-attention skip was used at 1/8 only, because of GPU memory §4.4. The gain (+1.4% DSC) has no table, and this card found no configuration switch for that variant in the released code.
• Single run, no error bars. The paper reports one number per cell and does not state seeds or variance tab. 1–5. The code fixes seed 1234 and sets deterministic cuDNN train.py train.py. Differences of about one DSC point between rows should be treated as uncertain.
• Metric edge cases. The code scores an empty prediction against a non-empty ground truth as DSC 0 and HD 0, and a non-empty prediction against an empty ground truth as DSC 1 and HD 0 utils.py. Both HD values are the best possible, so a missed organ or a false organ can both look favourable in HD. Check whether your protocol should count those cases separately.
• Background in the loss. The Dice term averages over all nine classes, background included utils.py. Evaluation averages over the eight organs only utils.py.
• Pretrained weights are external. The configs refer to ../model/vit_checkpoint/... paths that are not in the repository vit_seg_configs.py vit_seg_configs.py. Reproduction needs the ImageNet-21k checkpoints from elsewhere.
• Slice-by-slice 2D inference. Each 2D slice is predicted on its own, and the 3D volume is stacked from the slices §4.2. The released model takes one slice at a time, so no cross-slice context is used utils.py.
• Organ list typo. The Synapse text lists "spleen" twice in the eight-organ list §4.1. The table has one spleen column and eight organs.
• Hardware and runtime. Single RTX 2080 Ti §4.2. No wall-clock time or parameter count is given in the text.

Lineage

• Builds on: U-Net (Ronneberger et al., 2015) for the U-shaped encoder-decoder with skips §1 §3.2; ViT (Dosovitskiy et al., 2020) for the patch-token Transformer encoder §2 §3.1; ResNet-50 for the hybrid CNN stem §4.2.
• Compared against: V-Net, DARR, U-Net and AttnUNet (Schlemper et al., 2019) on Synapse tab. 1; R50-U-Net, R50-AttnUNet, ViT-CUP and R50-ViT-CUP on ACDC tab. 5.
• Superseded by: not known from the content of this wiki.