Memory Constrained Dynamic Subnetwork Update for Transfer Learning

TL;DR

MeDyate tackles on-device transfer learning under strict memory budgets by combining LaRa-based layer ranking with a dynamic, probability-weighted channel sampling strategy. It leverages the observed stability of channel importance during fine-tuning to explore the parameter space efficiently within memory constraints, updating channels across epochs rather than statically selecting them. Across CNN architectures and seven downstream tasks, MeDyate achieves state-of-the-art performance under extreme budgets (down to a few hundred kB RAM) and demonstrates notable reductions in activation memory and FLOPs, with transformer results indicating task-model-specific considerations. The approach provides a principled, budget-aware mechanism to enable practical on-device fine-tuning, though real-edge deployment remains to be validated with an on-device implementation.

Abstract

On-device neural network training faces critical memory constraints that limit the adaptation of pre-trained models to downstream tasks. We present MeDyate, a theoretically-grounded framework for memory-constrained dynamic subnetwork adaptation. Our approach introduces two key innovations: LaRa (Layer Ranking), an improved layer importance metric that enables principled layer pre-selection, and a dynamic channel sampling strategy that exploits the temporal stability of channel importance distributions during fine-tuning. MeDyate dynamically resamples channels between epochs according to importance-weighted probabilities, ensuring comprehensive parameter space exploration while respecting strict memory budgets. Extensive evaluation across a large panel of tasks and architectures demonstrates that MeDyate achieves state-of-the-art performance under extreme memory constraints, consistently outperforming existing static and dynamic approaches while maintaining high computational efficiency. Our method represents a significant step towards enabling efficient on-device learning by demonstrating effective fine-tuning with memory budgets as low as a few hundred kB of RAM.

Figures (11)

• Figure 1: Overview of the MeDyate framework. The offline phase (top) shows layer pre-selection based on LaRa rankings and memory budget constraints, while the online phase (bottom) illustrates dynamic channel sampling with importance-weighted probabilities during training epochs.
• Figure 2: Evolution of $u_e$ as a function of epochs for different $\alpha$ values.
• Figure 3: Spearman correlation of layer cumulated LaRa metric across seeds and datasets.
• Figure 4: Evolution of channel gradient norm T-test over time, MobileNetV2 fine-tuned on 3 downstream tasks.
• Figure 5: Evolution of channel gradient norm over time within a specific layer of a MobileNetV2 fine-tuned on CIFAR-10. Results are normalized per epoch for visualization.