ShERPA: Leveraging Neuron Alignment for Knowledge-preserving Fine-tuning

Machine learning models commonly face challenges in maintaining robustness under distribution shifts. An effective strategy to address this issue involves fine-tuning selected layers of pre-trained models for adaptation. However, the lack of clear criteria for selecting trainable layers poses a significant obstacle, inevitably distorting pre-trained knowledge amidst fine-tuning. This paper proposes a novel approach to this problem by analyzing the loss landscape of trained networks. By drawing insight from recent studies on neuron alignment, we conjecture that aligning models in their loss landscape will minimize the knowledge distortion during the fine-tuning process. Reflecting this, we introduce a novel fine-tuning framework, named ShERPA (Shifted basin for Enhanced Robustness via Permuted Activations), which shifts the training model towards the loss basin of a trained anchor model to encourage the preservation of pre-trained features. Empirical results demonstrate the effectiveness of ShERPA in enhancing Out-of-Distribution (OOD) robustness in multiple benchmarks (PACS, Terra Incognita, VLCS), outperforming the anchor without incurring additional overhead from gradient computations. Our work provides fresh perspectives in understanding how neural networks preserve and tune knowledge in the face of distribution shifts.