Hyperparameter Transfer for Dense Associative Memories

Abstract: Dense Associative Memory (DenseAM) is a promising family of AI architectures that is represented by a neural network performing temporal dynamics on an energy landscape. While hyperparameter transfer methods are well-studied for feed-forward networks, these methods have not been developed for settings in which weights are shared across layers and within the layer, which is common in DenseAMs. Additionally, DenseAMs utilize rapidly peaking activation functions that are rarely used in feed-forward architectures. The confluence of these aspects makes DenseAM a challenging framework for using existing methods for hyperparameter transfer. Our work initiates the development of hyperparameter transfer methods for this class of models. We derive explicit prescriptions for how the hyperparameters tuned on small models can be transferred to models trained at scale. We demonstrate excellent agreement between these theoretical findings and empirical results.

• The paper introduces a theoretical framework that derives precise scaling laws for initializing DenseAM networks and transferring hyperparameters.
• It demonstrates that centering activations and choosing optimizers like Adam over SGD are crucial for maintaining stability and preventing divergent dynamics.
• Empirical results on tasks such as denoising and MNIST validate that the prescribed transfer protocol produces a collapse in loss and internal statistics across model scales.

Hyperparameter Transfer in Dense Associative Memory Networks

Introduction

This paper investigates explicit hyperparameter (HP) transfer prescriptions for Dense Associative Memory (DenseAM) networks, a class of energy-based neural architectures characterized by tied weights across layers and steep activation functions. While dimensional analysis and $\mu$P-based scaling heuristics have enabled reliable HP transfer in feed-forward architectures and transformers, the structural properties and dynamic requirements of DenseAMs preclude direct application of prior HP transfer protocols. The research systematically derives scaling laws for HP transfer in DenseAMs and validates these prescriptions empirically, demonstrating both stability and dynamical consistency across scale.

DenseAMs generalize classical Hopfield networks by operating with significantly higher storage capacities and evolving network states via energy descent dynamics, thus enabling architectural modules such as Energy Transformers and energy-based GPT alternatives. The paper addresses the scaling behavior of initialization parameters and learning rates ($s_1$, $s_2$, $\eta_W$, $\eta_b$, $\eta_c$) with respect to model, data, and batch dimensions ($N$, $K$, $P$, $B$) and quantifies transfer efficacy for various activation and optimization schemes.

Theoretical Prescription for Hyperparameter Scaling

DenseAM networks are formalized by temporally evolving the state vector via a gradient step on the energy landscape. The update equation is

$s_1$0

where $s_1$1 and $s_1$2 are pointwise non-linearities, $s_1$3 denotes tied weights, and $s_1$4, $s_1$5 are biases. The proportional scaling regime is defined as $s_1$6, $s_1$7, $s_1$8, $s_1$9, with $s_2$0, $s_2$1, $s_2$2 fixed.

Three core desiderata for HP transfer are imposed:

1. Stability: All internal statistics (e.g., preactivations, outputs) are $s_2$3 at initialization.
2. Maximality: Magnitudes of key updates (preactivations, outputs) remain $s_2$4 over the first $s_2$5 optimization steps as $s_2$6.
3. Balance: Weight update components from different tied weight positions scale identically, ensuring consistent dynamics across scale.

For linear activations, analysis reveals HP transfer and collapse of training dynamics under:

• $s_2$7
• $s_2$8
• $s_2$9

Empirical results confirm consistency of loss curves and internal statistics across scales

Figure 1: Learning rate transfer and training loss dynamics for a linear DenseAM trained on a denoising objective.

For non-linear activations, especially those with nonzero mean (e.g., ReLU$\eta_W$0, softmax), centering pre- and post-activations is necessary to avoid instability due to emerging spectral spikes in the Gram matrix representation.

Figure 2: Learning rate transfer for a ReLU DenseAM trained on the denoising objective; centering is critical for transfer stability.

Empirical Validation and Activation Effects

ReLU-Type Activations

DenseAMs with ReLU$\eta_W$1 activations exhibit HP transfer only after centering both pre- and post-activations. Uncentered models show an instability originating from a rank-one spectral spike associated with activation mean, resulting in divergent dynamics with increasing scale. Centered models achieve loss and internal statistic collapse across scaled model/data dimensions.

Figure 3: Spike and bulk behaviors for centered and non-centered ReLU DenseAM trained by mini-batch SGD, demonstrating destabilizing spectral spikes for uncentered networks.

Figure 4: HP transfer for centered and non-centered ReLU DenseAM trained by mini-batch SGD; centering yields transfer and dynamical collapse.

For higher powers of ReLU ($\eta_W$2), Adam stabilizes transfer where SGD becomes unstable due to exacerbated nonlinearity-induced variance.

Softmax Activations

Softmax DenseAMs require scaling $\eta_W$3 for stable initialization, but SGD fails to deliver consistent transfer due to row amplification feedback. Adam optimizer, being scale-invariant, manages centering-induced stability and supports reliable HP transfer.

Figure 5: Learning rate transfer for DenseAM with centered vs. non-centered softmax activations trained with Adam.

SGD on softmax DenseAMs rapidly amplifies rows with concentrated softmax weights, causing localizations and numerical instability at large scales. Adam avoids this issue.

Figure 6: Learning rate transfer for softmax DenseAM trained with SGD and Adam; only Adam achieves transfer stability.

Figure 7: Training dynamics of softmax DenseAM with Adam in the proportional regime; participation ratio remains stable under Adam, diverges under SGD.

Scaling Regimes and Dataset-Model Co-Sizing

Two scaling regimes are characterized:

• Proportional scaling (all dimensions): Consistency and transfer across $\eta_W$4, $\eta_W$5, $\eta_W$6, $\eta_W$7; requires dynamic resizing of input features, e.g., coarse-graining for vision datasets.
• Width-only scaling ($\eta_W$8): Model scaling with fixed data dimensions; transfer persists for $\eta_W$9.

Empirical results demonstrate robust HP transfer for isotropic and anisotropic data distributions, including realistic datasets like MNIST after input dimension downsampling.

Figure 8: HP transfer for DenseAMs trained on MNIST using ReLU or softmax activations across proportional scale.

Figure 9: Comparison of denoisers across scales and downsampling factors for DenseAMs trained on MNIST; output agreement confirms transfer protocol efficacy.

Dynamical Consistency and Internal Statistic Collapse

Training dynamics for centered DenseAMs reveal that loss, internal representation statistics, and Gram matrix eigenvalue spectra collapse across scale under prescribed HP transfer, providing evidence for robust scaling limits.

Figure 10: Dynamical consistency across scales for MSE and preactivation updates in ReLU DenseAMs trained with SGD.

Figure 11: Training dynamics collapse for ReLU$\eta_b$0 DenseAMs with SGD; all update magnitudes remain invariant with increasing $\eta_b$1.

Implications and Outlook

The findings establish that HP transfer for DenseAMs demands deliberate scaling of initialization and learning rates, centering of activations, and optimizer selection tailored to the activation function's spectral properties. For energy-based architectures with tied weights and steep activations, failure to align mean and bulk statistics results in instabilities and transfer breakdowns. The practical prescription enables efficient scaling from small to large models, facilitating deployment of DenseAMs in increasingly demanding AI systems, and provides foundational methodology for future extensions to more complex energy-based modules and transformer variants.

Theoretically, the results support the universality of dynamical collapse in energy-based networks when HPs are scaled according to precise dimensional analysis, complementing scaling law studies in feed-forward and attention-based networks. Further analytic work is warranted to characterize training dynamics in the infinite scale limit and to extend transfer protocols to modular, residual, and multi-token energy architectures.

Conclusion

This paper delivers a foundational framework and empirically validated prescription for hyperparameter transfer in Dense Associative Memory networks, addressing the unique challenges posed by tied weights and steep activations. By combining theoretical dimensional analysis with centering and optimizer selection, robust transfer and dynamical consistency across scale are achieved. The results underpin future development of principled scaling laws for energy-based AI architectures and guide practical deployment in both memory retrieval and general computation contexts (2605.10164).