PaperlayerView exact revision on arXiv

Source-linked AI summary

TASTE: Throughput-Aware Batch Size Tuning for On-Device Edge Learning

Avik Bhatnagar, Federico Nicolas Peccia, Oliver Bringmann

arXiv:2609.07444v1cs.LGcs.AIcs.ARcs.PF

TL;DR

Resource-constrained edge CPUs make training-throughput optimization challenging because training stores intermediate activations and embedded devices lack high-bandwidth memory. TASTE uses Bayesian optimization to select batch sizes and adjusts training regimes for supervised and continual learning. The evaluated devices showed at least 1.5× higher training throughput, with optimized-batch continual-learning accuracy close to baseline.

  • Problem

    Edge training must optimize throughput within memory constraints, and existing throughput-optimization work largely targets server hardware and inference rather than edge-CPU training.

    Source: §I · p. 2

  • Method

    TASTE applies Bayesian optimization to choose a throughput-maximizing batch size and adapts training regimes for supervised and continual learning.

    Source: §I · p. 2, B. Batch Size Tuning for Training… · p. 3

  • Results

    Training throughput improved at least 1.5× across evaluated edge devices, while optimized-batch continual-learning accuracies stayed close to baseline and exceeded maximum-batch configurations.

    Source: A. Hardware Throughput… · p. 6, B. Impact of Optimized Batch Sizes on… · p. 7

  • Takeaways & Limitations

    In image classification, throughput-optimized batch sizes with tailored training regimes achieved learning performance comparable to reference implementations.

    Source: §I · p. 2

Abstract

from arXiv · show

The rise of privacy-preserving artificial intelligence (AI) has shifted the focus of model adaptation and personalization towards on-device learning, where deep learning models are finetuned directly on edge hardware using local user data. However, this shift requires optimization of deep learning training on resource-constrained hardware to maximize throughput while maintaining predictive accuracy. This paper introduces a novel technique for on-device model training that incorporates an efficient Bayesian optimization-based batch size tuning approach to maximize hardware throughput. To evaluate the impact of this hyperparameter on the learning dynamics, we investigated two distinct paradigms: standard supervised learning (SL) and online continual learning (CL). Experimental results across various edge devices demonstrate a throughput ceiling, beyond which increasing the batch size yields no additional throughput gains. The proposed tuning approach identifies the optimal batch size, which, when combined with gradient accumulation and linear learning rate scaling, achieves up to a 2X increase in training throughput on platforms such as Raspberry Pi 4 compared to maximum batch sizes, without compromising model accuracy. Furthermore, in the CL paradigm, we demonstrate that optimal batch sizes maintain the stability-plasticity balance required for incremental learning, effectively mitigating catastrophic forgetting while maximizing computational efficiency on edge-hardware.

I. INTRODUCTION

On-device learning can adapt models locally to shifting user contexts, but resource-constrained edge CPUs pose distinct training-throughput and memory challenges. TASTE combines Bayesian batch-size tuning with training-regime adjustments, reporting up to 2X higher throughput while maintaining comparable learning performance.

I. INTRODUCTION

  • On-device learning supports adaptation to distribution shifts while keeping sensitive data local and enabling personalization for user contexts.Source: §I · p. 1
  • Prior throughput-optimization work focuses mainly on server hardware and inference, while graph, data-movement, and optimization costs constrain its suitability for edge CPUs.Training also requires storing intermediate activations for backpropagation, so static pre-compiled approaches that respect memory and energy constraints may be more viable.Source: §I · p. 2
  • Embedded devices lack the large datasets, high-bandwidth memory, and rapid data transfers available in cloud or server training environments.Source: §I · p. 2
  • TASTE uses Bayesian optimization to find a throughput-maximizing batch size that need not be the hardware maximum, and adjusts training regimes to limit learning-performance degradation.Source: §I · p. 2
  • Up to 2X higher training throughput than using the maximum possible batch size was reported alongside learning performance comparable to reference implementations.Source: §I · p. 2

A. On-device Model Training

The on-device training implementation extends an inference model with backward-pass operations and deploys the resulting training graph using TVM optimizations.

A. On-device Model Training

  • ONNX training extends the inference forward graph with loss and gradient operators, enabling iterative parameter updates from local input data.Source: A. On-device Model Training · p. 2
  • For convolution layers, backward-pass gradients are expressed using standard convolution, transposed convolution, and reduced-sum operations.The forward output is computed from the input, weight kernel, and bias; the upstream loss gradient is used in backward computation.Source: A. On-device Model Training · pp. 2–3
  • TVM applies operator fusion, constant folding, and layout transformations to create a static forward-and-backward graph deployable on supported ARM and RISC-V backends.Its RPC infrastructure sends compiled binaries from a host to target edge platforms for execution and evaluation.Source: A. On-device Model Training · p. 3

B. Batch Size Tuning for Training Throughput Optimization

The tuning procedure searches device-feasible batch sizes using on-device throughput measurements. Bayesian optimization reduces search iterations while seeking a near-optimal configuration.

B. Batch Size Tuning for Training Throughput Optimization

  • The search space runs from batch size 1 through the memory-limited maximum, with even-numbered increments after 1.Exhaustive grid search can be costly, random search may not converge in limited iterations, and powers-of-two-only profiling may miss non-power-of-two peaks.Source: B. Batch Size Tuning for Training… · p. 3
  • Figure 1 depicts iterative on-device measurement and batch-size selection to identify the optimal configuration, bsopt.Source: Figure 1 · p. 3
  • Bayesian optimization with a Tree-structured Parzen Estimator selects candidate batch sizes from throughput measurements, using half the search-space size in iterations.The approach reports a twofold reduction in iterations versus exhaustive grid search while identifying near-optimal throughput configurations.Source: B. Batch Size Tuning for Training… · p. 3
  • Figure 2 plots training throughput in imgs/sec against batch size, marking bsopt for each platform and dataset and the mean improvement relative to bsmax.Source: Figure 2 · p. 4

C. Training Performance Evaluation

The evaluation considers static-data supervised learning and incremental continual learning, whose changing data stream creates a risk of forgetting earlier tasks. It adjusts training regimes alongside batch size to retain learning performance.

C. Training Performance Evaluation

  • Continual learning adapts to a non-stationary stream without prior access to the full dataset, and performance on previously learned tasks can degrade as new distributions arrive.Source: C. Training Performance Evaluation · p. 4
  • Task-incremental learning uses task identifiers and task-specific heads, whereas class-incremental learning must add classes without task IDs and preserve earlier-class accuracy.Source: C. Training Performance Evaluation · p. 4
  • Experience Replay uses a small episodic memory and is described as an effective, lightweight continual-learning baseline for resource-constrained on-device settings.Source: C. Training Performance Evaluation · p. 4
  • Learning rate, batch size, and dropout influence continual-learning accuracy and forgetting mitigation, while batch-size-to-learning-rate ratio affects supervised-learning generalization.Source: C. Training Performance Evaluation · p. 4
  • Figure 3 plots inference throughput in imgs/sec against batch size, marking bsopt for each platform and dataset and mean improvement relative to bsmax.Source: Figure 3 · p. 5
  • Continual learning pairs the obtained batch size with a linearly scaled learning rate, while supervised learning uses gradient accumulation to emulate larger batches within device memory limits.Accumulated mini-batch gradients trigger one weight update after all constituent mini-batches are processed.Source: C. Training Performance Evaluation · pp. 4–5

III. EXPERIMENTAL SETUP

The experiments evaluate ResNet-18 on three image-classification datasets and edge CPUs with varying computational and memory profiles. Supervised and online continual-learning setups compare tuned batch sizes with larger or reference configurations.

III. EXPERIMENTAL SETUP

  • Experiments used ResNet-18 on CIFAR-10, CIFAR-100, and TinyImageNet across edge CPUs with varying compute and memory profiles.The evaluation was CPU-only, including on the NVIDIA Jetson Xavier; the devices span ARM and RISC-V ISAs.Source: §III · p. 5, Table I · p. 5
  • For supervised learning, the tuned micro-batch was accumulated to an effective global batch size of 128; models trained for 200 epochs at learning rate 0.1.The backbone was ImageNet-pretrained ResNet-18 with its classifier replaced for each dataset.Source: §III · p. 5
  • Online continual learning used one pass over the training stream, with classes presented with task identifiers for TIL and without them for CIL.Experience Replay used a base mini-batch of 20, a replay buffer, and learning-rate scaling across batch sizes.Source: §III · pp. 5–6
  • The supervised accuracy comparison evaluates optimal batch sizes against maximum sizes 128 and 512 while using gradient accumulation.The figure covers CIFAR-10, CIFAR-100, and TinyImageNet.Source: Figure 4 · p. 6

A. Hardware Throughput Characterization and bsopt Identification

Training throughput peaks at an optimal batch size and then declines, unlike inference, where larger batches generally improve throughput. Across edge devices, tuning to the training peak improves throughput over maximum batch sizes.

A. Hardware Throughput Characterization and bsopt Identification

  • Across edge devices, training throughput peaks and then declines as batch size increases; the peak defines bsopt.The peak marks a hardware throughput ceiling for training.Source: A. Hardware Throughput… · p. 6
  • Training throughput increased by at least 1.5× across all devices with bsopt versus bsmax, peaking at approximately 2× on Raspberry Pi.Fig. 5 shows training speedup across CIFAR and TinyImageNet datasets and evaluated edge devices.Source: Figure 5 · p. 6, A. Hardware Throughput… · pp. 6–7
  • The training bottleneck is batch-scaled storage of ResNet-18 intermediate activations needed for backpropagation.Training throughput accounts for the computation and latency of loss and gradient calculations.Source: A. Hardware Throughput… · p. 6
  • Inference does not retain intermediate activations, and larger batches generally yield higher throughput, making maximum supported batches preferable.The tuned inference batch provides only a negligible average throughput increase over the maximum batch.Source: A. Hardware Throughput… · p. 6
  • Fig. 6 compares accuracy over cumulative classes in CIL or tasks in TIL across datasets for bsopt, batch 20, and bsmax.The plotted configurations use linear learning-rate scaling.Source: Figure 6 · p. 7

B. Impact of Optimized Batch Sizes on Training Performance

In online continual learning, tuned batch sizes achieve final-step accuracies close to the reference and outperform maximum batch sizes. In supervised learning, accuracy outcomes vary by dataset and batch size.

B. Impact of Optimized Batch Sizes on Training Performance

  • In supervised learning, tuned micro-batches with gradient accumulation yielded testing accuracies comparable to larger batch sizes.The continual-learning comparisons use incremental class or task presentations and a reference batch size of 20.Source: B. Impact of Optimized Batch Sizes on… · p. 7
  • For CIFAR-10, final test accuracy peaked at 86.76% ± 0.11% after 200 epochs, but batch size 4 yielded 69.61% ± 13.82% under unstable training.This is the reported outcome for the 10-class dataset.Source: B. Impact of Optimized Batch Sizes on… · p. 7
  • On CIFAR-100, batch size 10 achieved 62.97% ± 0.30% accuracy, exceeding batch sizes 128 at 59.91% ± 0.46% and 512 at 59.08%±0.31%.The paper suggests small-minibatch noise in Batch Normalization statistics may regularize this 100-class task.Source: B. Impact of Optimized Batch Sizes on… · p. 7
  • For TinyImageNet, accuracy peaked at 57.23% ± 0.25% with batch size 512, which the passage attributes to better representation of 200 classes in each update.Source: B. Impact of Optimized Batch Sizes on… · p. 7
  • In CIL, final-step accuracies for bsopt were 36.12%±2.23% vs. 34.48%±1.21% on CIFAR-10, 11.90%±0.97% vs. 10.98%±0.60% on CIFAR-100, and 11.57% ± 0.30% vs. 10.01% ± 0.38% on TinyImageNet, versus baseline.A similar trend holds for TIL, and bsopt configurations outperform bsmax; on CIFAR-10, bsmax degradation reached 11.67% ± 3.14% and 21.01% ± 3.67% in CIL, and 15.02% ± 11.72% and 28.54% ± 3.81% in TIL, respectively.Source: B. Impact of Optimized Batch Sizes on… · pp. 7–8
  • In continual learning, bsopt supports the stability-plasticity trade-off; larger batches mean fewer parameter updates and more reuse of fixed replay samples, potentially overfitting older classes.The proposed balance maintains plasticity for new information while providing stability to mitigate catastrophic forgetting.Source: B. Impact of Optimized Batch Sizes on… · p. 8

V. CONCLUSION & FUTURE WORK

TASTE uses Bayesian optimization to identify throughput-optimal batch sizes and reports gains without sacrificing accuracy. The study evaluates ResNet-18 and identifies broader validation and dynamic tuning as future work.

V. CONCLUSION & FUTURE WORK

  • TASTE achieved at least 1.5× training-throughput speedup across evaluated edge devices and approximately 2× on Raspberry Pi versus maximum batch size.The paper attributes the throughput ceiling to intermediate-activation memory bottlenecks during backpropagation.Source: §V · p. 8
  • Across supervised and online continual learning, tuned batch sizes achieved accuracy comparable to reference implementations with gradient accumulation and linear learning-rate scaling.The reported evaluation concerns image classification.Source: §V · p. 8
  • Evaluation used CNN-based ResNet-18; future work includes other architectures, application domains, heterogeneous accelerator platforms, dynamic batch selection, and adaptive optimizers.The paper notes that advanced optimizers could add computational overhead on resource-constrained devices.Source: §V · p. 8
Loading 2609.07444v1…