MAN++: Scaling Momentum Auxiliary Network for Supervised Local Learning in Vision Tasks

1. Problem Statement

Deep learning currently relies on End-to-End (E2E) backpropagation, which suffers from three critical limitations:

• Update Locking: Parameters are only updated after a full forward and backward pass, preventing parallel updates and reducing training efficiency.
• High GPU Memory Consumption: Storing activations for the entire network during backpropagation requires massive memory, limiting model depth and batch sizes.
• Biological Implausibility: The global error signal propagation does not mimic the localized learning observed in biological synapses.

Supervised Local Learning attempts to solve this by partitioning the network into gradient-isolated blocks, each trained with an auxiliary network and a local loss. However, existing local learning methods suffer from a significant performance gap compared to E2E training. This is caused by gradient isolation: earlier blocks are optimized only on local signals and lack feedback regarding how their representations will be utilized by downstream blocks, leading to suboptimal feature learning.

2. Methodology: Momentum Auxiliary Network++ (MAN++)

The authors propose MAN++, a scalable framework designed to bridge the performance gap in supervised local learning by enabling lightweight information exchange between adjacent blocks without breaking gradient isolation.

Core Components

MAN++ introduces two key mechanisms positioned between the main network and the auxiliary network of each block:

1. Exponential Moving Average (EMA) Module:

   • Instead of relying solely on local loss, the auxiliary network of block $j$ incorporates parameter information from the subsequent block ($j+1$).
   • It uses the EMA of parameters from block $j+1$ to update the auxiliary network of block $j$. This propagates "contextual" information from deeper layers back to earlier layers, mitigating the myopic nature of local learning.
   • To manage memory, the EMA source is limited to the first computational layer of the subsequent block (e.g., the first residual unit in ResNet or the first Transformer block in ViT).

2. Learnable Scale and Bias (LB) Module:

   • Directly applying EMA parameters can cause feature statistics mismatches because the distributions of features learned by isolated blocks differ.
   • To address this, MAN++ introduces a learnable scaling factor ($s_j$) and a learnable bias ($b_j$).
   • The update rule combines the standard local gradient descent with the EMA transfer, weighted by $s_j$:
     $$\gamma_j \leftarrow s_j \times \text{EMA}(\gamma_j, \theta_{j+1})$$
     $$(\gamma_j, b_j) \leftarrow (\gamma_j, b_j) - (2 - s_j) \times \eta_a \times \nabla L(\hat{y}_j, y)$$
   • This allows the network to dynamically balance the contribution of the global context (EMA) and the local objective, compensating for distribution shifts.

Architecture Integration

• Plug-and-Play: MAN++ is designed to be integrated into any supervised local learning framework (e.g., DGL, InfoPro).
• Granularity: The scaling and bias parameters are specific to each local block but shared within that block.
• Training Strategy: The paper utilizes Pipeline Parallel Local Learning (PPLL) for multi-GPU training, assigning one local block and its auxiliary head to a single GPU to maximize parallelism and minimize communication overhead.

3. Key Contributions

1. Novel Framework (MAN++): Proposes a momentum-based parameter transfer mechanism using EMA and learnable scale/bias to couple gradient-isolated blocks, effectively addressing the "short-sightedness" of local learning.
2. Plug-and-Play Design: The method is architecture-agnostic and can be seamlessly applied to various backbones (ResNet, ViT) and tasks (classification, detection, segmentation) with minimal additional memory overhead.
3. Theoretical Analysis: Provides theoretical bounds on parameter complexity, FLOPs, and memory usage, proving that the overhead is negligible for deep networks and that the method maintains convergence guarantees comparable to E2E training.
4. Comprehensive Evaluation: Extensive experiments across six datasets (CIFAR-10, SVHN, STL-10, ImageNet, COCO, Cityscapes) demonstrate state-of-the-art performance.

4. Experimental Results

The authors evaluated MAN++ against E2E training and existing local learning baselines (DGL, InfoPro).

• Accuracy:

  • Image Classification: On ImageNet, MAN++ significantly narrows the gap to E2E. For example, with ResNet-101 (K=4), it reduces Top-1 error by ~1.15% compared to InfoPro, achieving performance comparable to E2E. On CIFAR-10, it reduces test error by up to 38% compared to baselines.
  • Object Detection (COCO): MAN++ achieves higher mAP than E2E training while using significantly less memory. For ResNet-152, mAP improved by +0.4% while reducing GPU memory by ~74%.
  • Semantic Segmentation (Cityscapes): MAN++ consistently outperforms both E2E and InfoPro baselines, improving mean IoU by up to 1.8% on DeepLabV3-R101.

• Memory Efficiency:

  • MAN++ reduces peak GPU memory usage by 27% to 65% compared to E2E training.
  • For deep networks like ResNet-110, memory usage dropped from 9.26 GB (E2E) to 3.22 GB (MAN++).

• Training Speed:

  • While single-GPU training is slightly slower due to auxiliary heads, the PPLL multi-GPU strategy makes MAN++ training speeds comparable to or faster than E2E (e.g., ViT-Tiny on ImageNet: 458s/epoch vs. 503s for E2E).

• Ablation Studies:

  • Experiments confirm that both the EMA mechanism and the Learnable Bias are necessary. EMA provides global context, while the Bias corrects feature distribution mismatches. Their combination yields the best results.
  • CKA Analysis: Centered Kernel Alignment scores show that MAN++ produces feature representations much closer to E2E training than other local methods, indicating better semantic consistency.

5. Significance

• Scalability: MAN++ enables the training of very deep models on hardware with limited memory, making high-performance vision tasks accessible without requiring massive GPU clusters.
• Practicality: By achieving E2E-level accuracy with significantly reduced memory and comparable training speeds (via PPLL), MAN++ offers a viable, practical alternative to traditional backpropagation for large-scale vision tasks.
• Biological Plausibility: The method aligns more closely with biological learning principles (local updates with limited feedback) while overcoming the historical accuracy penalty associated with such approaches.

In conclusion, MAN++ successfully resolves the trade-off between the efficiency of local learning and the accuracy of end-to-end training, establishing a new standard for scalable supervised local learning in computer vision.