Gradient-Aware Balanced Feature Fusion (GABFusion)

• GABFusion is a modular technique that systematically mitigates gradient bias and quantization errors in multi-task deep neural architectures.
• It employs branch scaling and layer normalization to harmonize shallow and deep features, ensuring balanced gradient flow and improved convergence.
• The integrated ADA distillation module aligns attention distributions, significantly enhancing quantized model performance across various QAT schemes.

Gradient-Aware Balanced Feature Fusion (GABFusion) is a modular architectural technique designed to address the degradation in quantization-aware training (QAT) of multi-task deep neural networks, particularly under low-bit constraints. GABFusion systematically balances gradient contributions and feature quality at fusion points where disparate branches (such as shallow and deep features) are merged, thereby mitigating cumulative quantization errors and gradient starvation. The method incorporates a plug-and-play fusion module and an auxiliary feature-level distillation strategy—Attention Distribution Alignment (ADA)—to robustly improve quantized model performance across various QAT methods and architectures.

1. Quantization Error Propagation and Imbalance in Multi-Task Fusion

Quantization-aware training introduces layer-wise errors that propagate and compound through a network. For a layer $l$ with full-precision output $x_l$ and quantized output $\tilde{x}_l$, the quantization error $\delta_l = \tilde{x}_l - x_l$ evolves as: $\delta_l \approx J_l\,\delta_{l-1} + \epsilon_l$ where $J_l = \partial f_l/\partial x_{l-1}$ and $\epsilon_l$ represents the new quantization noise at layer $l$. This recursive relation shows that errors introduced by quantization accumulate across layers.

In architectures with multi-branch feature fusion (e.g., spatial and semantic pathways in detection heads), quantization noise affects deep branches more severely, leading to imbalanced gradients: $$\Bigl|\frac{\partial L}{\partial \tilde F_d}\Bigr| \gg \Bigl|\frac{\partial L}{\partial \tilde F_s}\Bigr|$$ where $\tilde F_d$ and $\tilde F_s$ are quantized deep and shallow features, respectively. This gradient bias causes the network to disproportionately update deep branch parameters, neglecting shallow representations critical for capturing fine details.

2. GABFusion Mechanism and Mathematical Formulation

GABFusion equalizes feature and gradient contributions at each fusion point through two primary steps:

1. Branch Scaling: Apply a learnable scalar $\alpha \in [0,1]$ to scale the shallow branch and $1-\alpha$ to scale the deep branch:

$$F'_s = \alpha\,\tilde F_s \qquad F'_d = (1-\alpha)\,\tilde F_d$$

2. Layer Normalization (LN): Concatenate the scaled features, then apply parameter-free LN across channels per spatial location. For $h = F'_{\text{cat}}(i,j) \in \mathbb{R}^{C_s+C_d}$,

$$\mathrm{LN}(h_k) = \frac{h_k - \mu}{\sigma}, \quad\text{where}\quad \mu = \frac{1}{C}\sum_k h_k,\ \sigma = \sqrt{\frac{1}{C}\sum_k (h_k-\mu)^2}$$

Gradient propagation through LN ensures: $$\frac{\partial L}{\partial h_k} = \frac{1}{\sigma} \left( \frac{\partial L}{\partial \mathrm{LN}(h_k)} - \frac{1}{C}\sum_j \frac{\partial L}{\partial \mathrm{LN}(h_j)} \right)$$ This normalization amplifies updates to underrepresented branches (typically shallow) and suppresses overrepresented ones (deep), thereby systematically balancing the training signal.

3. Architectural and Implementation Aspects

The GABFusion module is positioned immediately prior to each feature fusion (typically "Concat") node. Its specific workflow is:

• Inputs: $$\tilde F_s \in \mathbb{R}^{B\times C_s\times H\times W}$$, $$\tilde F_d \in \mathbb{R}^{B\times C_d\times H\times W}$$
• Branch scaling: Compute $F'_s$, $F'_d$
• Concatenation: $F_{\text{fused}} = [F'_s; F'_d]$
• Layer Normalization: Channel-wise, per spatial location
• Output: Shape $B\times(C_s+C_d)\times H\times W$

At inference, LayerNorm can be subsumed into neighboring layers or omitted, and $\alpha$ can be absorbed into subsequent convolutional weights. Thus, GABFusion is effectively "zero-cost" at deployment, adding no parameters or compute relative to the base fusion point.

4. Theoretical Guarantees and Optimization Properties

GABFusion’s construction provides explicit guarantees on gradient bias reduction at fusion points: $$\sum_k\frac{\partial L}{\partial h_k}=0,\quad \mathrm{Var}\left(\frac{\partial L}{\partial h_k}\right)\leq \mathrm{Var}\left(\frac{\partial L}{\partial \mathrm{LN}(h_k)}\right)$$ Thus, no branch consistently dominates the parameter update dynamics. This structure prevents "gradient starvation" of the shallow branch—an empirically verified source of accuracy degradation in quantized multi-task models.

Balanced gradient flow results in smoother loss curves and accelerated QAT convergence, and the approach is compatible with convergence analyses for QAT based on the Straight-Through Estimator (STE), since it reduces the variance in backpropagated gradients introduced by quantization.

5. Attention Distribution Alignment (ADA) Distillation

To preserve the semantic salience of features lost in quantization, ADA distillation aligns attention distributions between full-precision and quantized models without introducing additional trainable parameters.

• Channel Attention: SimAM attention weights per channel and spatial location $(i, j)$ are defined via inverse "energy":

$$E_{ij}^{-1} = \frac{(X_{ij} - \mu_x)^2}{4(\sigma_x^2 + \lambda)} + 0.5,\qquad A_{ij} = \sigma(E_{ij}^{-1})$$

where $X$ is a feature map, $\mu_x$ and $\sigma_x^2$ its channel mean and variance, and $\sigma$ is the sigmoid.

• Distribution Alignment: Normalize attention to a probability distribution $P_{ij}$ in the full-precision teacher and $Q_{ij}$ in the quantized model. ADA then minimizes $L_{\text{ADA}}$ (typically Jensen–Shannon divergence, or KL as ablated):

$L_{\rm ADA} = D_{\rm JS}(P \| Q)$

Ablation experiments show that JS divergence yields best results with N2UQ quantizers and KL with PACT/LSQ.

6. Empirical Performance and Experimental Protocol

Extensive multi-scenario benchmarking validates the effectiveness and generalizability of GABFusion and ADA:

• Backbones: YOLOv5s, YOLOv11s
• Quantization: N2UQ, PACT, LSQ, LSQ+ algorithms; bit-widths W4A4 (4b weights/activations), W3A3 (3b)
• Datasets: PASCAL VOC, MS COCO

Summary of Key Results:

Dataset / Metric	Full-precision	Baseline QAT (N2UQ)	GABFusion+ADA	$\Delta$ (Abs. Imp.)
VOC AP$_{50}$ (YOLOv5s)	85.9	82.1 (W4A4)	84.2 (W4A4)	+2.1
VOC AP$_{50}$ (YOLOv11s)	89.4	86.2 (W4A4)	87.4 (W4A4)	+1.2
COCO AP$_{50}$/mAP (YOLOv5s)	56.8/37.4	50.2/31.1	51.4/33.2	+1.2 / +2.1

The average mAP improvement on VOC is approximately 3.3%, and on COCO approximately 1.6%. When using YOLOv5s under 4-bit quantization, the accuracy gap to full-precision models is reduced to only 1.7%. Ablation on VOC confirms that both GABFusion and ADA contribute additively to performance gains across N2UQ, PACT, and LSQ quantizers.

7. Modularity, Integration, and Scope of Applicability

GABFusion is designed for maximal modularity and backward compatibility:

• Plug-and-Play: Implemented as a minimal add-on at feature-fusion points, requiring no change to backbone, neck, or head architecture.
• Quantizer-Agnostic: Operates with any QAT scheme (e.g., PACT, LSQ, N2UQ) without adjustment to their quantizers or training tricks.
• Inference-Efficient: All auxiliary computation exists only during training; deployment cost is identical to the original architecture.
• ADA Overhead: ADA distillation requires only a frozen full-precision teacher during training, with no effect during inference.

A plausible implication is that GABFusion and ADA can be applied to a broad range of multi-task and fusion-heavy architectures (e.g., detection, segmentation, and multi-modal tasks) undergoing aggressive quantization, provided the network exposes suitable fusion points for plug-in modules.

8. Relevance and Directions Relative to Broader Research

GABFusion addresses the recognized issue of multi-branch gradient starvation in quantized networks, supplementing standard numeric losses (e.g., CIoU) with structural interventions and feature-level supervision. By combining balanced optimization dynamics with semantic feature preservation, the method provides a theoretical and empirical framework for robust low-bit training in the presence of multi-branch fusion—a scenario increasingly critical as practical deployments of resource-efficient multi-task networks become prevalent. These advances locate GABFusion at the intersection of quantization theory, gradient optimization, and feature distillation, and define a reference methodology for subsequent research in balanced, quantization-friendly multi-task learning.