GABFusion: Rethinking Feature Fusion for Low-Bit Quantization of Multi-Task Networks

TL;DR

This work tackles the degradation of quantization-aware training (QAT) in multi-task networks by identifying gradient conflicts at feature fusion and task-specific discrepancies as key bottlenecks. It introduces Gradient-Aware Balanced Feature Fusion (GABFusion), which uses a learnable fusion weight and a post-fusion LayerNorm to balance gradient flow between shallow and deep features, and Attention Distribution Alignment (ADA), a feature-level distillation mechanism based on SimAM and divergence-based alignment to guide quantized models. The approach is validated across YOLO-based detectors and multiple QAT baselines on PASCAL VOC and COCO, achieving average mAP improvements of about $3.3\%$ and $1.6\%$, respectively, and narrowing the YOLOv5 4-bit gap to the full-precision model on VOC to $1.7\%$. The method is modular and architecture-agnostic, offering a practical, plug-in enhancement to existing QAT techniques for robust performance under low-bit constraints, with potential extension to transformer-based multi-task systems.

Abstract

Despite the effectiveness of quantization-aware training (QAT) in compressing deep neural networks, its performance on multi-task architectures often degrades significantly due to task-specific feature discrepancies and gradient conflicts. To address these challenges, we propose Gradient-Aware Balanced Feature Fusion (GABFusion), which dynamically balances gradient magnitudes and fuses task-specific features in a quantization-friendly manner. We further introduce Attention Distribution Alignment (ADA), a feature-level distillation strategy tailored for quantized models. Our method demonstrates strong generalization across network architectures and QAT algorithms, with theoretical guarantees on gradient bias reduction. Extensive experiments demonstrate that our strategy consistently enhances a variety of QAT methods across different network architectures and bit-widths. On PASCAL VOC and COCO datasets, the proposed approach achieves average mAP improvements of approximately 3.3% and 1.6%, respectively. When applied to YOLOv5 under 4-bit quantization, our method narrows the accuracy gap with the full-precision model to only 1.7% on VOC, showcasing its effectiveness in preserving performance under low-bit constraints. Notably, the proposed framework is modular, easy to integrate, and compatible with any existing QAT technique-enhancing the performance of quantized models without requiring modifications to the original network architecture.

Figures (6)

• Figure 1: Comparison of the accuracy drop of the YOLOv5 model under 4‑bit quantization across three representative QAT schemes, alongside the proposed approach. To ensure reproducibility, we include our reimplementations of the baseline methods in the Appendix.
• Figure 2: Architecture of the YOLOv5 object detection model.
• Figure 3: Measured average gradient magnitude at the feature fusion node in YOLOv5 during quantization.
• Figure 4: The work flow of ADA. DNN represents the full-precision baseline model, while QNN denotes quantized model. Highlight regions in the heatmap represent attention areas of critical semantic information.
• Figure 5: (a)Visualization of the proposed strategy during the feature fusion process. A subset of 32 representative channels is displayed for clarity. (b) Average gradient magnitude at the feature fusion node in YOLOv5 after applying the proposed optimization strategy under quantization.