Channel-Selective Normalization for Label-Shift Robust Test-Time Adaptation

Pedro Vianna, Muawiz Sajjad Chaudhary, Paria Mehrbod, An Tang, Guy Cloutier, Guy Wolf, Michael Eickenberg, Eugene Belilovsky
Proceedings of The 3rd Conference on Lifelong Learning Agents, PMLR 274:514-533, 2025.

Abstract

Deep neural networks have useful applications in many different tasks, however their performance can be severely affected by changes in the input distribution. For example, in the biomedical field, their performance can be affected by changes in the data (different machines, populations) between training and test datasets. To ensure robustness and generalization to real-world scenarios, test-time adaptation has been recently studied as an approach to adjust models to a new data distribution during inference. Test-time batch normalization is a simple and popular method that achieved compelling performance on domain shift benchmarks. It is implemented by recalculating batch normalization statistics on test batches. Prior work has focused on analysis with test data that has the same label distribution as the training data. However, in many practical applications this technique is vulnerable to label distribution shifts, sometimes causing catastrophic failure. This presents a risk in applying test time adaptation methods in deployment. We propose to tackle this challenge by only selectively adapting channels in a deep network, minimizing drastic adaptation that is sensitive to label shifts. Our selection scheme is based on two principles that we empirically motivate: (1) later layers of networks are more sensitive to label shift; (2) individual features can be sensitive to specific classes. We apply the proposed technique to three classification tasks, including CIFAR10-C, Imagenet-C, and diagnosis of fatty liver, where we explore both covariate and label distribution shifts. We find that our method allows to bring the benefits of TTA while significantly reducing the risk of failure common in imbalanced scenarios.

Cite this Paper

BibTeX
@InProceedings{pmlr-v274-vianna25a, title = {Channel-Selective Normalization for Label-Shift Robust Test-Time Adaptation}, author = {Vianna, Pedro and Chaudhary, Muawiz Sajjad and Mehrbod, Paria and Tang, An and Cloutier, Guy and Wolf, Guy and Eickenberg, Michael and Belilovsky, Eugene}, booktitle = {Proceedings of The 3rd Conference on Lifelong Learning Agents}, pages = {514--533}, year = {2025}, editor = {Lomonaco, Vincenzo and Melacci, Stefano and Tuytelaars, Tinne and Chandar, Sarath and Pascanu, Razvan}, volume = {274}, series = {Proceedings of Machine Learning Research}, month = {29 Jul--01 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v274/main/assets/vianna25a/vianna25a.pdf}, url = {https://proceedings.mlr.press/v274/vianna25a.html}, abstract = {Deep neural networks have useful applications in many different tasks, however their performance can be severely affected by changes in the input distribution. For example, in the biomedical field, their performance can be affected by changes in the data (different machines, populations) between training and test datasets. To ensure robustness and generalization to real-world scenarios, test-time adaptation has been recently studied as an approach to adjust models to a new data distribution during inference. Test-time batch normalization is a simple and popular method that achieved compelling performance on domain shift benchmarks. It is implemented by recalculating batch normalization statistics on test batches. Prior work has focused on analysis with test data that has the same label distribution as the training data. However, in many practical applications this technique is vulnerable to label distribution shifts, sometimes causing catastrophic failure. This presents a risk in applying test time adaptation methods in deployment. We propose to tackle this challenge by only selectively adapting channels in a deep network, minimizing drastic adaptation that is sensitive to label shifts. Our selection scheme is based on two principles that we empirically motivate: (1) later layers of networks are more sensitive to label shift; (2) individual features can be sensitive to specific classes. We apply the proposed technique to three classification tasks, including CIFAR10-C, Imagenet-C, and diagnosis of fatty liver, where we explore both covariate and label distribution shifts. We find that our method allows to bring the benefits of TTA while significantly reducing the risk of failure common in imbalanced scenarios.} }
Endnote
%0 Conference Paper %T Channel-Selective Normalization for Label-Shift Robust Test-Time Adaptation %A Pedro Vianna %A Muawiz Sajjad Chaudhary %A Paria Mehrbod %A An Tang %A Guy Cloutier %A Guy Wolf %A Michael Eickenberg %A Eugene Belilovsky %B Proceedings of The 3rd Conference on Lifelong Learning Agents %C Proceedings of Machine Learning Research %D 2025 %E Vincenzo Lomonaco %E Stefano Melacci %E Tinne Tuytelaars %E Sarath Chandar %E Razvan Pascanu %F pmlr-v274-vianna25a %I PMLR %P 514--533 %U https://proceedings.mlr.press/v274/vianna25a.html %V 274 %X Deep neural networks have useful applications in many different tasks, however their performance can be severely affected by changes in the input distribution. For example, in the biomedical field, their performance can be affected by changes in the data (different machines, populations) between training and test datasets. To ensure robustness and generalization to real-world scenarios, test-time adaptation has been recently studied as an approach to adjust models to a new data distribution during inference. Test-time batch normalization is a simple and popular method that achieved compelling performance on domain shift benchmarks. It is implemented by recalculating batch normalization statistics on test batches. Prior work has focused on analysis with test data that has the same label distribution as the training data. However, in many practical applications this technique is vulnerable to label distribution shifts, sometimes causing catastrophic failure. This presents a risk in applying test time adaptation methods in deployment. We propose to tackle this challenge by only selectively adapting channels in a deep network, minimizing drastic adaptation that is sensitive to label shifts. Our selection scheme is based on two principles that we empirically motivate: (1) later layers of networks are more sensitive to label shift; (2) individual features can be sensitive to specific classes. We apply the proposed technique to three classification tasks, including CIFAR10-C, Imagenet-C, and diagnosis of fatty liver, where we explore both covariate and label distribution shifts. We find that our method allows to bring the benefits of TTA while significantly reducing the risk of failure common in imbalanced scenarios.
APA
Vianna, P., Chaudhary, M.S., Mehrbod, P., Tang, A., Cloutier, G., Wolf, G., Eickenberg, M. & Belilovsky, E.. (2025). Channel-Selective Normalization for Label-Shift Robust Test-Time Adaptation. Proceedings of The 3rd Conference on Lifelong Learning Agents, in Proceedings of Machine Learning Research 274:514-533 Available from https://proceedings.mlr.press/v274/vianna25a.html.

Related Material