Overcoming Rank Collapse in Feedback Alignment

Technical Summary: Overcoming Rank Collapse in Feedback Alignment

Problem Statement
Backpropagation (BP) is the standard method for training neural networks but is often considered biologically implausible due to the "weight transport problem," which requires feedback weights to be the exact transpose of forward weights. Feedback Alignment (FA) was proposed as a biologically plausible alternative, utilizing fixed random feedback weights. While FA can train shallow Multi-Layer Perceptrons (MLPs) and very shallow Convolutional Neural Networks (CNNs) effectively, it fails to scale to deeper architectures. Previous attempts to improve FA, such as Direct Feedback Alignment (DFA) or training feedback weights, have not consistently closed the performance gap with BP on complex tasks like CIFAR-100 or in deep ResNets. The specific mechanism causing this failure in deep networks remains unclear.

Methodology and Investigation
The authors investigate the geometric properties of the gradients in FA-trained networks compared to BP-trained networks. They hypothesize that the failure of FA in deep networks is linked to the "effective rank" (a measure of effective dimensionality) of the gradient signals.

1. Diagnosis of Rank Collapse: The authors trained CNNs of varying depths (1 to 4 layers) on CIFAR-10 using both BP and FA. They observed that while FA models eventually align forward weights with random feedback weights, the gradient alignment remains low. Crucially, they measured the effective rank of the gradients and found that FA gradients collapse into a low-dimensional subspace much earlier and more severely than BP gradients. This low-rank geometry restricts parameter updates to a limited set of directions, hindering the exploration of the parameter space necessary for continued alignment.
2. Proposed Interventions: To counteract this rank collapse, the authors evaluated two mechanisms designed to increase the effective dimensionality of the updates and representations:
   • Muon Optimizer: An optimizer that adds an orthogonalization step to the momentum update. It sets all non-zero singular values of the momentum matrix to 1, effectively flattening the spectrum and increasing the rank of the update direction.
   • Hidden Activity Normalization: The application of Batch Normalization (BN) to hidden layers. Previous work suggests BN promotes orthogonal hidden representations, preventing them from collapsing into low-rank subspaces.
3. Experimental Setup: The study utilized CIFAR-10, CIFAR-100, STL-10, and Tiny ImageNet datasets. Architectures included custom CNNs, AlexNet, and ResNet-18. The authors compared standard optimizers (SGD, AdamW) against Muon, with and without Batch Normalization.

Key Results

• Diagnosis: In standard FA, the effective rank of gradients drops significantly as network depth increases, often stabilizing at very low values (e.g., ~12 for a 4-layer network), whereas BP gradients maintain high dimensionality. This low-dimensional trajectory correlates with poor convergence and low accuracy in deep FA models.
• Muon and BN Efficacy: Both Muon and BN independently mitigate the performance degradation of FA in deep networks.
  • On a 4-layer CNN, Muon improved FA test accuracy from 59.5% (AdamW) to 69.5%.
  • Both methods restored high-dimensional gradient trajectories and improved weight/gradient alignment.
• Scaling to Complex Tasks: The combination of Muon and Batch Normalization yielded the best results across all benchmarks:
  • STL-10: FA accuracy on ResNet-18 improved from 14.5% (baseline) to 64.2% (Muon + BN).
  • CIFAR-100: FA accuracy on ResNet-18 improved from 1.4% (baseline) to 46.1% (Muon + BN).
  • Tiny ImageNet: FA accuracy on ResNet-18 improved from 0.7% (baseline) to 24.3% (Muon + BN).
• Complementarity: The methods are complementary; Muon acts on the geometry of the parameter update, while BN acts on the forward representations. Their combination consistently outperformed either method alone.
• Control Experiments: The authors demonstrated that simply adding isotropic noise to updates (to increase dimensionality) did not improve FA, suggesting that the benefit comes from preserving useful update directions via orthogonalization, not merely increasing rank arbitrarily. Conversely, forcing updates to be low-rank (via low-rank SVD approximation) catastrophically degraded FA performance.

Significance and Claims
The paper claims to identify low-dimensional gradient dynamics (rank collapse) as a primary obstacle preventing Feedback Alignment from scaling to deep architectures. The authors argue that the failure of FA is not solely an alignment problem but fundamentally a dimensionality problem.

By inducing higher-dimensional update geometry through Muon and maintaining high-dimensional representations via Batch Normalization, the authors demonstrate that FA can be scaled to deeper CNNs and harder datasets, achieving performance previously unattainable for this class of biologically plausible algorithms. The work suggests that for FA to function effectively in deep networks, the optimization trajectory must be constrained to avoid collapsing into low-dimensional subspaces.

The authors maintain a modest stance regarding biological plausibility: while the mechanism of FA (random feedback) is biologically motivated, the specific interventions (Muon and Batch Normalization) are not presented as direct models of synaptic plasticity. Muon requires global spectral information of weight matrices, and BN relies on batch statistics, which may not be strictly local. However, the results provide a proof-of-concept that alternative credit-assignment rules can be made robust to depth if the geometric properties of the updates are properly managed.