About the work
A digital pathology slide is not an ordinary image. At high resolution, a single whole-slide image (WSI) can contain billions of pixels. Rather than feeding the entire slide into a neural network at once, modern systems divide it into hundreds or thousands of smaller image patches.
These patches are typically passed through a neural-network encoder, which turns each patch into a compact numerical representation called an embedding. A second model (an aggregator like an MIL (multiple instance learning) head) then combines embeddings of all patches to make a prediction for the WSI, for example, assigning a cancer grade.
There is a catch, though.
Training the entire system end-to-end (E2E) requires the GPU to retain large amounts of intermediate information (neural network activations) so it can propagate errors back and train the encoder. For hundreds of patches, this can quickly exceed the GPU's memory. As a result, a common solution is to freeze the encoder and train only the MIL model.
In a new paper accepted in Transactions on Machine Learning Research (TMLR), we investigate what happens when we remove this restriction.
But can we train the complete model E2E without requiring an enormous GPU, and does allowing the encoder to train actually improve performance?
Breaking a large computation into manageable pieces
We developed memory-efficient training strategies for both convolutional neural networks and modern Vision Transformer (ViT) based foundation models.
For convolutional networks, we split large intermediate tensors into smaller pieces. The GPU processes these pieces one at a time. The data that is not required for the current piece stays in the much larger system memory RAM. The method also handles operations such as Batch Normalisation, which require global statistics: information from the complete input rather than from isolated chunks.
For ViTs, we use a different strategy. Instead of storing the internal activations for every patch until the backward pass, the model first computes patch representations in small micro-batches. After the slide-level model has produced a loss, each micro-batch is recomputed during backpropagation. This lets us recover the full encoder gradient while keeping only a small number of encoder activations in GPU memory at a time.
Importantly, these techniques change how the computation is done, not what computation is done. The final result is the same as if the GPU were large enough to hold all activations in memory.
Freezing is not always the best choice, but neither is fine-tuning
We evaluated the approach on prostate-cancer and breast-cancer pathology datasets, comparing conventional frozen encoders with fully trainable models.
For a ResNet-18 encoder, E2E training consistently improved performance on the training domain. On the external prostate-cancer dataset TCGA-PRAD, the strongest results came from combining end-to-end training with WSI-level data augmentation.
We also studied H0-mini, a compact histopathology foundation model. When fine-tuned E2E with WSI-level augmentation and a regularisation technique called Learning without Forgetting (LwF), H0-mini could match or substantially outperform larger frozen foundation models on TCGA-PRAD.
However, the results also showed that fine-tuning is not automatically better. On the external CAMELYON16 dataset, frozen foundation models remained stronger. The study therefore argues against a simple rule that encoders should always be frozen or always fine-tuned.
Instead, the usefulness of encoder adaptation depends on the dataset size, the evaluation setting, and the complete training recipe.
Making encoder adaptation practical
The broader goal of the work is not to replace frozen models altogether, but to make full E2E training a practical option for researchers working with WSIs.
By reducing the GPU-memory barrier, the method lets researchers directly investigate how pretrained representations should be adapted to a particular pathology task, rather than freezing the encoder simply because full training does not fit in memory. For example, this method opens the door to using domain generalisation techniques which require adapting the encoder.
Paper: https://openreview.net/forum?id=ixDBhyLgWs
Code: https://github.com/undercutspiky/big-layers