Technical Summary: Hyperparameter Transfer for Dense Associative Memories

Problem Statement

Modern deep learning relies heavily on hyperparameter (HP) transfer, where optimal settings (e.g., initialization scales, learning rates) are identified on small models and datasets and then transferred to larger scales to avoid prohibitively expensive tuning. While theoretical frameworks for HP transfer exist for feed-forward networks (e.g., based on the $\mu$P heuristic), they have not been developed for Dense Associative Memories (DenseAMs).

DenseAMs present unique challenges for existing transfer methods due to three specific architectural features:

1. Weight Sharing: The energy function requires weights to be shared across layers and within operations, unlike standard feed-forward networks.
2. Residual Dynamics: The dynamical equations describing energy-descent natively contain residual connections.
3. Steep Activations: To achieve superior information storage scaling laws, DenseAMs utilize rapidly peaking activation functions (e.g., high-power ReLUs, softmax) that are rare in standard feed-forward architectures.

These factors create a complex landscape where standard dimensional analysis fails to predict how hyperparameters should scale with model size ($N, K$), dataset size ($P$), and batch size ($B$).

Methodology

The authors initiate the study of $\mu$P-type HP transfer for DenseAMs by deriving explicit prescriptions for scaling hyperparameters in the proportional scaling regime, defined as:
$$N \to \infty, \quad K = \kappa N, \quad P = \rho N, \quad B = \beta P$$
where $\kappa, \rho, \beta$ are fixed constants.

Model Definition

The study focuses on a shallow DenseAM update rule performing a single step of gradient descent on an energy landscape:
$$f(x) = s_2 W^\top \sigma_C [s_1 W g(x) + b] + c$$
where $W \in \mathbb{R}^{K \times N}$ are tied weights, $g$ is a fixed non-linearity, and $\sigma_C$ is a potentially centered activation function. The training objective is a denoising loss:
$$L(W) = \frac{1}{2B} \sum_{\mu=1}^B \| f_W(x_\mu + \epsilon_\mu) - x_\mu \|^2$$

Key Technical Modifications

The authors identify that standard parameterizations fail for DenseAMs with non-linear activations due to "mean-induced spikes" in the Gram matrix of activations. To address this, they introduce centering:

• Centering Operator: A projection matrix $C = I_K - \frac{1}{K}\mathbf{1}_K\mathbf{1}_K^\top$ is applied to pre- and post-activations.
• Necessity: For non-linear activations (ReLU, Softmax) with non-zero means, failing to center leads to a rank-one spike in the activation Gram matrix ($S^\top S$) that scales differently from the bulk eigenvalues. This mismatch causes numerical instability and prevents HP transfer. Centering removes this spike, ensuring all internal statistics scale consistently.

Derivation Approach

The authors adapt the $\mu$P heuristic by performing a dimensional analysis to satisfy three desiderata:

1. Stability: Initialization statistics (mean/variance) of pre-activations ($Z$) and outputs ($F$) must remain $O(1)$ as $N, K \to \infty$.
2. Maximality: Parameter updates ($\Delta Z, \Delta F$) after $T$ steps must remain $O(1)$.
3. Balance: The various terms contributing to updates (arising from the two appearances of tied weights $W$) must be balanced in magnitude.

By analyzing the gradients of the MSE loss and the resulting updates for both SGD and Adam optimizers, the authors derive specific scaling laws for initialization scales ($s_1, s_2$) and learning rates ($\eta_W, \eta_b, \eta_c$).

Key Contributions and Results

1. Explicit Scaling Prescriptions

The paper provides a table of scaling rules (Table 1) for different activation functions and optimizers. Key findings include:

• Linear Activations: In the proportional regime, $s_1 \sim N^{-1/2}$, $s_2 \sim K^{-1/2}$, and $\eta_W \sim K$ (for SGD).
• ReLU-type Activations: Centering is critical. Without centering, HP transfer fails due to instability. With centering, the linear prescriptions hold for SGD and Adam. Adam remains stable even for high-power ReLUs ($p=10$), whereas SGD becomes unstable at large learning rates.
• Softmax Activations:
  • SGD Instability: SGD suffers from numerical instability at large scales due to "row amplification feedback," where rows with larger softmax weights receive disproportionately large updates, leading to localization and divergence. The paper does not provide a stable SGD prescription for softmax.
  • Adam Success: Adam is scale-invariant and exhibits successful HP transfer for centered softmax DenseAMs with $s_1 \sim N^{-1/2}$, $s_2 \sim K^{1/2}$, and $\eta_W \sim 1$.

2. Empirical Validation

The authors validate these theoretical findings across multiple settings:

• Proportional Scaling: Experiments on isotropic Gaussian data and MNIST (via plaquette downsampling) demonstrate that training dynamics (loss curves, internal feature covariances) collapse across scales when the derived prescriptions are used.
• Anisotropic Data: The prescriptions work surprisingly well even for anisotropic and realistic data (MNIST), despite the theoretical derivations assuming isotropic inputs.
• Centering vs. Non-Centered: Empirical results (Figures 7, 8, 10, 11, 13) clearly show that non-centered models diverge or fail to transfer hyperparameters, while centered models exhibit consistent dynamics and successful transfer.

3. Dynamical Collapse

Beyond simple loss transfer, the paper demonstrates that the proposed parameterization leads to a collapse of training dynamics. Internal statistics, such as the RMS of weight updates and feature covariances, remain constant across different model sizes ($N, K$) when the scaling laws are applied.

Significance and Claims

The paper claims to be the first work to initiate the study of hyperparameter transfer for Dense Associative Memories. Its significance lies in:

• Bridging a Gap: It extends the theory of hyperparameter transfer from standard feed-forward networks to energy-based models with tied weights and steep activations.
• Practical Utility: It provides actionable, theoretically grounded rules for scaling DenseAMs, enabling practitioners to train large-scale models (e.g., Energy Transformers, NRGPT) without expensive large-scale hyperparameter tuning.
• Architectural Insight: It identifies centering as a fundamental requirement for the stability of DenseAMs with non-linear activations, a nuance previously unaddressed in the context of HP transfer.

The authors remain modest regarding the scope, noting that their analysis assumes isotropic data for derivations (though empirically robust otherwise) and that a full analytic description of training dynamics at infinite scale for non-linear DenseAMs remains a subject for future work. They do not claim to solve HP transfer for all modern architectures but specifically for the DenseAM family.