Exploiting Layer Normalization Fine-tuning in Visual Transformer Foundation Models for Classification

TL;DR

This paper investigates the role of LayerNorm in fine-tuning Vision Transformers for classification, revealing how LayerNorm shifts relate to domain shifts and proposing a rescaling method to improve transfer learning, especially with limited data.

Contribution

It introduces a novel analysis of LayerNorm shifts during fine-tuning, proposes a rescaling mechanism based on the Fine-tuning Shift Ratio, and demonstrates its effectiveness across various datasets and settings.

Findings

LayerNorm shifts indicate domain transfer quality.

The proposed rescaling improves fine-tuning performance.

OOD tasks show lower FSR and higher λ, especially with scarce data.

Abstract

LayerNorm is pivotal in Vision Transformers (ViTs), yet its fine-tuning dynamics under data scarcity and domain shifts remain underexplored. This paper shows that shifts in LayerNorm parameters after fine-tuning (LayerNorm shifts) are indicative of the transitions between source and target domains; its efficacy is contingent upon the degree to which the target training samples accurately represent the target domain, as quantified by our proposed Fine-tuning Shift Ratio ($FSR$). Building on this, we propose a simple yet effective rescaling mechanism using a scalar $\lambda$ that is negatively correlated to $FSR$ to align learned LayerNorm shifts with those ideal shifts achieved under fully representative data, combined with a cyclic framework that further enhances the LayerNorm fine-tuning. Extensive experiments across natural and pathological images, in both in-distribution (ID) and out-of-distribution (OOD) settings, and various target training sample regimes validate our framework. Notably, OOD tasks tend to yield lower $FSR$ and higher $\lambda$ in comparison to ID cases, especially with scarce data, indicating under-represented target training samples. Moreover, ViTFs fine-tuned on pathological data behave more like ID settings, favoring conservative LayerNorm updates. Our findings illuminate the underexplored dynamics of LayerNorm in transfer learning and provide practical strategies for LayerNorm fine-tuning.