Technical Summary: PILOT (Policy-Informed Learned Optimization for Adaptive Deep Network Training)

Problem Statement

Despite the centrality of optimization in deep learning, most existing optimizers rely on update structures with functional forms fixed prior to training. This static design limits their ability to adapt to the dynamic nature of the loss landscape, where training regimes shift between stable, noisy, and inconsistent states. While adaptive methods like Adam and AdamW improve robustness through fixed moment estimation and variance normalization, and newer methods like Lion and Sophia explore sign-based or curvature-aware updates, they all enforce a predefined trade-off between momentum, normalization, and update direction. Learned optimizers offer a broader design space but often require expensive offline search, meta-training, or task-specific generation. There remains a gap for an online adaptive mechanism that adjusts update behavior during standard training without these computational burdens.

Methodology

The paper proposes PILOT (Policy-Informed Learned OpTimizer), an online adaptive first-order optimizer that dynamically modulates its update behavior based on a signal of local training stability.

1. Gradient-Direction Agreement Signal

PILOT utilizes the agreement between consecutive stochastic gradients as a compact indicator of local optimization stability.

• Signal Definition: For successive gradients $g_t$ and $g_{t-1}$, a directional agreement signal $r_t$ is computed using numerically stabilized cosine similarity.
• Smoothing: To reduce short-term fluctuation, an exponential moving average $\rho_t$ is maintained.
• Interpretation: Positive values indicate consistent descent directions (smooth regions), values near zero suggest weak consistency (noise), and negative values indicate rapid changes or instability.

2. Learnable Policy

A compact polynomial policy maps the smoothed agreement signal $\rho_t$ to three scalar control variables:

• $p_{m,t}$: Controls momentum reliance (interpolation between momentum estimates and raw gradients).
• $p_{v,t}$: Controls variance-normalization strength (scaling of the update magnitude).
• $p_{s,t}$: Controls sign-based behavior (degree to which the update direction is determined by the sign of the gradient).

These variables are parameterized as univariate polynomials of degree $d$ followed by sigmoid activations, resulting in only $3(d+1)$ learnable coefficients. This formulation allows PILOT to continuously interpolate between Adam-like adaptive updates, raw-gradient updates, and sign-dominated updates.

3. Policy-Controlled Update Rule

The parameter update $\theta_{t+1}$ is constructed using first- and second-moment estimates ($\hat{m}_t, \hat{v}_t$). The policy variables modulate the combination of these estimates:

• The momentum estimate is interpolated with the current gradient based on $p_{m,t}$.
• The final update direction is scaled by the variance estimate raised to the power of $p_{v,t}$ and modified by the sign function controlled by $p_{s,t}$.
• This structure generalizes existing methods: setting specific fixed values for the policy variables recovers behaviors similar to Adam or sign-based optimizers.

4. Online Policy Learning

The policy coefficients are updated online during training using a one-step meta-gradient estimate. After observing the gradient of the subsequent mini-batch, the policy is updated to minimize the loss of that next step. This approach avoids offline search, separate meta-training phases, and backpropagation through the full optimization trajectory, keeping the computational overhead low.

Key Contributions

1. Online Adaptive Optimizer: Introduction of PILOT, which dynamically modulates momentum reliance, variance normalization, and sign-based behavior during training.
2. Landscape-Aware Signal: Utilization of gradient-direction agreement as a proxy for local stability, enabling the update rule to respond to changing training dynamics without explicit curvature estimation.
3. Compact Policy: Proposal of a polynomial policy with minimal learnable parameters ($3(d+1)$) that adapts optimizer behavior without offline search.
4. Unified Formulation: Demonstration that existing update behaviors (e.g., Adam-style, sign-based) can be represented as fixed configurations within the same general formulation.
5. Empirical Validation: Comprehensive evaluation showing PILOT's competitiveness and superiority in specific settings compared to strong baselines.

Experimental Results

PILOT was evaluated on FashionMNIST and CIFAR-10 using CNN and ResNet-18 architectures, comparing against Adam, AdamW, Lion, Sophia, and AdaBelief.

• Performance: PILOT consistently achieved the highest accuracy across all configurations.
  • FashionMNIST: 94.13% (CNN) and 95.71% (ResNet-18).
  • CIFAR-10: 81.94% (CNN) and 93.42% (ResNet-18).
• Stability: PILOT demonstrated the lowest loss variance in several settings (e.g., 0.0001 on CIFAR-10/ResNet-18) and maintained smoother gradient trajectories compared to baselines like Adam and AdamW, which exhibited larger spikes.
• Ablation Studies:
  • Momentum ($p_m$): Adaptive control of momentum reliance was crucial for performance on complex datasets (CIFAR-10), while fixed momentum sufficed for simpler settings.
  • Variance Normalization ($p_v$): Learning the normalization strength provided consistent advantages on CIFAR-10, whereas fixed Adam-style normalization was competitive on FashionMNIST.
  • Sign-Based Behavior ($p_s$): While disabling sign-based behavior sometimes improved loss minimization, the fully learned variant generally achieved better classification accuracy, suggesting a beneficial trade-off between stability and predictive performance.
• Transferability: A policy trained on CIFAR-10 and frozen for use on FashionMNIST outperformed AdamW and nearly matched a fresh PILOT instance trained on the target dataset, suggesting the learned policy captures broadly applicable optimization dynamics.

Significance and Claims

The paper claims that PILOT bridges the gap between rigid, hand-designed optimizers and computationally expensive learned optimizers. By learning a small online policy conditioned on gradient-direction agreement, PILOT provides a practical mechanism for adaptive behavior while retaining a simple first-order optimization framework.

The authors posit that the optimization regime changes over training, and a single fixed update rule is often insufficient. PILOT demonstrates that adapting the update structure during training can improve performance across both compact and deeper convolutional models. The results suggest that the learned polynomial encodes optimization dynamics that generalize beyond the specific dataset on which they were learned, offering a promising direction for flexible optimizer design in deep neural network training.

Limitations and Future Work: The authors acknowledge that current evaluations are limited to image-classification benchmarks and convolutional architectures. The generality of PILOT for larger datasets, transformer-based architectures, and sequence models remains unestablished. Future work aims to evaluate PILOT on Vision Transformers and language models, analyze the evolution of policy variables, and incorporate additional online signals such as gradient variance or curvature approximations.