Neural LoFi with Backward Coupling: A Spectral Theory of Cross-Layer Feature Learning
Organizations: Information, Learning, and Physics Laboratory (IDEPHICS) EPFL, Lausanne, Switzerland
Abstract
During joint training, backpropagation allows later layers to reshape earlier representations, a feedback mechanism absent from purely layer-wise learning. We extend Neural Low-Degree Filtering (Neural LoFi), a recently proposed layer-wise spectral surrogate for feature learning, to capture this interaction. Starting from the exact backpropagation gradients, we derive a local coupling between a neuron's incoming weights and its corresponding next-layer column. In a controlled small-initialization regime with smooth activations and fixed surroundings, these dynamics amplify singular modes of a cross-layer correlation operator. Each growing mode pairs an earlier feature with an upper-layer response through shared neuron coefficients. This motivates a backward-coupling algorithm in which upper-layer sensitivities generate sample-dependent effective labels for earlier spectral steps, while corrected features refine upward transmission. A function-space variational principle characterizes the selected feature-response pairs by maximizing the label correlation of their product under kernel-norm constraints. For three-layer networks, we derive self-consistent continuous-time dynamics for the filtered hidden-layer features, the activation kernels, and the network output under the proposed backward coupling. In a hierarchical multi-index model, we establish a feature-recovery separation: backward coupling recovers lower-layer directions invisible to the original forward label-weighted statistic. The linear specialization recovers residual-dependent singular-mode dynamics, while paired modes provide a dynamical mechanism for linked low-rank covariance components consistent with the rainbow-network picture. Experiments with fully connected and convolutional networks on binary CIFAR-10 show improved accuracy over forward Neural LoFi at matched widths.
Figures & tables
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Stage | Predicted sample scale |
|---|---|
| Visible recovery | |
| Upper recovery | |
| Learned backward score | |
| Hidden recovery given the score | |
| Corrected upper Forward pass |
Neural Network Training Dynamics
Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.
See all 507 papers