# AYLA: Architecting a loss landscape in shallow neural networks to accelerate feature recovery

arXiv: https://arxiv.org/abs/2504.01875 · HTML: https://mlsanity.com/paper/2504.01875/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 30 references that are in the map.
- First posted on arXiv 2 Apr 2025.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning](https://mlsanity.com/paper/1810.01075/) (2018)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Towards moderate overparameterization: global convergence guarantees for training shallow neural networks](https://mlsanity.com/paper/1902.04674/) (2019)
- [An Exponential Learning Rate Schedule for Deep Learning](https://mlsanity.com/paper/1910.07454/) (2019)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [Asymptotics of Wide Networks from Feynman Diagrams](https://mlsanity.com/paper/1909.11304/) (2019)
- [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://mlsanity.com/paper/2305.18270/) (2023)
- [Learning ReLUs via Gradient Descent](https://mlsanity.com/paper/1705.04591/) (2017)
- [The dynamics of representation learning in shallow, non-linear autoencoders](https://mlsanity.com/paper/2201.02115/) (2022)
- [Is Deeper Better only when Shallow is Good?](https://mlsanity.com/paper/1903.03488/) (2019)
- [Convergence of gradient descent for deep neural networks](https://mlsanity.com/paper/2203.16462/) (2022)
- [A Dynamical Central Limit Theorem for Shallow Neural Networks](https://mlsanity.com/paper/2008.09623/) (2020)
- [Stability and Generalization Analysis of Gradient Methods for Shallow Neural Networks](https://mlsanity.com/paper/2209.09298/) (2022)
- [Minimum Width of Leaky-ReLU Neural Networks for Uniform Universal Approximation](https://mlsanity.com/paper/2305.18460/) (2023)
- [Why Shallow Networks Struggle to Approximate and Learn High Frequencies](https://mlsanity.com/paper/2306.17301/) (2023)
- [Adaptive Loss Scaling for Mixed Precision Training](https://mlsanity.com/paper/1910.12385/) (2019)
- [Normalization effects on deep neural networks](https://mlsanity.com/paper/2209.01018/) (2022)
- [How Does Gradient Descent Learn Features -- A Local Analysis for Regularized Two-Layer Neural Networks](https://mlsanity.com/paper/2406.01766/) (2024)
- [Post-Workshop Report on Science meets Engineering in Deep Learning, NeurIPS 2019, Vancouver](https://mlsanity.com/paper/2007.13483/) (2020)

## Abstract (arXiv)

Feature learning in shallow neural networks exhibits rich yet fragile dynamics, including prolonged plateaus, abrupt phase transitions, and sensitivity to optimization hyperparameters. While recent theoretical work has characterized these behaviors through the geometry of loss landscapes, saddle escape mechanisms, and emergent scaling laws, practical methods for actively shaping these dynamics remain limited. In this paper, we introduce AYLA, a principled loss reparameterization framework that dynamically modulates gradient magnitudes during training without altering the location of stationary points or optimal solutions. AYLA applies a smooth, sigmoid-controlled power-law transformation to empirical loss, yielding a state-dependent effective learning rate that accelerates descent in flat or saddle-dominated regions while stabilizing late-stage optimization. Crucially, AYLA preserves all critical points of the original objective, acting solely as a monotone transformation that reshapes optimization trajectories rather than objectives. We evaluate AYLA in controlled teacher student settings using two-layer tanh networks trained on synthetic Gaussian data. Across stochastic gradient descent and multiple loss-exponent schedules, AYLA consistently improves feature recovery. This evidence is observed in terms of weight alignment, per-neuron cosine similarity, hidden-activation correlation, and spectral properties of learned representations, while AYLA maintains competitive or faster loss convergence. Spectral analyses further demonstrate that AYLA mitigates rank collapse and promotes richer internal representations, signaling a transition from lazy to active feature-learning regimes. AYLA offers a lightweight, theoretically grounded way to improve shallow-network optimization, especially in resource-limited or noise-sensitive settings.