Resonance-Aware Learning Rate Tuning

Production-ready spectral resonance frequency-based optimizer scheduling for distributed transformer training.

Getting Started

• What: the optimizers subtree inside Gnosis.
• Why: it keeps this local concern documented where readers will look before editing or using it.
• How: read this entry first, then follow the local commands, child links, or file list below.
• Next: go back to the parent README when you need the wider package context.

Overview

This module implements two key components:

1. ResonanceFrequencyCalculator: Computes the natural resonance frequency (ω₀) of a transformer model from its architecture
2. AdaptiveResonanceLRScheduler: Dynamically adjusts learning rate based on spectral properties (σ₁/σ₂ ratio) during training

The system models layer dynamics as driven harmonic oscillators, where the resonance frequency depends on the model's spectral properties (singular values, condition number). By tracking the singular value ratio during training, we detect dimensionality collapse and proactively adjust learning rate to maintain stability.

Key Features

• Spectral Analysis: Computes σ₁/σ₂ ratio across layers to detect instability
• Phase Classification: Tracks training phases (INITIALIZATION → STABLE → DRIFT → CLIFF → RECOVERY)
• Adaptive Learning Rate: Adjusts LR in response to spectral drift and loss cliffs
• Cliff Detection: Triggers aggressive LR reduction when loss suddenly spikes
• Comprehensive Logging: Exports metrics to JSON for analysis and visualization
• Production-Ready: Type hints, error handling, extensive docstrings, logging

Installation

# Install PyTorch (required)
pip install torch>=2.0.0

# Clone repo and add to path
export PYTHONPATH="/path/to/gnosis/distributed-inference:$PYTHONPATH"

Quick Start

Basic Usage

from optimizers import (
    ResonanceFrequencyCalculator,
    AdaptiveResonanceLRScheduler,
    compute_spectral_ratio,
)
import torch.optim as optim

# 1. Calculate optimal learning rate from model architecture
calc = ResonanceFrequencyCalculator(
    hidden_dim=3072,      # Phi-3
    num_layers=32,
    attention_heads=32,
    intermediate_dim=8192,
)

optimal_lr = calc.compute_optimal_learning_rate(
    num_training_steps=100000
)
print(f"Optimal LR: {optimal_lr:.2e}")  # e.g., 5.23e-05

# 2. Initialize optimizer with predicted learning rate
optimizer = optim.AdamW(model.parameters(), lr=optimal_lr)

# 3. Create adaptive scheduler
scheduler = AdaptiveResonanceLRScheduler(
    optimizer=optimizer,
    base_lr=optimal_lr,
    target_sigma_ratio=2.0,
    enable_resonance_tuning=True,
    warmup_steps=100,
)

# 4. Training loop with resonance tuning
for epoch in range(num_epochs):
    for batch in dataloader:
        loss = forward_pass(batch)
        backward_pass(loss)
        optimizer.step()

        # Compute spectral ratio (per-layer σ₁/σ₂)
        sigma_ratio, diagnostics = compute_spectral_ratio(model)

        # Update scheduler
        metrics = scheduler.step(
            loss=loss.item(),
            spectral_ratios=[sigma_ratio] * model.num_layers,
            spectral_diagnostics=diagnostics,
        )

        # Log/monitor if needed
        if step % 100 == 0:
            print(f"Step {step}: LR={scheduler.current_lr:.2e}, "
                  f"σ₁/σ₂={sigma_ratio:.2f}, phase={metrics.phase.value}")

# 5. Export diagnostics
scheduler.export_metrics("training_metrics.json")

API Reference

ResonanceFrequencyCalculator

Computes natural resonance frequency from model architecture.

Constructor

calc = ResonanceFrequencyCalculator(
    hidden_dim: int,                    # Model hidden dimension
    num_layers: int,                    # Number of transformer layers
    attention_heads: int,               # Number of attention heads
    intermediate_dim: Optional[int] = None,  # FFN intermediate dimension
    vocab_size: int = 32000,            # Vocabulary size
)

Key Methods

# Compute ω₀ (resonance frequency)
omega_0 = calc.compute_resonance_frequency()  # Returns float (rad/step)

# Predict optimal learning rate
optimal_lr = calc.compute_optimal_learning_rate(
    base_scale: float = 1e-3,           # Baseline LR scale
    num_training_steps: int = 100000,   # Total training steps
)  # Returns float

# Get architecture diagnostics
diag = calc.get_diagnostics()  # Returns Dict[str, float]

Example: Architecture Diagnostics

calc = ResonanceFrequencyCalculator(3072, 32, 32, 8192)
diag = calc.get_diagnostics()
# {
#   "hidden_dim": 3072,
#   "num_layers": 32,
#   "attention_heads": 32,
#   "intermediate_dim": 8192,
#   "head_dim": 96,
#   "vocab_size": 32000,
#   "resonance_frequency": 0.0456,
#   "spectral_condition_number": 123.4,
#   "estimated_l2_norm": 45.6,
# }

AdaptiveResonanceLRScheduler

Dynamically adjusts learning rate based on spectral metrics during training.

Constructor

scheduler = AdaptiveResonanceLRScheduler(
    optimizer: Optional[Optimizer] = None,
    base_lr: float = 1e-3,                      # Initial learning rate
    target_sigma_ratio: float = 2.0,            # Target σ₁/σ₂ ratio
    enable_resonance_tuning: bool = True,       # Enable adaptive tuning
    cliff_threshold: float = 0.2,               # Loss spike threshold (20%)
    warmup_steps: int = 100,                    # Warmup steps
    decay_rate: float = 0.95,                   # Recovery decay rate
    min_lr: float = 1e-6,                       # Minimum learning rate
    max_lr: float = 1e-2,                       # Maximum learning rate
)

Key Methods

# Perform one scheduler step
metrics = scheduler.step(
    loss: float,                                # Current loss
    spectral_ratios: Optional[List[float]] = None,  # Per-layer σ₁/σ₂
    spectral_diagnostics: Optional[Dict] = None,    # Detailed per-layer metrics
)  # Returns ResonanceMetrics

# Set or replace optimizer
scheduler.set_optimizer(new_optimizer)

# Get recent metrics
metrics = scheduler.get_last_metrics()  # Returns Optional[ResonanceMetrics]

# Export metrics to JSON
scheduler.export_metrics("metrics.json")

# Get diagnostics summary
diag = scheduler.get_diagnostics()  # Returns Dict[str, Any]

Training Phases

The scheduler transitions through 5 phases based on spectral and loss metrics:

ResonanceMetrics

Diagnostics dataclass containing per-step metrics.

@dataclass
class ResonanceMetrics:
    timestamp: int                      # Unix timestamp
    step: int                           # Global step count
    loss: float                         # Current loss
    current_lr: float                   # Applied learning rate
    mean_sigma_ratio: float             # Mean σ₁/σ₂ across layers
    min_sigma_ratio: float              # Minimum σ₁/σ₂
    max_sigma_ratio: float              # Maximum σ₁/σ₂
    sigma_ratio_drift: float            # |σ₁/σ₂ - target|
    phase: ResonancePhase               # Current training phase
    cliff_triggered: bool               # Whether cliff was detected
    lr_adjustment_factor: float         # LR multiplier applied
    spectral_layer_diagnostics: Dict    # Per-layer details

    # Methods
    def to_dict(self) -> Dict           # Serialize to dictionary

Utility Functions

# Compute σ₁/σ₂ ratio for model
sigma_ratio, per_layer_diag = compute_spectral_ratio(
    model: Optional[nn.Module] = None,
    layer_index: Optional[int] = None,  # Analyze specific layer (optional)
)
# Returns: (mean_ratio: float, diagnostics: Dict[int, Dict])

Command-Line Interface

Resonance Calculator Demo

python optimizers/resonance_scheduler.py \
  --hidden-dim 3072 \
  --num-layers 32 \
  --attention-heads 32 \
  --target-sigma-ratio 2.0 \
  --output-json calc_diag.json

# Output:
# === Resonance Frequency Calculator ===
# {
#   "hidden_dim": 3072,
#   "resonance_frequency": 0.0456,
#   ...
# }
# === Simulated Training Run (1000 steps) ===
# ...

Full Training Integration Example

python optimizers/training_integration_example.py \
  --model phi3 \
  --hidden-dim 3072 \
  --num-layers 32 \
  --attention-heads 32 \
  --batch-size 32 \
  --num-steps 1000 \
  --target-sigma-ratio 2.0 \
  --enable-resonance-tuning \
  --compute-spectral \
  --output-metrics metrics.json \
  --log-level INFO

# Output:
# INFO - Training configuration: {...}
# INFO - Model initialized: 7,243,632,640 parameters
# INFO - Optimal learning rate from resonance: 5.23e-05
# INFO - Starting training for 1000 steps...
# INFO - Step 100: loss=1.2345, lr=5.23e-05, phase=stable, σ₁/σ₂=2.15
# ...

Theory

Resonance Frequency

The natural resonance frequency models the characteristic timescale at which weight matrices evolve during training:

ω₀ = √(κ/m)

Where:

• κ = spectral condition number (λ_max / λ_min)
• m = effective mass (related to layer dimensions)

For a transformer weight matrix W ∈ ℝ^(d_out × d_in):

• λ_max ≈ √(d_in) + √(d_out) (due to initialization + gradient accumulation)
• λ_min ≈ 1 (conservative estimate after training)
• κ typically ranges from 10 to 1000 depending on architecture

Optimal Learning Rate

The predicted optimal learning rate is scaled inversely by ω₀:

lr_opt = base_scale / ω₀ × damping_factor

Where damping_factor = log₁₀(num_training_steps) / 5 accounts for training duration.

Spectral Ratio Monitoring

During training, we track σ₁/σ₂ (ratio of largest to 2nd largest singular values) per layer. When this ratio drifts from the target:

• Too high (> 1.5 × target): Indicates rank collapse, reduce LR
• Too low (< 0.7 × target): Indicates excessive noise, increase LR
• Nominal (±15% of target): Maintain current LR

Cliff Detection

When loss increases suddenly (> 20% in one step), the scheduler:

1. Cuts LR by 50%
2. Enters RECOVERY phase
3. Gradually restores LR while monitoring spectral stability

Integration Guide

Step 1: Compute Optimal Learning Rate

Before training, compute the optimal learning rate:

from optimizers import ResonanceFrequencyCalculator

calc = ResonanceFrequencyCalculator(
    hidden_dim=model.config.hidden_size,
    num_layers=model.config.num_hidden_layers,
    attention_heads=model.config.num_attention_heads,
)

optimal_lr = calc.compute_optimal_learning_rate(
    num_training_steps=len(train_loader) * num_epochs
)

Step 2: Initialize Optimizer with Predicted LR

import torch.optim as optim

optimizer = optim.AdamW(
    model.parameters(),
    lr=optimal_lr,
    weight_decay=0.01,  # Recommended
    betas=(0.9, 0.999),
)

Step 3: Create Scheduler

from optimizers import AdaptiveResonanceLRScheduler

scheduler = AdaptiveResonanceLRScheduler(
    optimizer=optimizer,
    base_lr=optimal_lr,
    target_sigma_ratio=2.0,
    enable_resonance_tuning=True,
    warmup_steps=int(0.1 * len(train_loader)),  # 10% of first epoch
)

Step 4: Training Loop

for epoch in range(num_epochs):
    for step, batch in enumerate(train_loader):
        # Forward + backward
        logits = model(**batch)
        loss = criterion(logits, batch["labels"])
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        # Compute spectral metrics (optional but recommended)
        sigma_ratio, layer_diags = compute_spectral_ratio(model)

        # Update scheduler
        metrics = scheduler.step(
            loss=loss.item(),
            spectral_ratios=[sigma_ratio] * model.config.num_hidden_layers,
            spectral_diagnostics=layer_diags,
        )

        # Log if needed
        if (step + 1) % 100 == 0:
            print(f"Epoch {epoch}, Step {step}: "
                  f"loss={loss:.4f}, lr={scheduler.current_lr:.2e}, "
                  f"phase={metrics.phase.value}")

        # Optional: Log to W&B
        # wandb.log(metrics.to_dict(), step=global_step)

# Export final metrics
scheduler.export_metrics("training_metrics.json")

Performance Tuning

Spectral Ratio Computation

Computing SVD for every step is expensive. Strategies to reduce cost:

1. Skip computation in early steps:

   if step > 100 and step % 10 == 0:
       sigma_ratio, _ = compute_spectral_ratio(model)

2. Compute on subset of layers:

   # Only analyze FFN layers
   sigma_ratios = []
   for name, param in model.named_parameters():
       if "dense" in name:
           sigma_ratios.append(compute_spectral_ratio(model, ...))

3. Batch-wise approximation (cheaper):

   # Estimate from weight matrix norms
   sigma_approx = torch.norm(param, 2) / torch.norm(param, 2)

Target Sigma Ratio

The target σ₁/σ₂ ratio can be tuned based on model architecture:

Logging & Monitoring

JSON Export Format

[
  {
    "timestamp": 100,
    "step": 100,
    "loss": 2.3456,
    "current_lr": 5.23e-05,
    "mean_sigma_ratio": 2.15,
    "min_sigma_ratio": 1.8,
    "max_sigma_ratio": 2.5,
    "sigma_ratio_drift": 0.15,
    "phase": "stable",
    "cliff_triggered": false,
    "lr_adjustment_factor": 1.0,
    "spectral_layer_diagnostics": {
      "0": {"sigma_1": 45.2, "sigma_2": 21.0, "ratio": 2.15, ...},
      "1": {"sigma_1": 43.1, "sigma_2": 20.5, "ratio": 2.10, ...},
      ...
    }
  },
  ...
]

Visualization

Example Python script to plot training dynamics:

import json
import matplotlib.pyplot as plt

with open("training_metrics.json") as f:
    metrics = json.load(f)

steps = [m["step"] for m in metrics]
losses = [m["loss"] for m in metrics]
lrs = [m["current_lr"] for m in metrics]
ratios = [m["mean_sigma_ratio"] for m in metrics]

fig, axes = plt.subplots(2, 2, figsize=(12, 8))
axes[0, 0].plot(steps, losses)
axes[0, 0].set_ylabel("Loss")
axes[0, 1].plot(steps, lrs)
axes[0, 1].set_ylabel("Learning Rate")
axes[1, 0].plot(steps, ratios)
axes[1, 0].axhline(y=2.0, color="r", linestyle="--", label="Target")
axes[1, 0].set_ylabel("σ₁/σ₂ Ratio")
axes[1, 1].hist([m["phase"] for m in metrics], bins=10)
axes[1, 1].set_ylabel("Phase Distribution")

plt.tight_layout()
plt.savefig("training_analysis.png")

Troubleshooting

Issue: SVD Computation Too Slow

Solution: Reduce frequency or use batch approximation:

# Compute every 10 steps instead of every step
if step % 10 == 0:
    sigma_ratio, _ = compute_spectral_ratio(model)
else:
    sigma_ratio = scheduler.target_sigma_ratio  # Use default

Issue: Learning Rate Not Changing

Solution: Ensure enable_resonance_tuning=True and check warmup:

# Debug: print scheduler state
print(scheduler.phase, scheduler.step_count, scheduler.warmup_steps)
# If step_count < warmup_steps, LR is ramping; wait or reduce warmup_steps

Issue: Spectral Ratio Always NaN

Solution: Check model contains weight matrices:

# Verify model has parameters with dim >= 2
for name, param in model.named_parameters():
    if param.dim() >= 2:
        print(f"{name}: {param.shape}")

References

• Golub, G. H., & Van Loan, C. F. (1996). Matrix computations (3rd ed.)
• Smith, S. L., & Topin, N. (2019). "Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates"
• Buley, T. et al. (2026). "Resonance-Aware Learning Rate Tuning for Distributed Inference"

License

MPL-2.0 (Multi-Platform License)

Author

Taylor Buley (taylorbuley@gmail.com)