Source-linked AI summary
Partial transfusion: on the expressive influence of trainable batch norm parameters for transfer learning
Fahdi Kanavati, Masayuki Tsuneki
TL;DR
ImageNet transfer learning is widely used for medical images, though its performance advantage over training from scratch remains uncertain. This paper compares fine-tuning only batch-normalisation affine parameters with full fine-tuning across medical-imaging datasets and architectures. The affine-only approach has similar performance in almost all experiments and is reported to converge faster.
Problem
ImageNet transfer learning is widely used for medical imaging, but its evaluation-performance advantage over training from scratch remains uncertain.
Source: §1 · p. 1Method
The study compares batch-normalisation affine-only fine-tuning with full fine-tuning and other methods across medical-imaging datasets and four architectures.
Source: §1 · p. 2Results
Batch-normalisation affine-only fine-tuning performs similarly to full fine-tuning in almost all experiments, and fine-tuning improves performance over feature extraction.
Source: §4 · pp. 4–5Takeaways & Limitations
Affine-only fine-tuning is reported to converge faster while achieving similar performance to full-model fine-tuning across the studied medical-imaging datasets.
Source: §5 · p. 5, §6 · p. 6The study did not conduct hyperparameter-search optimisation or investigate the primary mechanism behind the sufficiency of batch-normalisation parameter tuning.
Source: §5 · p. 6
Abstract
from arXiv · showhide
Transfer learning from ImageNet is the go-to approach when applying deep learning to medical images. The approach is either to fine-tune a pre-trained model or use it as a feature extractor. Most modern architecture contain batch normalisation layers, and fine-tuning a model with such layers requires taking a few precautions as they consist of trainable and non-trainable weights and have two operating modes: training and inference. Attention is primarily given to the non-trainable weights used during inference, as they are the primary source of unexpected behaviour or degradation in performance during transfer learning. It is typically recommended to fine-tune the model with the batch normalisation layers kept in inference mode during both training and inference. In this paper, we pay closer attention instead to the trainable weights of the batch normalisation layers, and we explore their expressive influence in the context of transfer learning. We find that only fine-tuning the trainable weights (scale and centre) of the batch normalisation layers leads to similar performance as to fine-tuning all of the weights, with the added benefit of faster convergence. We demonstrate this on a variety of seven publicly available medical imaging datasets, using four different model architectures.
1 Introduction
ImageNet transfer learning is widely used for medical imaging, but its performance advantage over training from scratch remains uncertain. This paper examines whether tuning batch-normalisation affine parameters can match full fine-tuning.
1 Introduction
- Fine-tuning typically updates all layers or a subset, while other work has learned which layers to tune based on the input.The cited approaches use low learning rates to avoid destroying pre-trained weights.Source: §1 · p. 1
- ImageNet transfer learning’s advantage over training from scratch remains uncertain given natural–medical image differences.Prior evidence cited in the introduction found little evaluation-performance benefit on some chest X-ray and diabetic-retinopathy tasks, with lightweight models trained from scratch performing comparably to larger fine-tuned models.Source: §1 · p. 1
- Batch normalisation has trainable affine parameters, scale γ and offset β, as well as non-trainable moving statistics used at inference.Updating moving averages with limited data can create a discrepancy between training and evaluation; the paper instead investigates the affine parameters’ expressive influence.Source: §1 · p. 2
- The paper compares five methods across nine datasets from seven publicly available medical-imaging datasets and four architectures.Its reported findings include similar performance for batch-normalisation affine-only and full fine-tuning, improved performance from fine-tuning versus feature extraction, and acceptable random-weight results on some datasets.Source: §1 · p. 2
2 Datasets
The experiments cover nine datasets derived from seven public medical-imaging datasets, spanning chest imaging, blood cells, retinal OCT, histopathology, and colorectal histology. The dataset descriptions specify distinct sampling, labeling, and split choices.
2 Datasets
- The study uses nine datasets originating from seven publicly available medical-imaging datasets.The described tasks include chest X-rays, thoracic disease classification, malaria cell images, OCT, lymph-node metastases, and colorectal histology.Source: §2 · p. 2
- CheXpert is restricted to five pathologies and frontal images, with uncertain labels excluded and patient overlap prevented across the balanced splits.The sampled sets contain 11,633 training, 1,036 validation, and 1,350 test images.Source: §2 · p. 2
- Patch Camelyon 16 uses original train, validation, and test splits, with an additional smaller training version; colorectal histology uses a random 70/15/15 split.Patch Camelyon 16’s smaller version has 6,400 training and 960 validation images.Source: §2 · pp. 2–3
- Patch Camelyon 17 groups breast-cancer metastasis patches by hospital, using one hospital for training and other hospitals for validation and separate test sets.The full dataset includes 450,000 patches; a smaller training version uses 6,400 training and 960 validation patches.Source: §2 · pp. 2–3
3 Method
The method compares transfer-learning and random-weight training variants using four CNN architectures. It distinguishes trainable batch-normalisation affine parameters from the layer’s mode-dependent moving statistics.
3.1 Models
- Each pre-trained model replaces its classification head with global average pooling when needed, dropout, and a dataset-specific fully connected classifier.The classifier uses softmax for categorical labels and sigmoid otherwise; four architectures are evaluated, with their trainable batch-normalisation parameter counts listed in Table 1.Source: §3.1 · p. 3, Table 1 · p. 3
3.2 Batch normalisation layer
- Batch-normalisation scale γ and offset β are trainable, while mean µ and standard deviation σ are estimated from data and moving averages are used in inference.The authors distinguish a layer being trainable (affine parameters can change) from being in training mode (µ and σ are computed from the current batch and their averages updated); limited-data updates can cause training–evaluation discrepancies.Source: §3.2 · p. 4
3.3 Training
- Five methods comprise feature extraction, full fine-tuning, affine-only batch-normalisation tuning, moving-average updates, and random-weight affine-only tuning.The fine-tuning variants keep batch normalisation in inference mode; FC-BN trains the FC layer for one epoch before training it with the BN affine parameters, while FC-full first trains the FC layer to validation-loss plateau.Source: §3.3 · p. 4
- Training uses Adam, an initial learning rate of 1e-3 except 1e-5 for full-model fine-tuning, step decay, and early stopping after six unimproved epochs.All datasets use batch size 32 except Camelyon at 64, with minimal translation, zoom, and dataset-dependent flipping augmentation.Source: §3.3 · p. 4
4 Results
Across nearly all experiments, tuning only batch-normalisation parameters performs similarly to full-model fine-tuning. Fine-tuning generally improves on feature extraction, whereas updating moving averages or using random weights usually performs less well.
4 Results
- Fine-tuning with FC-then-full or FC-then-BN improves performance over feature extraction, while updating moving averages tends to degrade it versus keeping them fixed.The reported results are based on three repetitions per method, model, and dataset, using test-set ROC AUC averaged across labels or test subsets where applicable.Source: §4 · pp. 4–5
- FC-then-BN achieves similar ROC AUC to FC-then-full in almost all experiments across five methods, nine datasets, and four architectures.The exceptions with slight underperformance are InceptionV3 on CheXpert and both small Patch Camelyon datasets, and ResNet50V2 on CheXpert; InceptionV3 has the fewest trainable BN parameters but is comparable on some datasets.Source: §4 · p. 4, Figure 2 · p. 5
- Random-weight BN-affine tuning mostly has the lowest performance, though DenseNet121 achieves acceptable results on Chest X-ray and Malaria.EfficientNetB3 with random weights can remain saturated and fail to train; training-mode BN allowed training on some datasets but still gave it the lowest performance among the four models.Source: §4 · p. 5
5 Discussion
Across medical imaging datasets spanning x-ray, OCT, and histopathology, fine-tuning only batch norm affine parameters yields performance similar to fine-tuning all model parameters. It converges faster, while the two tested training orders perform similarly overall.
5 Discussion
- Fine-tuning only batch norm affine parameters achieves performance similar to fine-tuning all model parameters across varied medical imaging datasets.The datasets span x-ray, OCT, and histopathology.Source: §5 · p. 5
- Fine-tuning fewer parameters allows faster convergence and use of a higher learning rate.Source: §5 · p. 5
- The study did not optimize hyperparameters beyond settings used in common practice, given the large number of experiments.Source: §5 · p. 6
- The study did not investigate the primary mechanism behind why fine-tuning batch norm parameters is sufficient.Source: §5 · p. 6
6 Conclusion
Across varied medical imaging datasets, fine-tuning only batch norm affine parameters performs similarly to fine-tuning all model parameters. This approach uses fewer trainable parameters and converges faster, with results observed across four architectures, particularly DenseNet121.
6 Conclusion
- Fine-tuning only batch norm affine parameters achieves performance similar to fine-tuning all model parameters across varied medical imaging datasets.Source: §6 · p. 6
- The approach converges faster while training fewer parameters, without loss in performance.The conclusion attributes faster convergence to use of a higher learning rate and fine-tuning fewer parameters.Source: §6 · p. 6
- The result was observed with four model architectures, particularly DenseNet121, highlighting the expressive power of scaling and offsetting pretrained convolutional-layer outputs.Source: §6 · p. 6
A Performances on test sets using DenseNet121
The appendix tables report ROC AUC mean and standard deviation results for DenseNet121 on four medical imaging datasets, with outcomes organized by dataset-specific test sets, labels, or classes.
A Performances on test sets using DenseNet121
- Patch Camelyon 17 results report ROC AUC mean and standard deviation for three test sets originating from different hospitals.Source: Table A2 · p. 7
- CheXpert results report ROC AUC mean and standard deviation for five labels.Source: Table A3 · p. 8
- OCT small results report ROC AUC mean and standard deviation for four labels.Source: Table A4 · p. 8
- Colorectal Histology results report ROC AUC mean and standard deviation for eight labels.Source: Table A5 · p. 8
B Random weights
With random weights, EfficientNetB3 could not train batch norm parameters on some datasets because outputs remained mostly saturated and gradients could not flow. The paper compares batch norm operating modes and suggests training mode may improve conditions for training.
B Random weights
- EfficientNetB3 with random weights could not train batch norm parameters on some datasets when saturated outputs prevented gradients from flowing.Source: B Random weights · p. 8
- Setting batch norm layers to training mode allowed the model to train, potentially because additional activation standardization reduced saturation.Source: B Random weights · p. 8
- Figure B3 compares ROC AUC for random-weight DenseNet121 and EfficientNetB3 with trainable batch norm in inference mode versus training mode.Source: B Random weights · p. 8, Figure B3 · p. 9
C Training method with batch norm layer
Three strategies vary when the batch norm parameters are trained alongside the FC layer. FC-BN and FC-then-BN perform similarly, while FC-then-BN-lr occasionally performs worse.
C Training method with batch norm layer
- FC-BN trains the FC layer and batch norm parameters together from the start; FC-then-BN trains the FC layer for one epoch before adding batch norm parameters.Both use an initial learning rate of 1e-3.Source: C Training method with batch norm layer · p. 9
- FC-then-BN-lr trains the FC layer until validation performance stops improving, then trains it with batch norm parameters at an initial learning rate of 1e-5.Source: C Training method with batch norm layer · p. 9
- FC-BN and FC-then-BN show similar performance, while FC-then-BN-lr is occasionally lower, suggesting some benefit from a higher learning rate.Figure C4 reports ROC AUC for DenseNet121 and EfficientNetB3 across the three strategies.Source: C Training method with batch norm layer · p. 9, Figure C4 · p. 10
D Batch norm in training mode
The section compares batch norm layers in training versus inference mode across three main transfer-learning methods. Inference mode performs best overall for all three.
D Batch norm in training mode
- Each of the three main transfer-learning methods is evaluated with batch norm layers in either training or inference mode.Source: D Batch norm in training mode · p. 10, Figure D5 · p. 11
- Setting batch norm layers to inference mode performs best overall for all three main transfer-learning methods.Figure D5 reports ROC AUC for this comparison.Source: D Batch norm in training mode · p. 10, Figure D5 · p. 11
- The comparison uses ROC AUC to assess the two batch norm modes across the three methods.Source: Figure D5 · p. 11