AXIOM: Activation-Conditioned Momentum as Associative Memory

Classical momentum-based optimizers, such as Stochastic Gradient Descent (SGD) with momentum and Adam accumulate gradient statistics isotropically, treating every direction in parameter space identically and remaining blind to the geometry of the activations that produced each gradient. Motivated by the Nested Learning view of optimization as a system of associative memories, we in this study, reformulate momentum as the solution of a regularized least squares objective that maps an input activation to its local surprise signal, rather than an exponential moving average, yielding a closed form optimizer we call AXIOM (Activation-guided eXpressive Iterative Optimization with associative Memory). We derive the resulting update independently through a Sherman–Morrison expansion and through Fenchel–Rockafellar duality, showing that both routes converge on the same rank-one, activation-conditioned correction. We characterize AXIOM as a discrete-time linear time-varying system, prove global uniform exponential stability of its unforced dynamics, and establish an O(1/T) + O(σ 2 g ) nonconvex convergence rate for the forced stochastic system, matching classical firstorder guarantees. To test whether activation-aware momentum provides practical benefit under heterogeneous and non-stationary feature geometries, we introduce Wavelet2VecNet-a framework that couples pretrained Self-Supervised Learning (SSL)-based contextual embeddings (Wav2Vec2) with a two-dimensional wavelet packet transform decomposition for compact utterance-level speech representation learning, and evaluate it across five speech processing tasks namely, dysarthric speech severity-level assessment, infant cry pathology classification, environmental acoustic scene and sound classification, and deepfake speech source attribution. AXIOM matches or exceeds Adam and SGD on four of five tasks, with the largest gains (+7 to +14 accuracy points) on the most acoustically heterogeneous tasks, and no advantage on a near-saturated verification task – a pattern consistent with our theoretical account of when activation-conditioned momentum should help. We report ablations across wavelet bases and feature backbones, including a case where AXIOM’s advantage does not persist, and discuss the limitations of the present activation-only, linear-layer analysis.