TLDR: GABFusion and Attention Distribution Alignment (ADA) are proposed to improve low-bit quantization in multi-task neural networks like YOLO. GABFusion dynamically balances gradient magnitudes and fuses task-specific features in a quantization-friendly way, while ADA uses feature-level distillation to align attention distributions between full-precision and quantized models. This approach significantly reduces accuracy degradation, achieving notable mAP improvements and narrowing the performance gap with full-precision models, especially under challenging low-bit constraints.
In the rapidly evolving world of artificial intelligence, the ability to deploy powerful deep neural networks on resource-constrained devices is becoming increasingly crucial. One key technique for achieving this is quantization, which compresses models by converting their high-precision floating-point weights and activations into lower-bit integer equivalents. While highly effective for single-task networks, low-bit quantization (especially at 4 bits or less) often leads to significant performance drops in multi-task architectures, such as object detectors like the YOLO family.
Researchers Zhaoyang Wang and Dong Wang from Beijing Jiaotong University have identified the core challenges behind this performance degradation: task-specific feature discrepancies and gradient conflicts during the quantization-aware training (QAT) process. To address these issues, they propose a novel framework called Gradient-Aware Balanced Feature Fusion (GABFusion) combined with Attention Distribution Alignment (ADA).
The Quantization Conundrum in Multi-Task Networks
Traditional quantization methods, whether Post-Training Quantization (PTQ) or Quantization-Aware Training (QAT), have largely focused on image classification. However, multi-task networks, which simultaneously handle tasks like object classification and bounding box regression, present unique challenges. For instance, in YOLO models, shallow layers are vital for fine-grained spatial details needed for regression, while deeper layers capture high-level semantics for classification. During quantization, the noise introduced at each layer accumulates, leading to severe gradient oscillations and, more critically, an imbalance in gradient updates at feature fusion points. This means the network might prioritize optimizing deep semantic features while neglecting the shallow, local features essential for accurate object localization.
GABFusion: Balancing the Feature Act
GABFusion is designed to dynamically balance gradient magnitudes and fuse task-specific features in a quantization-friendly manner. It introduces two main components:
- Learnable Dynamic Scaling Factors: Before features from different branches (e.g., shallow and deep) are fused, GABFusion applies learnable weight parameters. These parameters adaptively control the contribution ratio of each feature branch, ensuring a more balanced input to the fusion layer.
- Gradient-Balancing Normalization: Immediately after feature fusion, a Layer Normalization module is inserted. This acts as a plug-in component that re-centers and re-scales gradients, effectively amplifying gradients from the shallow branch and suppressing those from the deep branch. This process corrects the optimization bias, encouraging both branches to contribute equally to the learning process. A notable advantage is that this Layer Normalization can be removed during inference, adding no computational overhead to the deployed model.
ADA: Aligning Attention for Better Learning
Beyond numerical accuracy, the paper highlights that a model’s decision-making capacity is driven by its feature representations. Conventional loss functions, which focus on numerical deviations, often fall short in guiding quantized multi-task networks. To overcome this, the researchers propose Attention Distribution Alignment (ADA), a feature-level distillation strategy.
ADA leverages a parameter-free attention mechanism, SimAM, to extract attention distributions from both a full-precision ‘teacher’ model and the quantized ‘student’ model. These distributions, which highlight important feature regions, are then aligned using a divergence metric like Jensen-Shannon (JS) divergence. By minimizing the discrepancy between these attention patterns, ADA guides the quantized model to recover crucial semantic information and improve its robustness and discriminability.
Also Read:
- Streamlining Person Re-Identification for Edge Devices with One-Shot Knowledge Transfer
- Tiny Bits, Big Impact: Quantized AI Controllers for Real-Time Robotics
Impressive Results and Broad Applicability
The effectiveness of GABFusion and ADA was rigorously tested on standard benchmarks like PASCAL VOC and COCO datasets, using YOLOv5s and YOLOv11s as baseline models. The proposed strategy consistently enhanced a variety of state-of-the-art QAT methods (N2UQ, PACT, LSQ, LSQ+) across different network architectures and challenging low-bit settings (4-bit and 3-bit quantization).
On PASCAL VOC, the method achieved average mAP improvements of approximately 3.3% and 1.6% on COCO. Notably, when applied to YOLOv5 under 4-bit quantization, GABFusion and ADA narrowed the accuracy gap with the full-precision model to only 1.7% on VOC. This demonstrates its significant capability in preserving performance under severe low-bit constraints. The approach also showed strong compatibility, further improving performance when integrated with other techniques like EMA.
The research provides a deep analysis into the root causes of performance degradation in low-bit quantization for multi-task networks, identifying misaligned feature fusion and biased gradient optimization as critical bottlenecks. The modular nature of GABFusion and ADA means they can be easily integrated with any existing QAT technique without requiring modifications to the original network architecture, making them highly practical for real-world applications. For more details, you can read the full research paper here.
