Latency-Constrained DNN Architecture Learning for Edge Systems using Zerorized Batch Normalization

Abstract: Deep learning applications have been widely adopted on edge devices, to mitigate the privacy and latency issues of accessing cloud servers. Deciding the number of neurons during the design of a deep neural network to maximize performance is not intuitive. Particularly, many application scenarios are real-time and have a strict latency constraint, while conventional neural network optimization methods do not directly change the temporal cost of model inference for latency-critical edge systems. In this work, we propose a latency-oriented neural network learning method to optimize models for high accuracy while fulfilling the latency constraint. For efficiency, we also introduce a universal hardware-customized latency predictor to optimize this procedure to learn a model that satisfies the latency constraint by only a one-shot training process. The experiment results reveal that, compared to state-of-the-art methods, our approach can well-fit the 'hard' latency constraint and achieve high accuracy. Under the same training settings as the original model and satisfying a 34 ms latency constraint on the ImageNet-100 dataset, we reduce GoogLeNet's latency from 40.32 ms to 34 ms with a 0.14% accuracy reduction on the NVIDIA Jetson Nano. When coupled with quantization, our method can be further improved to only 0.04% drop for GoogLeNet. On the NVIDIA Jetson TX2, we compress VGG-19 from 119.98 ms to 34 ms and even improve its accuracy by 0.5%, and we scale GoogLeNet up from 20.27 ms to 34 ms and achieve higher accuracy by 0.78%. We also open source this framework at https://github.com/ntuliuteam/ZeroBN

• The paper introduces a hardware-aware DNN learning framework with zeroized batch normalization to directly meet latency constraints in edge systems.
• It leverages a learning-based latency predictor to substitute slow on-device measurements and guide dynamic channel pruning and model scaling.
• Empirical results demonstrate minimized accuracy loss—with occasional gains—even under aggressive pruning and real-time latency targets.

Latency-Constrained Neural Network Optimization for Edge Systems via Zeroized Batch Normalization

Introduction

The increasing deployment of Deep Neural Networks (DNNs) at the edge has induced acute demand for models that are both accurate and capable of real-time inference within stringent latency constraints. Existing neural network compression and architecture optimization strategies—such as pruning, quantization, and model scaling—generally operate on indirect performance metrics (FLOPs, parameter counts), which are notoriously ill-correlated with wall-clock latency on heterogeneous edge devices. Moreover, these methods are coupled with significant training overhead due to multi-stage pipelines comprising pre-training, iterative modifications/evaluations on real hardware, and post-pruning fine-tuning steps. This work, "Latency-Constrained DNN Architecture Learning for Edge Systems using Zerorized Batch Normalization" (2607.06922), introduces a hardware-aware DNN learning framework with explicit latency optimization as the principal objective, leveraging a one-shot training paradigm and a universal hardware-customized latency predictor.

Framework Overview

The proposed pipeline encompasses: (1) compact learning for model reduction, (2) a learning-based latency prediction model to substitute for real-device measurements during training, and (3) a scaling approach to expand underutilized models where feasible. The workflow is illustrated in (Figure 1). The latency predictor enables the host to avoid device-executed latency profiling, drastically increasing optimization efficiency.

Figure 1: The workflow of hardware-aware neural network optimization. With a latency predictor, the host can avoid time-consuming on-device measurements during the process, extremely improving efficiency.

The end-to-end system (Figure 2) ingests a pre-existing model (mismatched in latency to target device) and emits a compressed or expanded model that maximizes accuracy while strictly adhering to a specified latency constraint.

Figure 2: The overview of the model learning framework. The input is a mismatched original model and this framework outputs a model with high accuracy under the latency constraint.

Compact Learning via Zeroized Batch Normalization

Formulation

Unlike standard three-stage pruning schedules, compact learning operates directly on model architectures via dynamic channel and layer elimination during training, obviating the need for pre-trained weights or iterative retraining. The core mechanism exploits the scaling parameter $\gamma$ of Batch Normalization (BN) layers as a proxy for channel importance, applying sparsity-inducing regularization and dynamically zeroizing $\gamma$/$\beta$ after a fixed number of epochs. Channels with low-magnitude $\gamma$ are zeroed, effecting channel-wise pruning. This is accompanied by a "Zero-Recovery" cycle where pruned channels may be restored if their gradient dynamics dictate relevance, leveraging nonzero momentum in contemporary optimizers to escape local minima.

Figure 3: The process of the compact learning approach. The input is a redundant model and the output is a compact model that meets the system latency constraint. The gray dotted boxes represent the process of quantization.

To support structured layer-level pruning, auxiliary branch layers are pre-inserted such that the importance of each layer may be evaluated globally and removed while maintaining connectivity, as in (Figure 4). The pruning is thus sensitive to both width and depth, with empirical findings confirming that latency reductions from width and depth pruning are model-dependent.

Figure 4: The preprocessing for layer pruning by adding branch layers.

A distinguishing attribute is the reversible pruning during training (Figure 5): zeroized channels can be reinstated as the optimizer's momentum induces gradient updates away from zero, circumventing irrecoverable pruning seen in prior art such as Network Slimming (NS) and Soft Filter Pruning (SFP).

Figure 5: Channel importance changes during the learning process. Zeroized channels (white blocks) can be easily recovered.

Complexity and Efficiency

The approach inserts negligible computational overhead over standard training. The main additional cost is $O(n\log n)$ sorting per epoch—subordinate compared to $O(\sum_l M_l^2 K_l^2 C_{l-1} C_l)$ for convolutions. No additional pre-training or retraining is needed.

Machine Learning-Based Hardware Latency Prediction

A critical innovation is the replacement of device latency measurements (which are slow and inefficient) with a trained regression model—specifically, a shallow backpropagation neural network—mapping layer configurations to latency (Figure 6).

Figure 6: The framework of the latency predictor: model construction (left), prediction usage (right).

Training data for the latency predictor is generated via both regular and random sampling over architecture configurations, and the latency responses are post-processed for a normal-distribution-like label distribution (Figure 7) to facilitate stable regression.

Figure 7: The latency distribution before/after transform. The normal distribution can improve the accuracy of the latency predictor.

This predictor generalizes to unseen architectures within the sampled parameter space, achieving an average prediction error of 6.12%, clearly outperforming previous predictors (e.g., nn-Meter, Justus et al.). It provides real-time latency guidance within the compact learning loop.

Latency-Critical Model Scaling

When the initial model is over-compact relative to the device's latency budget, performance is left untapped. The framework uniformly scales depth and width (both channels and layers) by an automatically chosen expansion factor, followed by compact learning to reel the bloated model back into the latency constraint (Figure 8). This eliminates expensive microscopic searches for optimal scaling factors typically employed in works such as EfficientNet.

Figure 8: The process of the latency-critical model scaling approach.

Quantization Integration

Post-training quantization (FP16 supported on tested hardware) further reduces inference latency, allowing retention of more parameters under the same latency constraint and thus higher accuracy. The latency predictor is retrained to capture quantized inference characteristics.

Empirical Results

Compression and Accuracy

Experiments on CIFAR-10 and ImageNet-100 across VGG-19, ResNet-164/50, DenseNet-40, and GoogLeNet show that the proposed method consistently outperforms SFP, FPGM, PGMPF, NS, and OTO both at aggressive 50%/70% pruning and under latency constraints (Table 1, Figure 9). With the same or less training budget, accuracy drop is minimized, and in some settings, accuracy increases upon pruning—contradicting the prevalent accuracy-efficiency tradeoff. The method remains robust across models with varying architectural redundancy.

Figure 9: The accuracy w.r.t the variation of $k$ and $s$ values.

Architecture Adaptation

Notably, the output compact architectures differ contextually: on large-scale datasets (ImageNet-100), the framework automatically retains greater depth (Figure 10); for over-parameterized cases, compact learning alone suffices, while scaling is exploited where models are initially too constrained.

Figure 10: Model architecture from layer pruning of VGG-19. Model a is for CIFAR-10, Model b is for ImageNet-100.

Latency-Constrained Optimization

The framework enables both model compression and expansion tuned to device-specific latency constraints (e.g., 34ms, corresponding to typical real-time requirements). Applied to NVIDIA Jetson TX2 and Nano, the approach can compress VGG-19 from 119.98ms to 34ms with a 0.5% accuracy gain, and scale up GoogLeNet from 20.27ms to 34ms achieving a 0.78% accuracy gain. Quantized models (FP16) further reduce the required compression ratio and increase accuracy, as predicted by the latency predictor. Overhead is marginal: e.g., compact learning adds only 1.1% to baseline training cost.

Ablations and Hyperparameter Sensitivity

The method introduces two hyperparameters—start epoch $s$ and interval $k$ for zero/recovery cycles—which are empirically shown to be robust: accuracy varies little over a wide range and defaults are automated (Figure 9).

Implications and Future Directions

This framework exemplifies direct hardware-aware DNN optimization that is extensible beyond latency: the same predictor architecture can guide constraints for energy or memory, forming a foundation for resource-constrained AutoML and NAS systems. It is not constrained to human-designed models; NAS architectures and black-box hardware can be integrated. Additionally, the technique opens directions for joint application with knowledge distillation and automated model search, as well as composable multi-objective optimization.

Conclusion

This work presents a hardware-tailored, latency-critical DNN learning framework that integrates compact learning driven by zeroized batch normalization, efficient latency prediction, and universal model scaling. The resultant models faithfully satisfy latency constraints while often improving accuracy and reducing engineering cost in edge deployment, substantially advancing practical neural network adaptation for heterogeneous, real-time systems.

---

Reference:

"Latency-Constrained DNN Architecture Learning for Edge Systems using Zerorized Batch Normalization" (2607.06922)