Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Study of Training Dynamics for Memory-Constrained Fine-Tuning

Aël Quélennec, Nour Hezbri, Pavlo Mozharovskyi, Van-Tam Nguyen, Enzo Tartaglione

TL;DR

This work tackles memory-efficient fine-tuning of large models by introducing TraDy, a dynamic channel sampling scheme that updates a task-relevant subnetwork within a strict memory budget. Grounded in heavy-tailed gradient theory, TraDy uses a Reweighted Gradient Norm to identify memory-efficient update opportunities and exploits architecture-driven layer importance to fix a subset of layers while dynamically resampling channels within them. Empirical results across CNNs and Transformer models demonstrate high sparsity in activations and weights, substantial FLOP reductions, and improved or competitive accuracy compared to static approaches, validating both the theoretical and practical viability of memory-constrained on-device learning. The approach holds promise for scalable on-device adaptation of large models, with potential extensions to broader architectures and optimization-theoretic analyses.

Abstract

Memory-efficient training of deep neural networks has become increasingly important as models grow larger while deployment environments impose strict resource constraints. We propose TraDy, a novel transfer learning scheme leveraging two key insights: layer importance for updates is architecture-dependent and determinable a priori, while dynamic stochastic channel selection provides superior gradient approximation compared to static approaches. We introduce a dynamic channel selection approach that stochastically resamples channels between epochs within preselected layers. Extensive experiments demonstrate TraDy achieves state-of-the-art performance across various downstream tasks and architectures while maintaining strict memory constraints, achieving up to 99% activation sparsity, 95% weight derivative sparsity, and 97% reduction in FLOPs for weight derivative computation.

Study of Training Dynamics for Memory-Constrained Fine-Tuning

TL;DR

This work tackles memory-efficient fine-tuning of large models by introducing TraDy, a dynamic channel sampling scheme that updates a task-relevant subnetwork within a strict memory budget. Grounded in heavy-tailed gradient theory, TraDy uses a Reweighted Gradient Norm to identify memory-efficient update opportunities and exploits architecture-driven layer importance to fix a subset of layers while dynamically resampling channels within them. Empirical results across CNNs and Transformer models demonstrate high sparsity in activations and weights, substantial FLOP reductions, and improved or competitive accuracy compared to static approaches, validating both the theoretical and practical viability of memory-constrained on-device learning. The approach holds promise for scalable on-device adaptation of large models, with potential extensions to broader architectures and optimization-theoretic analyses.

Abstract

Memory-efficient training of deep neural networks has become increasingly important as models grow larger while deployment environments impose strict resource constraints. We propose TraDy, a novel transfer learning scheme leveraging two key insights: layer importance for updates is architecture-dependent and determinable a priori, while dynamic stochastic channel selection provides superior gradient approximation compared to static approaches. We introduce a dynamic channel selection approach that stochastically resamples channels between epochs within preselected layers. Extensive experiments demonstrate TraDy achieves state-of-the-art performance across various downstream tasks and architectures while maintaining strict memory constraints, achieving up to 99% activation sparsity, 95% weight derivative sparsity, and 97% reduction in FLOPs for weight derivative computation.
Paper Structure (22 sections, 2 theorems, 13 equations, 13 figures, 5 tables, 1 algorithm)

This paper contains 22 sections, 2 theorems, 13 equations, 13 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Method
  4. Problem Formulation and Notations
  5. Heavy-tailed Theory and Gradient Norm Metric
  6. Layers Behavior During Fine-tuning
  7. Dynamic Channel Sampling
  8. Experiments
  9. Setup
  10. Gradient Study
  11. Main Results
  12. Conclusion
  13. Limitations
  14. Layer Ranking Consistency Detailed Analysis
  15. Additional Experimental Results
  16. ...and 7 more sections

Key Result

Proposition 3.1

The relative ranking of layers to their reweighted gradient norm remains largely invariant over time during training and across different downstream tasks. This ranking is primarily determined by the network architecture rather than dataset-specific characteristics.

Figures (13)

  • Figure 1: TRady dynamically reselects the subgraph to update within the memory budget $B_{\text{mem}}$.
  • Figure 2: Evolution of stochastic gradient heavy-tailed index $\alpha_t$.
  • Figure 3: Spearman correlation of layer gradient norm across seeds and datasets.
  • Figure 4: T-test of channel gradient norm across seeds and datasets.
  • Figure 5: T-test comparisons of average final test accuracies across multiple experimental dimensions.
  • ...and 8 more figures

Theorems & Definitions (2)

  • Proposition 3.1
  • Proposition 3.2