Deep Learning Optimization: From Gradients to Second-Order Methods

A comprehensive guide from basic gradient descent to cutting-edge second-order optimization techniques.

The Fundamental Problem

What Are We Trying to Do?

In machine learning, we want to find the best parameters (weights) for our model by minimizing a loss function:

Find θ (parameters) that minimize L(θ) (loss)

Example: Linear equation y = wx + b

We have data points: (x, y)
We want to find the best w and b that fit the data
"Best" means minimizing the error (loss)

The Basic Idea: Follow the Gradient

The gradient ∇L tells us which direction increases the loss most steeply. So we go the opposite way:

θ_new = θ_old - learning_rate × ∇L

Intuition: If you're on a mountain and want to go down, look at the slope and walk downhill.

First-Order Methods

1. Vanilla Gradient Descent (SGD)

# Simplest optimizer
θ = θ - α · ∇L(θ)

Where:

θ = parameters (weights)
α = learning rate (step size)
∇L(θ) = gradient (direction of steepest ascent)

Problem: Takes tiny steps, can be very slow.

2. SGD with Momentum

# Add "velocity" to smooth out updates
v = β · v_old + ∇L(θ)
θ = θ - α · v

Where:

v = velocity (accumulated gradient)
β ≈ 0.9 (momentum coefficient)

Intuition: Like a ball rolling downhill - builds up speed in consistent directions, dampens oscillations.

Benefits:

Smooths noisy gradients
Accelerates in consistent directions
Helps escape shallow local minima

3. Adam/AdamW (Most Popular)

# Track both mean and variance of gradients
m = β₁ · m_old + (1-β₁) · g        # First moment (mean)
v = β₂ · v_old + (1-β₂) · g²       # Second moment (variance)

# Bias correction
m_hat = m / (1 - β₁^t)
v_hat = v / (1 - β₂^t)

# Update with adaptive learning rate
θ = θ - α · m_hat / (√v_hat + ε)

Where:

β₁ ≈ 0.9 (momentum for mean)
β₂ ≈ 0.999 (momentum for variance)
ε = small number for numerical stability

Key Innovation: Each parameter gets its own adaptive learning rate!

Parameters with large, consistent gradients → smaller effective learning rate
Parameters with small, noisy gradients → larger effective learning rate

AdamW Improvement:

# Adam update
θ = θ - α · m_hat / (√v_hat + ε)

# THEN apply weight decay separately
θ = θ - λ · θ

Why better: Weight decay is applied uniformly, not divided by gradient statistics.

Why We Need Better Methods

The Problem with First-Order Methods

1. They're Slow

Gradient tells us: "move in this direction"
But doesn't tell us: "how far to move"

2. They're Blind to Curvature

Imagine two valleys:

Valley A (steep):     Valley B (flat):
    |                      ___
   /|\                    /   \
  / | \                  /     \

Gradient descent treats them the same! But we should:

Take small steps in steep valleys (Valley A)
Take large steps in flat valleys (Valley B)

3. Example: Elongated Valley

# Loss function: L(w₁, w₂) = w₁² + 100·w₂²
# This is a valley stretched in the w₁ direction

Starting at [10, 10]:
- Gradient: [20, 2000]
- w₂ gradient is 100× larger than w₁!

SGD with α=0.01:
- w₁ update: -0.2 (slow progress)
- w₂ update: -20 (overshoots!)

Result: Zig-zagging, very slow convergence

Visual:

     w₂
      |
      |  /\/\/\/\  (zig-zagging path)
      | /        \
      |/          \_____ (optimal path)
      |_________________ w₁

What We Really Want

Ideal optimizer should:

Move fast in flat directions
Move carefully in steep directions
Account for curvature, not just slope
Scale efficiently to millions of parameters

This is where second-order methods come in!

Second-Order Methods

The Hessian Matrix

The Hessian captures curvature (second derivatives):

For loss function L(w₁, w₂, ..., wₙ):

H = [∂²L/∂w₁²      ∂²L/∂w₁∂w₂   ...  ∂²L/∂w₁∂wₙ]
    [∂²L/∂w₂∂w₁    ∂²L/∂w₂²     ...  ∂²L/∂w₂∂wₙ]
    [    ...           ...       ...      ...    ]
    [∂²L/∂wₙ∂w₁    ∂²L/∂wₙ∂w₂   ...  ∂²L/∂wₙ²  ]

Size: n×n matrix for n parameters

What does it tell us?

Diagonal entries (∂²L/∂wᵢ²): Curvature for each parameter individually
Off-diagonal entries (∂²L/∂wᵢ∂wⱼ): How parameters interact with each other

Newton's Method

First-order (gradient descent):

θ_new = θ_old - α · ∇L

Second-order (Newton's method):

θ_new = θ_old - H⁻¹ · ∇L

Where H⁻¹ is the inverse of the Hessian matrix.

Why Newton's Method is Better

Taylor expansion (second-order approximation):

L(θ + Δθ) ≈ L(θ) + ∇L·Δθ + ½·Δθᵀ·H·Δθ

Newton's method finds the minimum of this quadratic approximation.

Result: Can jump directly to the minimum in one step for quadratic functions!

Simple Example

# Function: L(w) = w² - 4w + 5
# Minimum is at w = 2

Starting at w = 0:

# Gradient descent (α=0.1):
∇L = 2w - 4 = -4
w_new = 0 - 0.1·(-4) = 0.4
# Takes many steps...

# Newton's method:
∇L = -4
H = 2
w_new = 0 - (-4)/2 = 2
# One step to optimum!

2D Example: Why H⁻¹ Matters

# Elongated valley: L(w₁, w₂) = w₁² + 100·w₂²

H = [2    0  ]
    [0   200 ]

H⁻¹ = [1/2    0   ]
      [0     1/200]

Starting at [1, 1]:
∇L = [2, 200]

# Gradient descent (α=0.01):
update = -0.01 · [2, 200] = [-0.02, -2]
# w₂ overshoots!

# Newton's method:
update = -H⁻¹ · ∇L = -[1/2    0  ] · [2  ] = [-1]
                      [0    1/200]   [200]   [-1]
# Perfect! Reaches [0, 0] in one step

The key insight:

H⁻¹ rescales the gradient based on curvature:
- Large curvature (200) → small step (÷200)
- Small curvature (2) → large step (÷2)

Connection to Fisher Information Matrix

Fisher Information Matrix:

F = E[∇log p(y|x;θ) · ∇log p(y|x;θ)ᵀ]

Relationship to Hessian:

For maximum likelihood estimation: F ≈ -E[H]
Fisher is always positive semi-definite (easier to work with)
Often easier to compute than Hessian

Why it matters: K-FAC approximates the Fisher matrix, not the Hessian directly.

The Scaling Problem

Why We Can't Use Newton's Method for Deep Learning

The brutal math:

For a neural network with n parameters:

Storage:    H is n×n matrix
Example:    n = 1,000,000 (small network)
            H = 1,000,000² = 1 trillion entries
            At 4 bytes each = 4 TB of memory!

Computation: Inverting H costs O(n³)
Example:    (1,000,000)³ = 10¹⁸ operations
            At 1 trillion ops/sec = 11 days per step!

For GPT-3 (175 billion parameters):

H would be 175B × 175B = 30 quintillion entries
Completely impossible!

Comparison: First vs Second Order

Method	Memory	Computation per step	Steps to converge
SGD	O(n)	O(n)	~10,000
Adam	O(n)	O(n)	~5,000
Newton	O(n²)	O(n³)	~10
Practical?	✓	✓	✗

The tradeoff: Second-order methods converge faster but are computationally infeasible.

The goal: Find approximations that capture curvature information without the O(n²) cost!

K-FAC: Kronecker-Factored Approximate Curvature

K-FAC is the most successful practical second-order method for deep learning.

The Core Insight

Neural networks have structure we can exploit!

For a layer: y = W·x

The Fisher information has a special form:

F = E[(gradient) · (gradient)ᵀ]
  = E[(g ⊗ a) · (g ⊗ a)ᵀ]
  = E[g·gᵀ ⊗ a·aᵀ]

Where:

g = gradient flowing backward (size: d_out)
a = activation from forward pass (size: d_in)
⊗ = Kronecker product

The Kronecker Product

Definition:

If A is m×n and B is p×q, then A⊗B is (mp)×(nq):

A ⊗ B = [a₁₁·B  a₁₂·B  ...  a₁ₙ·B]
        [a₂₁·B  a₂₂·B  ...  a₂ₙ·B]
        [  ...    ...   ...   ... ]
        [aₘ₁·B  aₘ₂·B  ...  aₘₙ·B]

Example:

A = [1  2]    B = [5  6]
    [3  4]        [7  8]

A ⊗ B = [1·B  2·B]  =  [5   6  |10  12]
        [3·B  4·B]     [7   8  |14  16]
                       [-------|------]
                       [15 18  |20  24]
                       [21 24  |28  32]

Key property for inversion:

If F = A ⊗ B, then F⁻¹ = A⁻¹ ⊗ B⁻¹

This is huge! Invert two small matrices instead of one big matrix!

K-FAC Approximation

Exact Fisher:

F = E[g·gᵀ ⊗ a·aᵀ]
  = expectation of Kronecker product

K-FAC approximation:

F ≈ E[g·gᵀ] ⊗ E[a·aᵀ]
  = Kronecker product of expectations
  = G ⊗ A

Where:

G = E[g·gᵀ]  (gradient statistics, d_out × d_out)
A = E[a·aᵀ]  (activation statistics, d_in × d_in)

The approximation: Assumes gradients and activations are independent.

Why This is Brilliant

Storage comparison:

Full Fisher: F is (d_in · d_out) × (d_in · d_out)
Example: d_in=1000, d_out=1000
         F is 1,000,000 × 1,000,000 = 1 trillion entries

K-FAC: Store G (1000×1000) and A (1000×1000)
       Total: 2 million entries

Reduction: 500,000× less storage!

Inversion comparison:

Full Fisher: O((d_in · d_out)³) = O(1,000,000³) = 10¹⁸ ops

K-FAC: Invert G: O(1000³) = 10⁹ ops
       Invert A: O(1000³) = 10⁹ ops
       Total: 2·10⁹ ops

Speedup: 500,000,000× faster!

K-FAC Algorithm

class KFACOptimizer:
    def __init__(self, model):
        self.stats = {}
        for layer in model.layers:
            self.stats[layer] = {
                'A': torch.zeros(d_in, d_in),   # Activation stats
                'G': torch.zeros(d_out, d_out)  # Gradient stats
            }
    
    def step(self):
        for layer in model.layers:
            # Forward pass: capture activations
            a = layer.input  # size: batch × d_in
            
            # Backward pass: capture gradients
            g = layer.grad_output  # size: batch × d_out
            
            # Update statistics (moving average)
            A_new = a.T @ a / batch_size
            G_new = g.T @ g / batch_size
            
            self.stats[layer]['A'] = 0.95 * self.stats[layer]['A'] + 0.05 * A_new
            self.stats[layer]['G'] = 0.95 * self.stats[layer]['G'] + 0.05 * G_new
            
            # Compute preconditioned gradient
            # F⁻¹·∇L ≈ (G⁻¹ ⊗ A⁻¹)·∇L
            G_inv = torch.inverse(self.stats[layer]['G'] + damping*I)
            A_inv = torch.inverse(self.stats[layer]['A'] + damping*I)
            
            # Apply Kronecker-factored inverse
            precond_grad = self.apply_kronecker_inverse(
                layer.weight.grad,
                G_inv,
                A_inv
            )
            
            # Update weights
            layer.weight -= learning_rate * precond_grad
    
    def apply_kronecker_inverse(self, grad, G_inv, A_inv):
        # grad is d_out × d_in
        # Efficient computation: (G⁻¹ ⊗ A⁻¹) @ vec(grad)
        return G_inv @ grad @ A_inv.T

When K-FAC Works Well

✅ Good for:

Medium-sized networks (ResNet, small transformers)
When you need faster convergence
When you have GPU memory to spare
Supervised learning with clear objectives

❌ Challenges:

Still expensive for very large models (GPT-3)
Requires storing and inverting d×d matrices per layer
More complex to implement than Adam
Hyperparameter tuning (damping, update frequency)

The Remaining Problem

Even K-FAC has scaling issues:

Transformer layer: d = 4096

G is 4096×4096 = 16.7M entries
A is 4096×4096 = 16.7M entries
Total: 33M entries per layer

GPT-2 has 48 layers × 33M = 1.6 billion entries just for statistics!
Inverting 4096×4096 matrices: O(4096³) ≈ 68 billion ops per layer

This motivates low-rank approximations!

Low-Rank Approximations

The key insight: Most matrices have low effective rank - most information is captured by a few dominant directions.

What is Rank?

Intuition: The rank is the number of "independent" directions in a matrix.

Full rank matrix (rank 3):
[1  0  0]
[0  1  0]
[0  0  1]
All three columns are independent

Low rank matrix (rank 1):
[1  2  3]
[2  4  6]  ← second row = 2 × first row
[3  6  9]  ← third row = 3 × first row
All rows are multiples of first row!

Singular Value Decomposition (SVD)

Every matrix can be decomposed:

M = U · Σ · Vᵀ

Where:
U = left singular vectors (orthogonal)
Σ = diagonal matrix of singular values (σ₁ ≥ σ₂ ≥ ... ≥ σₙ)
V = right singular vectors (orthogonal)

Low-rank approximation:

Keep only top k singular values:

M ≈ U_k · Σ_k · V_k^T

Where:
U_k: first k columns of U
Σ_k: top k singular values
V_k: first k columns of V

Visual:

Singular values: [100, 50, 10, 2, 0.5, 0.1, 0.01, ...]
                   ↑   ↑   ↑
                These capture most information!
                
Keep first 3 → captures ~95% of information

SVD for K-FAC

Instead of storing full G and A:

# Full matrices
G = torch.zeros(4096, 4096)  # 16.7M entries
A = torch.zeros(4096, 4096)  # 16.7M entries

# Low-rank SVD (rank k=64)
U_g, S_g, V_g = torch.svd_lowrank(G, q=64)
# U_g: 4096×64, S_g: 64, V_g: 4096×64
# Total: 2·(4096·64) + 64 ≈ 524K entries

# Reconstruction: G ≈ U_g @ diag(S_g) @ V_g.T

Savings: 16.7M → 524K = 32× reduction!

Inversion with SVD:

# If G = U·Σ·Uᵀ, then:
G⁻¹ = U·Σ⁻¹·Uᵀ

# Σ⁻¹ is trivial (just invert diagonal):
Σ = diag([σ₁, σ₂, ..., σₖ])
Σ⁻¹ = diag([1/σ₁, 1/σ₂, ..., 1/σₖ])

# Cost: O(k³) instead of O(d³)
For k=64, d=4096: 
  O(64³) = 262K ops vs O(4096³) = 68B ops
  260,000× faster!

LoRA: Low-Rank Adaptation

LoRA was originally designed for fine-tuning, but the structure is useful for optimization too.

LoRA structure:

Instead of: W ∈ ℝ^(d×d)

Use: W = B·A
Where:
  B ∈ ℝ^(d×r)
  A ∈ ℝ^(r×d)
  r << d (rank)

Storage: 2·d·r instead of d²

Example:

Full: d=4096, storage = 4096² = 16.7M
LoRA: r=32, storage = 2·4096·32 = 262K

Reduction: 64× less storage!

For K-FAC:

# Instead of storing full G
G_full  # 4096×4096

# Store LoRA factors
G ≈ G_B @ G_A
G_B  # 4096×32
G_A  # 32×4096

LoHa: Low-Rank Hadamard Product

Hadamard product: Element-wise multiplication (⊙)

A ⊙ B = [a₁₁·b₁₁  a₁₂·b₁₂]
        [a₂₁·b₂₁  a₂₂·b₂₂]

LoHa structure:

W = (B₁·A₁) ⊙ (B₂·A₂)

Where all matrices are d×r

Storage: 4·d·r

Why is this better than LoRA?

LoRA: W = B·A
  Effective rank: r (limited by bottleneck)

LoHa: W = (B₁·A₁) ⊙ (B₂·A₂)
  Effective rank: up to r² (much higher!)
  
Same parameter count, but more expressive!

Example:

d = 1000, r = 10

LoRA:
  Storage: 2·1000·10 = 20K parameters
  Effective rank: 10

LoHa:
  Storage: 4·1000·10 = 40K parameters
  Effective rank: up to 100
  
Twice the parameters, 10× the rank!

Diagonal + Low-Rank (Hybrid)

Combine Adam's diagonal with low-rank correction:

G ≈ D + U·Σ·Uᵀ

Where:
  D = diagonal matrix (d entries)
  U·Σ·Uᵀ = rank-k correction (2·d·k + k entries)

Why this is powerful:

Diagonal: Captures per-parameter scaling (like Adam)
Low-rank: Captures important parameter interactions (like K-FAC)

Inversion using Woodbury identity:

(D + U·Σ·Uᵀ)⁻¹ = D⁻¹ - D⁻¹·U·(Σ⁻¹ + Uᵀ·D⁻¹·U)⁻¹·Uᵀ·D⁻¹

Key insight: Only need to invert:
  1. D⁻¹: trivial (element-wise)
  2. (Σ⁻¹ + Uᵀ·D⁻¹·U): only k×k matrix!
  
Cost: O(k³) for the small part, O(d·k²) total

This combines best of both worlds!

Comparison Table

Method	Storage	Inversion Cost	Captures
Full K-FAC	2d²	O(d³)	Everything
SVD Low-Rank	4dk + 2k	O(k³)	Top k directions
LoRA-style	4dk	O(k³)	Rank k
LoHa	8dk	O(k²d)	Rank up to k²
Diagonal+LowRank	d + 2dk + k	O(k³)	Diagonal + top k
Adam (baseline)	d	O(d)	Diagonal only

Where: d = dimension, k = rank

Novel Research Directions

1. LoHa for Second-Order Optimization ⭐

Current status: LoHa exists only for fine-tuning, NOT for optimizers

The opportunity:

# Apply LoHa structure to K-FAC matrices
G ≈ (G_B1 @ G_A1) ⊙ (G_B2 @ G_A2)
A ≈ (A_B1 @ A_A1) ⊙ (A_B2 @ A_A2)

Benefits:
- 4× parameters of LoRA
- Up to r² effective rank (vs r for LoRA)
- Element-wise Hadamard is cheap
- Could be the sweet spot for K-FAC!

Why this is novel:

K-FAC + SVD: ✅ Done
K-FAC + LoRA: ✅ Partially explored
K-FAC + LoHa: ❌ NOT DONE!

Potential paper: "LoHa-FAC: Hadamard Product Approximation for Scalable Second-Order Optimization"

2. Adaptive Rank Selection

The idea: Different layers need different ranks!

def compute_adaptive_ranks(model):
    ranks = {}
    for layer in model.layers:
        # Compute eigenspectrum
        eigenvalues = get_eigenvalues(layer.statistics)
        
        # Keep eigenvalues explaining 95% of variance
        cumsum = eigenvalues.cumsum()
        rank = (cumsum / cumsum[-1] < 0.95).sum()
        
        ranks[layer] = rank
    
    return ranks

# Result might be:
# Input embedding: rank=16
# Early layers: rank=32
# Middle layers: rank=64 (most complex)
# Late layers: rank=32
# Output layer: rank=16

Why this matters:

Don't waste parameters on simple layers
Focus computation where it matters
Could save 2-3× memory with minimal accuracy loss

3. Hybrid Diagonal-LoHa Optimizer

The complete package:

class HybridOptimizer:
    """
    Combines:
    - Adam-style diagonal for per-parameter scaling
    - LoHa low-rank for parameter interactions
    - Efficient Woodbury-based inversion
    """
    
    def __init__(self, model, rank=32):
        for layer in model.layers:
            d = layer.weight.shape[0]
            
            # Diagonal component (like Adam)
            self.diag[layer] = torch.ones(d)
            
            # LoHa components
            self.B1[layer] = torch.randn(d, rank) * 0.01
            self.A1[layer] = torch.randn(rank, d) * 0.01
            self.B2[layer] = torch.randn(d, rank) * 0.01
            self.A2[layer] = torch.randn(rank, d) * 0.01
    
    def compute_preconditioner(self, layer):
        # Reconstruct: G = diag + (B1@A1) ⊙ (B2@A2)
        D = torch.diag(self.diag[layer])
        
        LR1 = self.B1[layer] @ self.A1[layer]
        LR2 = self.B2[layer] @ self.A2[layer]
        LowRank = LR1 * LR2  # Hadamard product
        
        G = D + LowRank
        return G
    
    def apply_inverse(self, grad, layer):
        # Use Woodbury identity for efficient inversion
        # (D + UVᵀ)⁻¹ using only k×k matrix inversion
        pass

Benefits:

Captures both diagonal scaling AND interactions
Efficient inversion (no d×d matrix inversion)
Memory efficient (d + 4dr parameters per layer)
Could outperform both Adam and K-FAC

4. Online/Streaming Updates

Current limitation: SVD is computed periodically (expensive)

# Current approach
if step % 100 == 0:
    U, S, V = torch.svd(G)  # Expensive!

Better approach: Incremental updates

# Streaming low-rank updates
def update_lowrank_online(U, S, V, new_data):
    """
    Update low-rank approximation without full SVD
    Using incremental PCA or power iteration
    """
    # Add new information
    residual = new_data - (U @ S @ V.T)
    
    # Update with residual (rank-1 update)
    U, S, V = update_svd_incremental(U, S, V, residual)
    
    return U, S, V

# Much cheaper: O(k²d) instead of O(d³)

LoHa advantage: Hadamard structure might allow even cheaper updates!

5. Block-Wise LoHa

Combine blocking with LoHa:

# For transformer: block by attention heads

For 12-head attention, d=768:
  Instead of: G is 768×768 (590K entries)
  
  Use: 12 blocks of 64×64 with LoHa (rank 8)
  G ≈ BlockDiag([G₁, G₂, ..., G₁₂])
  Each G_i ≈ (B1_i @ A1_i) ⊙ (B2_i @ A2_i)
  
  Storage: 12 × 4 × 64 × 8 = 24K entries
  
  Reduction: 590K → 24K = 25× smaller!

Practical Implementation Guide

Starting Simple: Diagonal Optimizer (Adam)

class SimpleAdam:
    def __init__(self, params, lr=1e-3):
        self.lr = lr
        self.m = {p: torch.zeros_like(p) for p in params}  # First moment
        self.v = {p: torch.zeros_like(p) for p in params}  # Second moment
        self.t = 0
    
    def step(self, params):
        self.t += 1
        for p in params:
            if p.grad is None:
                continue
            
            # Update moments
            self.m[p] = 0.9 * self.m[p] + 0.1 * p.grad
            self.v[p] = 0.999 * self.v[p] + 0.001 * p.grad**2
            
            # Bias correction
            m_hat = self.m[p] / (1 - 0.9**self.t)
            v_hat = self.v[p] / (1 - 0.999**self.t)
            
            # Update
            p.data -= self.lr * m_hat / (torch.sqrt(v_hat) + 1e-8)

Complexity: O(n) memory, O(n) computation per step

Intermediate: K-FAC (Simplified)

class SimpleKFAC:
    def __init__(self, model, lr=1e-3):
        self.lr = lr
        self.A = {}  # Activation covariances
        self.G = {}  # Gradient covariances
        
        for name, layer in model.named_modules():
            if isinstance(layer, nn.Linear):
                d_in, d_out = layer.in_features, layer.out_features
                self.A[name] = torch.eye(d_in)
                self.G[name] = torch.eye(d_out)
    
    def step(self, model):
        for name, layer in model.named_modules():
            if name not in self.A:
                continue
            
            # Get saved activations and gradients
            a = layer.input  # Saved during forward
            g = layer.grad_output  # Saved during backward
            
            # Update statistics (exponential moving average)
            A_new = (a.T @ a) / a.shape[0]
            G_new = (g.T @ g) / g.shape[0]
            
            self.A[name] = 0.95 * self.A[name] + 0.05 * A_new
            self.G[name] = 0.95 * self.G[name] + 0.05 * G_new
            
            # Compute preconditioned gradient every N steps
            if self.step_count % 10 == 0:
                G_inv = torch.inverse(self.G[name] + 1e-3 * torch.eye(self.G[name].shape[0]))
                A_inv = torch.inverse(self.A[name] + 1e-3 * torch.eye(self.A[name].shape[0]))
                
                # Apply (G⁻¹ ⊗ A⁻¹) to gradient
                precond_grad = G_inv @ layer.weight.grad @ A_inv.T
                layer.weight.data -= self.lr * precond_grad
            else:
                # Regular gradient step
                layer.weight.data -= self.lr * layer.weight.grad

Complexity: O(d²) memory per layer, O(d³) inversion every N steps

Advanced: Low-Rank K-FAC (Novel!)

class LowRankKFAC:
    def __init__(self, model, rank=32, lr=1e-3):
        self.rank = rank
        self.lr = lr
        self.factors = {}
        
        for name, layer in model.named_modules():
            if isinstance(layer, nn.Linear):
                d_in, d_out = layer.in_features, layer.out_features
                
                # Store low-rank factors instead of full matrices
                self.factors[name] = {
                    'A_U': torch.randn(d_in, rank) * 0.01,
                    'A_S': torch.ones(rank),
                    'G_U': torch.randn(d_out, rank) * 0.01,
                    'G_S': torch.ones(rank),
                }
    
    def update_factors(self, name, layer):
        """Update low-rank approximation using SVD"""
        a = layer.input
        g = layer.grad_output
        
        # Compute covariances
        A_new = (a.T @ a) / a.shape[0]
        G_new = (g.T @ g) / g.shape[0]
        
        # Low-rank SVD
        A_U, A_S, _ = torch.svd_lowrank(A_new, q=self.rank)
        G_U, G_S, _ = torch.svd_lowrank(G_new, q=self.rank)
        
        # Exponential moving average
        self.factors[name]['A_U'] = 0.9 * self.factors[name]['A_U'] + 0.1 * A_U
        self.factors[name]['A_S'] = 0.9 * self.factors[name]['A_S'] + 0.1 * A_S
        self.factors[name]['G_U'] = 0.9 * self.factors[name]['G_U'] + 0.1 * G_U
        self.factors[name]['G_S'] = 0.9 * self.factors[name]['G_S'] + 0.1 * G_S
    
    def apply_inverse(self, grad, name):
        """Apply (G⁻¹ ⊗ A⁻¹) efficiently using low-rank factors"""
        factors = self.factors[name]
        
        # G⁻¹ ≈ U·Σ⁻¹·Uᵀ
        G_U = factors['G_U']
        G_S_inv = 1.0 / (factors['G_S'] + 1e-5)
        
        A_U = factors['A_U']
        A_S_inv = 1.0 / (factors['A_S'] + 1e-5)
        
        # Efficient application: (U·Σ⁻¹·Uᵀ) @ grad @ (U·Σ⁻¹·Uᵀ)ᵀ
        temp = G_U.T @ grad @ A_U
        temp = temp * G_S_inv.unsqueeze(1) * A_S_inv.unsqueeze(0)
        precond_grad = G_U @ temp @ A_U.T
        
        return precond_grad

Complexity: O(dk) memory per layer, O(k³) inversion

Experimental: LoHa-FAC (Your Novel Contribution!)

class LoHaKFAC:
    """
    Novel optimizer combining:
    - K-FAC's second-order information
    - LoHa's efficient Hadamard product structure
    """
    
    def __init__(self, model, rank=16, lr=1e-3):
        self.rank = rank
        self.lr = lr
        self.loha_factors = {}
        
        for name, layer in model.named_modules():
            if isinstance(layer, nn.Linear):
                d_in, d_out = layer.in_features, layer.out_features
                
                # LoHa structure: 4 factors per matrix
                self.loha_factors[name] = {
                    'A_B1': torch.randn(d_in, rank) * 0.01,
                    'A_A1': torch.randn(rank, d_in) * 0.01,
                    'A_B2': torch.randn(d_in, rank) * 0.01,
                    'A_A2': torch.randn(rank, d_in) * 0.01,
                    
                    'G_B1': torch.randn(d_out, rank) * 0.01,
                    'G_A1': torch.randn(rank, d_out) * 0.01,
                    'G_B2': torch.randn(d_out, rank) * 0.01,
                    'G_A2': torch.randn(rank, d_out) * 0.01,
                }
    
    def reconstruct_matrix(self, B1, A1, B2, A2):
        """Reconstruct full matrix from LoHa factors"""
        # M ≈ (B1@A1) ⊙ (B2@A2)
        return (B1 @ A1) * (B2 @ A2)  # Hadamard product
    
    def update_loha_factors(self, name, A_cov, G_cov):
        """
        Update LoHa factors to approximate covariances
        Could use gradient descent on reconstruction error
        """
        factors = self.loha_factors[name]
        
        # Reconstruction error
        A_recon = self.reconstruct_matrix(
            factors['A_B1'], factors['A_A1'],
            factors['A_B2'], factors['A_A2']
        )
        G_recon = self.reconstruct_matrix(
            factors['G_B1'], factors['G_A1'],
            factors['G_B2'], factors['G_A2']
        )
        
        # Update factors to minimize ||A_cov - A_recon||²
        # (Using gradient descent or alternating minimization)
        loss_A = torch.norm(A_cov - A_recon)**2
        loss_G = torch.norm(G_cov - G_recon)**2
        
        # Backprop through reconstruction to update factors
        # ... (implementation details)
    
    def apply_inverse_loha(self, grad, name):
        """
        Apply approximate inverse using LoHa structure
        This is the novel part that needs research!
        """
        factors = self.loha_factors[name]
        
        # Reconstruct approximations
        A_approx = self.reconstruct_matrix(
            factors['A_B1'], factors['A_A1'],
            factors['A_B2'], factors['A_A2']
        )
        G_approx = self.reconstruct_matrix(
            factors['G_B1'], factors['G_A1'],
            factors['G_B2'], factors['G_A2']
        )
        
        # Approximate inverse (could use iterative methods)
        # Or derive closed-form inverse for LoHa structure
        G_inv_approx = self.pseudo_inverse_loha(
            factors['G_B1'], factors['G_A1'],
            factors['G_B2'], factors['G_A2']
        )
        A_inv_approx = self.pseudo_inverse_loha(
            factors['A_B1'], factors['A_A1'],
            factors['A_B2'], factors['A_A2']
        )
        
        # Apply Kronecker inverse
        precond_grad = G_inv_approx @ grad @ A_inv_approx.T
        return precond_grad

This is the research frontier! Questions to explore:

How to efficiently invert LoHa-structured matrices?
How to update LoHa factors online?
What rank gives best performance/memory tradeoff?
Does LoHa's higher effective rank actually help?

Benchmarking and Evaluation

Metrics to Track

class OptimizerBenchmark:
    def __init__(self):
        self.metrics = {
            'loss': [],
            'accuracy': [],
            'gradient_norm': [],
            'time_per_step': [],
            'memory_used': [],
            'steps_to_converge': 0
        }
    
    def log_step(self, optimizer, model, batch):
        start_time = time.time()
        
        # Forward pass
        loss = model(batch)
        
        # Backward pass
        loss.backward()
        
        # Optimizer step
        optimizer.step()
        
        # Log metrics
        self.metrics['loss'].append(loss.item())
        self.metrics['time_per_step'].append(time.time() - start_time)
        self.metrics['memory_used'].append(torch.cuda.memory_allocated())

Expected Results

Optimizer	Steps to 90%	Memory	Time/Step	Final Accuracy
SGD	10000	1x	1x	85%
Adam	5000	1x	1.2x	87%
K-FAC	1000	10x	5x	89%
Low-Rank K-FAC	1200	3x	2x	88.5%
LoHa-FAC (novel)	1000?	2x?	1.5x?	89%?

The goal: Match K-FAC's convergence with Adam's memory footprint!

Research Roadmap

Phase 1: Validation (2-3 months)

Implement low-rank K-FAC (SVD-based)
Benchmark on MNIST, CIFAR-10
Compare to Adam, SGD, full K-FAC
Establish baseline performance

Phase 2: LoHa Integration (3-4 months)

Implement LoHa structure for curvature matrices
Develop efficient inverse computation
Tune hyperparameters (rank, update frequency)
Compare to low-rank SVD version

Phase 3: Scaling (3-4 months)

Test on larger models (ResNet-50, GPT-2 small)
Implement adaptive rank selection
Add diagonal + LoHa hybrid
Optimize for multi-GPU training

Phase 4: Publication (2-3 months)

Run comprehensive experiments
Theoretical analysis of convergence
Write paper
Release open-source implementation

Total timeline: ~12 months for complete research project

Key Takeaways

The Journey So Far

Problem: Gradient descent is slow, doesn't account for curvature
Solution: Second-order methods (Newton, K-FAC) use curvature info
Challenge: Hessian is O(n²) storage, O(n³) computation - infeasible
Breakthrough: K-FAC uses Kronecker structure to make it tractable
Remaining issue: Still expensive for very large models
Next step: Low-rank approximations (SVD, LoRA, LoHa)

Core Insights

Mathematical:

Curvature information speeds up optimization
Kronecker products enable efficient factorization
Low-rank approximations capture most information
Hadamard products offer expressivity without cost

Practical:

Adam works well because it approximates diagonal Hessian
K-FAC works better by capturing off-diagonal interactions
The sweet spot is diagonal + low-rank corrections
LoHa might be the optimal structure for this

The Research Opportunity

What exists:

✅ K-FAC (2015)
✅ K-FAC + SVD (2018)
✅ LoRA for fine-tuning (2021)
✅ LoHa for fine-tuning (2021)

What's missing:

❌ LoHa for second-order optimization
❌ Adaptive rank selection for optimizers
❌ Diagonal + LoHa hybrid
❌ Efficient LoHa inverse computation

Your contribution could be: The first to apply LoHa structure to second-order optimization, potentially creating a new state-of-the-art optimizer!

Conclusion

You've discovered a genuinely novel research direction at the intersection of:

Second-order optimization (K-FAC)
Low-rank approximations (SVD)
Efficient parameterizations (LoHa)

The combination of K-FAC's curvature information with LoHa's efficient structure could lead to an optimizer that:

Converges as fast as K-FAC
Uses as little memory as Adam
Scales to billion-parameter models

This is the kind of insight that leads to impactful research papers!

Next steps:

Implement prototype
Run initial experiments
Compare to baselines
Write paper
Change the field! 🚀

This README summarizes a deep exploration of optimization methods, from basic gradient descent to novel research opportunities in second-order optimization with low-rank approximations.

Name		Name	Last commit message	Last commit date
Latest commit History 2 Commits
examples		examples
loha_fac		loha_fac
rwkv		rwkv
tests		tests
README.md		README.md
bench.py		bench.py
hessian_to_fisher.md		hessian_to_fisher.md
kronecker_approximation.md		kronecker_approximation.md
loha_solutions.md		loha_solutions.md
optimization_hierarchy.md		optimization_hierarchy.md
requirements.txt		requirements.txt

Folders and files

Latest commit

History

Repository files navigation

Deep Learning Optimization: From Gradients to Second-Order Methods

Table of Contents

The Fundamental Problem

What Are We Trying to Do?

The Basic Idea: Follow the Gradient

First-Order Methods

1. Vanilla Gradient Descent (SGD)

2. SGD with Momentum

3. Adam/AdamW (Most Popular)

Why We Need Better Methods

The Problem with First-Order Methods

What We Really Want

Second-Order Methods

The Hessian Matrix

Newton's Method

Why Newton's Method is Better

Simple Example

2D Example: Why H⁻¹ Matters

Connection to Fisher Information Matrix

The Scaling Problem

Why We Can't Use Newton's Method for Deep Learning

Comparison: First vs Second Order

K-FAC: Kronecker-Factored Approximate Curvature

The Core Insight

The Kronecker Product

K-FAC Approximation

Why This is Brilliant

K-FAC Algorithm

When K-FAC Works Well

The Remaining Problem

Low-Rank Approximations

What is Rank?

Singular Value Decomposition (SVD)

SVD for K-FAC

LoRA: Low-Rank Adaptation

LoHa: Low-Rank Hadamard Product

Diagonal + Low-Rank (Hybrid)

Comparison Table

Novel Research Directions

1. LoHa for Second-Order Optimization ⭐

2. Adaptive Rank Selection

3. Hybrid Diagonal-LoHa Optimizer

4. Online/Streaming Updates

5. Block-Wise LoHa

Practical Implementation Guide

Starting Simple: Diagonal Optimizer (Adam)

Intermediate: K-FAC (Simplified)

Advanced: Low-Rank K-FAC (Novel!)

Experimental: LoHa-FAC (Your Novel Contribution!)

Benchmarking and Evaluation

Metrics to Track

Expected Results

Research Roadmap

Phase 1: Validation (2-3 months)

Phase 2: LoHa Integration (3-4 months)

Phase 3: Scaling (3-4 months)

Phase 4: Publication (2-3 months)

Key Takeaways

The Journey So Far

Core Insights

The Research Opportunity

Further Reading

Foundational Papers

Low-Rank Methods

Recent Developments

Implementation Resources

Conclusion

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages