ElasticZO: A Memory-Efficient On-Device Learning with Combined Zeroth- and First-Order Optimization

• URL: http://arxiv.org/abs/2501.04287v1
• Date: Wed, 08 Jan 2025 05:25:14 GMT
• Title: ElasticZO: A Memory-Efficient On-Device Learning with Combined Zeroth- and First-Order Optimization
• Authors: Keisuke Sugiura, Hiroki Matsutani,
• Abstract summary: We propose ZO-based on-device learning (ODL) methods for full-precision and 8-bit quantized deep neural networks (DNNs)<n> ElasticZO achieves 5.2-9.5% higher accuracy with only 0.072-1.7% memory overhead, and can handle fine-tuning tasks as well as full training.<n> ElasticZO-INT8 achieves integer arithmetic-only ZO-based training for the first time, by incorporating a novel method for computing quantized ZO gradients from integer cross-entropy loss values.
• Score: 0.9444784653236158
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Zeroth-order (ZO) optimization is being recognized as a simple yet powerful alternative to standard backpropagation (BP)-based training. Notably, ZO optimization allows for training with only forward passes and (almost) the same memory as inference, making it well-suited for edge devices with limited computing and memory resources. In this paper, we propose ZO-based on-device learning (ODL) methods for full-precision and 8-bit quantized deep neural networks (DNNs), namely ElasticZO and ElasticZO-INT8. ElasticZO lies in the middle between pure ZO- and pure BP-based approaches, and is based on the idea to employ BP for the last few layers and ZO for the remaining layers. ElasticZO-INT8 achieves integer arithmetic-only ZO-based training for the first time, by incorporating a novel method for computing quantized ZO gradients from integer cross-entropy loss values. Experimental results on the classification datasets show that ElasticZO effectively addresses the slow convergence of vanilla ZO and shrinks the accuracy gap to BP-based training. Compared to vanilla ZO, ElasticZO achieves 5.2-9.5% higher accuracy with only 0.072-1.7% memory overhead, and can handle fine-tuning tasks as well as full training. ElasticZO-INT8 further reduces the memory usage and training time by 1.46-1.60x and 1.38-1.42x without compromising the accuracy. These results demonstrate a better tradeoff between accuracy and training cost compared to pure ZO- and BP-based approaches, and also highlight the potential of ZO optimization in on-device learning.

Related papers

• Tensor-Compressed and Fully-Quantized Training of Neural PDE Solvers [10.320585073024455]
  We present a framework that enables scalable and energy-efficient PINN training on edge devices.<n>This work enables real-time PDE solving on edge devices and paves the way for energy-efficient scientific computing at scale.
  arXiv  Detail & Related papers  (2025-12-10T00:00:34Z)
• Decorrelation Speeds Up Vision Transformers [0.6924349411126935]
  Masked Autoencoder (MAE) pre-training of vision transformers (ViTs) yields strong performance in low-label regimes but comes with substantial computational costs.<n>We address this by integrating Decorrelated Backpropagation (DBP) into MAE pre-training, an optimization method that iteratively reduces input correlations at each layer to accelerate convergence.<n>On ImageNet-1K pre-training with ADE20K fine-tuning, DBP-MAE reduces wall-clock time to baseline performance by 21.1%, lowers carbon emissions by 21.4% and improves segmentation mIoU by 1.1
  arXiv  Detail & Related papers  (2025-10-16T13:13:12Z)
• Tri-Accel: Curvature-Aware Precision-Adaptive and Memory-Elastic Optimization for Efficient GPU Usage [0.6511750267058007]
  Tri-Accel is a unified optimization framework that co-adapts three acceleration strategies along with adaptive parameters during training.<n>On CIFAR-10 with ResNet-18 and EfficientNet-B0, Tri-Accel achieves up to 9.9% reduction in training time and 13.3% lower memory usage.<n>Compared to static mixed-precision training, Tri-Accel maintains 78.1% accuracy while reducing memory footprint from 0.35GB to 0.31GB on standard hardware.
  arXiv  Detail & Related papers  (2025-08-23T05:38:42Z)
• The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws [51.608402959163925]
  We present the first systematic exploration of optimal sparse pre-training configurations for large language models. We find that initiating pruning at 25% of total training compute and concluding at 75% achieves near-optimal final evaluation loss. We propose a new scaling law that modifies the Chinchilla scaling law to use the average parameter count over pre-training.
  arXiv  Detail & Related papers  (2025-01-21T20:23:22Z)
• YOSO: You-Only-Sample-Once via Compressed Sensing for Graph Neural Network Training [9.02251811867533]
  YOSO (You-Only-Sample-Once) is an algorithm designed to achieve efficient training while preserving prediction accuracy. YOSO not only avoids costly computations in traditional compressed sensing (CS) methods, such as orthonormal basis calculations, but also ensures high-probability accuracy retention.
  arXiv  Detail & Related papers  (2024-11-08T16:47:51Z)
• AutoScale: Automatic Prediction of Compute-optimal Data Composition for Training LLMs [61.13296177652599]
  This paper demonstrates that the optimal composition of training data from different domains is scale-dependent.<n>We introduce *AutoScale*, a novel, practical approach for optimizing data compositions at potentially large training data scales.<n>Our evaluation on GPT-2 Large and BERT pre-training demonstrates *AutoScale*'s effectiveness in improving training convergence and downstream performance.
  arXiv  Detail & Related papers  (2024-07-29T17:06:30Z)
• Towards Continual Learning Desiderata via HSIC-Bottleneck Orthogonalization and Equiangular Embedding [55.107555305760954]
  We propose a conceptually simple yet effective method that attributes forgetting to layer-wise parameter overwriting and the resulting decision boundary distortion. Our method achieves competitive accuracy performance, even with absolute superiority of zero exemplar buffer and 1.02x the base model.
  arXiv  Detail & Related papers  (2024-01-17T09:01:29Z)
• AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
  We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models. AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
  arXiv  Detail & Related papers  (2023-10-16T09:04:28Z)
• Efficient and Flexible Neural Network Training through Layer-wise Feedback Propagation [49.44309457870649]
  We present Layer-wise Feedback Propagation (LFP), a novel training principle for neural network-like predictors. LFP decomposes a reward to individual neurons based on their respective contributions to solving a given task. Our method then implements a greedy approach reinforcing helpful parts of the network and weakening harmful ones.
  arXiv  Detail & Related papers  (2023-08-23T10:48:28Z)
• Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed) Neural Networks [15.188785164091987]
  Backward propagation (BP) is widely used to compute the gradients in neural network training. It is hard to implement BP on edge devices due to the lack of hardware and software resources to support automatic differentiation. This paper presents a completely BP-free framework that only requires forward propagation to train realistic neural networks.
  arXiv  Detail & Related papers  (2023-08-18T23:56:50Z)
• Robust Learning with Progressive Data Expansion Against Spurious Correlation [65.83104529677234]
  We study the learning process of a two-layer nonlinear convolutional neural network in the presence of spurious features. Our analysis suggests that imbalanced data groups and easily learnable spurious features can lead to the dominance of spurious features during the learning process. We propose a new training algorithm called PDE that efficiently enhances the model's robustness for a better worst-group performance.
  arXiv  Detail & Related papers  (2023-06-08T05:44:06Z)
• Layer-wise Adaptive Step-Sizes for Stochastic First-Order Methods for Deep Learning [8.173034693197351]
  We propose a new per-layer adaptive step-size procedure for first-order optimization methods in deep learning. The proposed approach exploits the layer-wise curvature information contained in the diagonal blocks of the Hessian in deep neural networks (DNNs) to compute adaptive step-sizes (i.e., LRs) for each layer. Numerical experiments show that SGD with momentum and AdamW combined with the proposed per-layer step-sizes are able to choose effective LR schedules.
  arXiv  Detail & Related papers  (2023-05-23T04:12:55Z)
• 8-bit Optimizers via Block-wise Quantization [57.25800395197516]
  Statefuls maintain statistics over time, e.g., the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past values. This state can be used to accelerate optimization compared to plain gradient descent but uses memory that might otherwise be allocated to model parameters. In this paper, we develop first gradients that use 8-bit statistics while maintaining the performance levels of using 32-bit gradient states.
  arXiv  Detail & Related papers  (2021-10-06T15:43:20Z)
• Efficient training of physics-informed neural networks via importance sampling [2.9005223064604078]
  Physics-In Neural Networks (PINNs) are a class of deep neural networks that are trained to compute systems governed by partial differential equations (PDEs) We show that an importance sampling approach will improve the convergence behavior of PINNs training.
  arXiv  Detail & Related papers  (2021-04-26T02:45:10Z)
• FracTrain: Fractionally Squeezing Bit Savings Both Temporally and Spatially for Efficient DNN Training [62.932299614630985]
  We propose FracTrain that integrates progressive fractional quantization which gradually increases the precision of activations, weights, and gradients.<n>FracTrain reduces computational cost and hardware-quantified energy/latency of DNN training while achieving a comparable or better (-0.12%+1.87%) accuracy.
  arXiv  Detail & Related papers  (2020-12-24T05:24:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.