Source-linked AI summary
ACC-UNet: A Completely Convolutional UNet model for the 2020s
Nabil Ibtehaz, Daisuke Kihara
TL;DR
This paper asks whether a purely convolutional UNet can compete with transformer-based UNets. ACC-UNet uses convolutional blocks for hierarchical neighborhood context and multi-level feature fusion; across five benchmarks, it performs best among the compared categories, with reported Dice-score improvements over larger state-of-the-art models.
Problem
The paper asks whether a UNet based solely on convolution can compete with transformer-based UNets.
Source: §1 · p. 2Method
ACC-UNet adapts transformer-inspired context aggregation and multi-level encoder-feature fusion using convolutional blocks and modified skip connections.
Source: §1 · p. 2, §2.2 · p. 5, §2.4 · p. 6Results
Across five datasets, ACC-UNet performed best in the compared categories, with Dice-score improvements of 0.13%, 0.10%, 0.63%, 0.90%, and 0.27% over much larger state-of-the-art models.
Source: §3.3 · p. 7Takeaways & Limitations
The experiments suggest that combining CNN inductive bias with transformer-inspired design has the potential to improve UNet models.
Source: §4 · p. 9Concat operations slow the model; replacing them is identified as a possible solution.
Source: §4 · p. 9
Abstract
from arXiv · showhide
This decade is marked by the introduction of Vision Transformer, a radical paradigm shift in broad computer vision. A similar trend is followed in medical imaging, UNet, one of the most influential architectures, has been redesigned with transformers. Recently, the efficacy of convolutional models in vision is being reinvestigated by seminal works such as ConvNext, which elevates a ResNet to Swin Transformer level. Deriving inspiration from this, we aim to improve a purely convolutional UNet model so that it can be on par with the transformer-based models, e.g, Swin-Unet or UCTransNet. We examined several advantages of the transformer-based UNet models, primarily long-range dependencies and cross-level skip connections. We attempted to emulate them through convolution operations and thus propose, ACC-UNet, a completely convolutional UNet model that brings the best of both worlds, the inherent inductive biases of convnets with the design decisions of transformers. ACC-UNet was evaluated on 5 different medical image segmentation benchmarks and consistently outperformed convnets, transformers, and their hybrids. Notably, ACC-UNet outperforms state-of-the-art models Swin-Unet and UCTransNet by $2.64 \pm 2.54\%$ and $0.45 \pm 1.61\%$ in terms of dice score, respectively, while using a fraction of their parameters ($59.26\%$ and $24.24\%$). Our codes are available at https://github.com/kiharalab/ACC-UNet.
1 Introduction
UNet’s skip connections help restore spatial information, while later designs addressed a semantic gap and incorporated transformer or hybrid components. ACC-UNet asks whether a purely convolutional UNet can compete with transformer-based UNets, and its five-benchmark experiments suggest its modifications can improve UNet models.
1 Introduction
- UNet skip connections pass spatial information to the decoder, while U-Net++ [26] and MultiResUNet modify these connections to address a likely encoder-decoder semantic gap.U-Net++ uses dense connections; MultiResUNet adds convolutional blocks along the skip connection.Source: §1 · p. 2
- UNet development expanded from convolutional models to transformer-based Swin-Unet, convolution-transformer hybrids such as UCTransNet, and all-transformer designs such as SMESwin Unet.Swin-Unet replaces convolutional blocks with Swin Transformer blocks; hybrid models use transformers along skip connections.Source: §1 · p. 2
- Inspired by ConvNet designs that enabled ResNets to outperform Swin Transformers, the paper asks whether a purely convolutional UNet can compete with transformer-based UNets.The proposed approach uses patch-based context aggregation and fuses encoder feature maps from multiple levels in skip connections.Source: §1 · p. 2
- Five-benchmark experiments suggest the proposed modifications have the potential to improve UNet models.The paper evaluates ACC-UNet against representative models from five UNet classes.Source: §1 · p. 3
2 Method
The method adapts transformer-inspired long-range context and multi-level feature fusion to convolutional UNet blocks and skip connections. It combines these blocks in a vanilla UNet backbone to form ACC-UNet.
2 Method
- The method designs two convolutional blocks inspired by transformer operations and integrates them into a vanilla UNet backbone as ACC-UNet.The resulting model replaces classical convolutional blocks with HANC and modifies skip connections with MLFC.Source: §2 · p. 3, §2.4 · p. 6
2.1 A high-level view of transformers in UNet
- Transformer-based UNets use long-range context, inverted bottlenecks, and shortcut connections, while their skip connections can combine features from multiple encoder levels.The paper uses these design choices as motivation for modifying convolutional blocks and UNet skip connections.Source: Leveraging the long-range dependency… · p. 4, Adaptive Multi-level feature… · p. 4
2.2 Hierarchical Aggregation of Neighborhood Context (HANC)
- To induce long-range context and expressivity, HANC uses pointwise and depthwise convolutions, expanding channels with an inverted bottleneck and aggregating neighborhood means and maxima hierarchically.Increasing k provides more contextual information without requiring larger convolutional kernels; a shortcut connection is also included.Source: §2.2 · pp. 4–5
2.3 Multi Level Feature Compilation (MLFC)
- MLFC resizes feature maps from different encoder levels to a common size, concatenates and summarizes them with pointwise convolution, then merges the result with the corresponding encoder feature map.This operation is performed individually at each level to enrich its feature map with multi-level information.Source: §2.3 · p. 5
2.4 ACC-UNet
- ACC-UNet replaces vanilla UNet convolutional blocks with HANC, modifies skip connections with residual blocks, and stacks three MLFC blocks.The configuration uses k = 3 for most levels, with k = 1 at the bottleneck and k = 2 at the adjacent level; the model has 16.77M parameters, roughly 2M more than vanilla UNet.Source: §2.4 · p. 6
- ACC-UNet combines convolutional inductive bias with HANC’s approximate self-attention and MLFC’s aggregation of features from different encoder levels.The HANC blocks replace classical convolutional blocks, while MLFC blocks modify skip connections.Source: §2.4 · p. 6
3 Experiments
The experiments evaluate ACC-UNet on five public segmentation datasets under a shared training protocol, compare it with representative UNet models, and ablate its design choices.
3.1 Datasets
- Experiments covered five public datasets across different tasks and modalities; GlaS retained its original test split, while the other datasets used a randomly selected 20% test split and three shuffled repeats.Images and masks were resized to 224 × 224; the remaining data were split into 60% training and 20% validation.Source: §3.1 · p. 6
3.2 Implementation Details
- Training used batch size 12, 1000 epochs with 100-epoch early stopping, Adam at 10^-3 with cosine annealing, combined cross-entropy and Dice loss, and random flips and rotations.Swin-UNet-based models instead used SGD because Adam performed poorly.Source: §3.2 · p. 6
3.3 Comparisons with State-of-the-Art Methods
- ACC-UNet ranked best across all five datasets with much lower parameter counts, improving Dice by 0.13%, 0.10%, 0.63%, 0.90%, and 0.27%, respectively, over larger state-of-the-art models.The comparison reports test-set Dice scores.Source: §3.3 · p. 7, Table 1 · p. 7
3.4 Comparative Qualitative Results on the Five Datasets
- Qualitatively, ACC-UNet follows lesion boundaries, separates a finger from a polyp, rejects a false-positive BUSI nodule, models COVID lung gaps, and recovers GlaS glands.The COVID example had a 2.9% higher Dice score than the second-best method.Source: §3.4 · pp. 7–8, Figure 3 · p. 8
3.5 Ablation Study
- On CVC-ClinicDB, depthwise convolution plus a 4x bottleneck raised Dice from 87.77% to 88.26% while reducing parameters from 7.8M to 7.5M.Adding HANC with k=3 increased parameters by 340% for a 1.1% Dice increase; shortcuts raised performance by 2.16%. Removing MLFC yielded 91.9% Dice, while removing HANC yielded 90.96% with 25% more filters.Source: §3.5 · p. 8, Figure 4 · p. 9
4 Conclusions
ACC-UNet combines CNN inductive biases with transformer-inspired long-range and multi-level feature accumulation. The authors report potential to improve UNet models, while noting concatenation-related slowdown and further transformer-inspired design choices for future work.
4 Conclusions
- ACC-UNet combines CNN inductive bias with long-range and multi-level feature accumulation inspired by transformers.Source: §4 · p. 9
- The experiments suggest that this combination has the potential to improve UNet models.Source: §4 · p. 9
- Concatenation operations slow the model, and the authors propose replacing them while leaving layer normalization, GELU, and AdamW for future exploration.Source: §4 · p. 9
A Remarks on Training Time
ACC-UNet was comparatively slow because its architecture was not thoroughly optimized, but preallocating tensors avoids most concat overhead and reduces training time.
A Remarks on Training Time
- ACC-UNet was comparatively slow because the architecture was not thoroughly optimized.The work prioritized bringing transformer attributes to CNNs over thorough architecture optimization.Source: A Remarks on Training Time · p. 12
- 28.43% lower training time from preallocated-tensor population makes ACC-UNet comparable to UCTransNet.The inference time also benefits from this optimization.Source: Table 2 · p. 12, A Remarks on Training Time · p. 12
- Table 2 reports training and inference time on the BUSI dataset.Source: Table 2 · p. 12