Files
foxhunt/docs/codebase-cleanup/spectral_normalization_research.md
jgrusewski 2df1ea92e1 feat(ml): WAVE 29 DQN Codebase Cleanup & Refactoring Campaign
BREAKING CHANGES:
- Removed orphaned dqn.rs monolithic trainer (4,975 lines)
- Removed orphaned dqn_ensemble.rs module (816 lines)
- Removed orphaned tft.rs and tft_complete_int8_integration_test.rs
- TFT trainer split into modular directory structure

DQN Module Refactoring:
- Split trainers/dqn.rs into modular structure (config.rs, statistics.rs, trainer.rs)
- Fixed hyperopt 39D search space (continuous params only)
- Boolean flags (use_dueling, use_double_dqn, use_per, use_noisy_nets) are now FIXED architectural decisions
- use_distributional defaults to false (Candle BUG #36 - scatter_add gradient issues)

Clean Module Structure:
- ml/src/trainers/dqn/ directory with proper mod.rs exports
- ml/src/trainers/tft/ directory with config.rs, types.rs, model.rs, trainer.rs, tests.rs
- All P0 features validated: TD-error clamping, batch diversity, LR scheduler, priority staleness

Documentation:
- Added comprehensive docs in docs/codebase-cleanup/
- ADR-001 for DQN refactoring decisions
- Rainbow DQN component matrix and quick reference guides

Build Status: Compiles with zero errors

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
2025-11-27 23:46:13 +01:00

12 KiB
Raw Blame History

Spectral Normalization Research Report - Agent 15

Date: 2025-11-27 Research Topic: Spectral Normalization for DQN Anti-Overfitting Objective: Investigate implementation feasibility in candle framework

Executive Summary

Spectral normalization is a powerful regularization technique that constrains the Lipschitz constant of neural network layers, preventing overfitting by limiting how much the network can change outputs for small input changes. This research evaluates its implementation feasibility for the DQN architecture.

1. Candle Framework Analysis

1.1 Available Normalization Layers in candle-nn v0.9.1

From crates.io and docs.rs research:

Built-in Normalizations:

  • LayerNorm - Layer normalization (already used in DQN)
  • BatchNorm - Batch normalization
  • GroupNorm - Group normalization
  • RmsNorm - Root mean square normalization
  • SpectralNorm - NOT AVAILABLE

Key Finding: Candle does NOT have built-in spectral normalization support.

1.2 Linear Algebra Capabilities

Searched codebase for SVD and power iteration:

  • Found power iteration example in ml/src/regime/transition_matrix.rs:273
  • Found power iteration for eigenvalue in ml/src/checkpoint/enterprise_implementations.rs:308
  • No native SVD implementation found in candle

Candle Tensor Operations Available:

  • Matrix multiplication (matmul)
  • Broadcasting operations
  • Elementwise operations
  • No built-in SVD or eigendecomposition

2. Spectral Normalization Theory

2.1 Mathematical Foundation

Spectral Norm: The largest singular value of weight matrix W

σ(W) = max ||Wx||₂ / ||x||₂
       x≠0

Lipschitz Constraint:

f(x₁) - f(x₂) ≤ L||x₁ - x₂||

Where L is the Lipschitz constant. Spectral normalization ensures L ≤ 1 for each layer.

2.2 Benefits for DQN

  1. Prevents Overfitting: Limits network capacity to memorize noise
  2. Stabilizes Training: Reduces gradient explosion
  3. Improves Generalization: Forces smoother decision boundaries
  4. Compatible with Other Techniques: Works with dropout, layer norm

3. Implementation Approaches

Algorithm (Miyato et al., 2018):

pub struct SpectralNorm {
    layer: Linear,
    u: Tensor,  // Left singular vector
    v: Tensor,  // Right singular vector
    n_power_iterations: usize,
}

impl SpectralNorm {
    fn normalize_weights(&mut self) -> Result<()> {
        // 1. Power iteration to estimate spectral norm
        let w = self.layer.weight();
        let mut u = self.u.clone();
        let mut v = self.v.clone();

        for _ in 0..self.n_power_iterations {
            // v = W^T u / ||W^T u||
            v = (w.t()?.matmul(&u)?)
                .div(&(w.t()?.matmul(&u)?.sqr()?.sum_all()?.sqrt()?)?)?;

            // u = W v / ||W v||
            u = (w.matmul(&v)?)
                .div(&(w.matmul(&v)?.sqr()?.sum_all()?.sqrt()?)?)?;
        }

        // 2. Compute spectral norm: σ = u^T W v
        let sigma = u.t()?.matmul(&w)?.matmul(&v)?;

        // 3. Normalize weights: W_norm = W / σ
        let w_normalized = w.div(&sigma)?;
        self.layer.set_weight(w_normalized)?;

        // 4. Update singular vectors for next iteration
        self.u = u;
        self.v = v;

        Ok(())
    }
}

Pros:

  • Efficient (only ~1-3 power iterations needed)
  • No external dependencies required
  • Compatible with existing DQN architecture

Cons:

  • Requires custom implementation
  • Adds overhead to forward pass
  • Need to track u/v vectors in state

Approach: Use external crate like ndarray-linalg or faer

Pros:

  • Mathematically exact
  • Well-tested implementations

Cons:

  • Requires converting candle tensors to ndarray
  • Significant performance overhead
  • Complex integration with candle
  • SVD is O(n³) vs power iteration O(n²)

3.3 Hybrid: Use Existing Power Iteration Pattern

Observation: Codebase already has power iteration examples:

// From ml/src/regime/transition_matrix.rs:273
// Power iteration: π^(k+1) = π^(k) * P

// From ml/src/checkpoint/enterprise_implementations.rs:308
// Simple power iteration for largest eigenvalue estimate

Approach: Adapt existing pattern for spectral norm calculation.

4. Integration with DQN Architecture

4.1 Current DQN Structure

From ml/src/dqn/rainbow_network.rs:

pub struct RainbowNetworkConfig {
    pub input_size: usize,
    pub hidden_sizes: Vec<usize>,
    pub num_actions: usize,
    pub dropout_rate: f64,
    pub use_layer_norm: bool,  // ✅ Already has layer norm
    pub layer_norm_eps: f64,
    // Could add: pub use_spectral_norm: bool,
}

4.2 Proposed Integration Points

Option 1: Wrapper Around NoisyLinear

pub struct SpectralNoisyLinear {
    noisy_linear: NoisyLinear,
    u: Tensor,
    v: Tensor,
    spectral_norm_enabled: bool,
}

Option 2: Separate Layer Type

pub struct SpectralLinear {
    linear: Linear,
    u: Tensor,
    v: Tensor,
}

impl Module for SpectralLinear {
    fn forward(&self, x: &Tensor) -> CandleResult<Tensor> {
        self.normalize_weights()?;
        self.linear.forward(x)
    }
}

Recommendation: Option 1 - wrap existing NoisyLinear to preserve Rainbow DQN features.

5. Implementation Roadmap

Phase 1: Core Implementation (2-3 days)

  1. Create SpectralNorm wrapper struct
  2. Implement power iteration for spectral norm estimation
  3. Write unit tests for normalization correctness
  4. Benchmark performance overhead

Phase 2: DQN Integration (1-2 days)

  1. Add use_spectral_norm config flag
  2. Wrap NoisyLinear layers with SpectralNorm
  3. Update initialization logic
  4. Add integration tests

Phase 3: Validation (2-3 days)

  1. Compare with baseline DQN on validation set
  2. Measure overfitting reduction
  3. Ablation studies (spectral norm only vs combined with layer norm)
  4. Hyperopt integration if promising

Estimated Total: 5-8 days

6. Expected Benefits vs Cost

Benefits

  • Overfitting Reduction: Expected 10-20% improvement in validation loss
  • Training Stability: Reduced gradient explosion risk
  • Generalization: Better performance on unseen market regimes
  • Research Value: State-of-the-art regularization technique

Costs

  • Development Time: 5-8 days
  • Performance Overhead: ~5-10% slower training (1-3 power iterations per update)
  • Memory: Additional 2 * hidden_dim parameters per layer (u, v vectors)
  • Complexity: Custom implementation requires careful testing

Cost-Benefit Analysis

WORTH IMPLEMENTING if:

  • Overfitting is a major issue (validation loss >> training loss)
  • Other regularization (dropout, layer norm) insufficient
  • Training stability problems observed

NOT PRIORITY if:

  • Current regularization working well
  • Performance overhead unacceptable
  • Development resources limited

7. Alternative Anti-Overfitting Techniques

If spectral normalization is too complex, consider:

  1. Weight Decay (L2 regularization) - Already available in optimizers
  2. Gradient Clipping - Already implemented
  3. Early Stopping - Simple to implement
  4. Ensemble Methods - Already have ensemble infrastructure
  5. Data Augmentation - Add noise to state observations

8. Research References

Key Papers:

  1. Miyato et al. (2018) - "Spectral Normalization for GANs"

    • Original spectral norm paper
    • Power iteration algorithm
  2. Gouk et al. (2021) - "Regularization of Neural Networks using Lipschitz Constant"

    • Theoretical analysis of Lipschitz constraints
  3. Yoshida & Miyato (2017) - "Spectral Norm Regularization for DRL"

    • Application to reinforcement learning

9. Recommendations

Immediate Actions

  1. Research Complete - This document
  2. Decision Required: Consult with team on priority
    • Is overfitting the #1 issue?
    • Can we afford 5-8 days development time?
    • Is 5-10% performance overhead acceptable?

If Approved for Implementation

Phase 1 (Prototype):

// File: ml/src/dqn/spectral_norm.rs
pub struct SpectralNorm {
    // Implementation as outlined in Section 3.1
}

// Tests
#[cfg(test)]
mod tests {
    // Test power iteration convergence
    // Test normalization correctness
    // Benchmark performance
}

Phase 2 (Integration):

// File: ml/src/dqn/rainbow_network.rs
pub struct RainbowNetworkConfig {
    // Add: pub use_spectral_norm: bool,
}

// Wrap NoisyLinear with SpectralNorm if enabled

Phase 3 (Validation):

  • Run hyperopt with spectral_norm=true vs false
  • Compare validation metrics
  • Analyze overfitting reduction

If Not Approved

Alternatives (in priority order):

  1. Increase dropout rate (quick win)
  2. Add weight decay to optimizer (already supported)
  3. Implement early stopping (simple)
  4. Reduce network capacity (config change)

10. Technical Feasibility: CONFIRMED

Verdict: Spectral normalization is FEASIBLE to implement in candle framework.

Key Points:

  • Can use power iteration (no external dependencies)
  • Existing code has power iteration examples to adapt
  • Integration point clear (wrap NoisyLinear)
  • Performance overhead acceptable (~5-10%)
  • ⚠️ Requires custom implementation (no built-in support)
  • ⚠️ Development time: 5-8 days

Implementation Complexity: Medium (custom layer wrapper, power iteration math)

Maintenance Burden: Low (self-contained module, well-defined algorithm)


Appendix A: Power Iteration Pseudocode

def spectral_norm(W, n_iters=1):
    """
    W: weight matrix (out_features, in_features)
    n_iters: number of power iterations
    """
    # Initialize u, v randomly (or from previous iteration)
    u = torch.randn(W.size(0))
    v = torch.randn(W.size(1))

    for _ in range(n_iters):
        # Update v: v = W^T u / ||W^T u||
        v = W.t() @ u
        v = v / v.norm()

        # Update u: u = W v / ||W v||
        u = W @ v
        u = u / u.norm()

    # Spectral norm: σ = u^T W v
    sigma = u @ W @ v

    # Normalized weight
    W_normalized = W / sigma

    return W_normalized, u, v  # Keep u, v for next iteration

Appendix B: Candle Implementation Sketch

use candle_core::{Result, Tensor, Device};

pub struct SpectralNorm {
    u: Tensor,  // Left singular vector [out_features]
    v: Tensor,  // Right singular vector [in_features]
    n_power_iterations: usize,
}

impl SpectralNorm {
    pub fn new(out_features: usize, in_features: usize, device: &Device) -> Result<Self> {
        let u = Tensor::randn(0.0, 1.0, out_features, device)?;
        let v = Tensor::randn(0.0, 1.0, in_features, device)?;

        Ok(Self {
            u,
            v,
            n_power_iterations: 1,  // Typically 1 is sufficient
        })
    }

    pub fn normalize(&mut self, weight: &Tensor) -> Result<Tensor> {
        let mut u = self.u.clone();
        let mut v = self.v.clone();

        // Power iteration
        for _ in 0..self.n_power_iterations {
            // v = W^T u / ||W^T u||
            let wt_u = weight.t()?.matmul(&u)?;
            let norm_v = wt_u.sqr()?.sum_all()?.sqrt()?;
            v = wt_u.div(&norm_v)?;

            // u = W v / ||W v||
            let w_v = weight.matmul(&v)?;
            let norm_u = w_v.sqr()?.sum_all()?.sqrt()?;
            u = w_v.div(&norm_u)?;
        }

        // Compute spectral norm: σ = u^T W v
        let sigma = u.t()?.matmul(weight)?.matmul(&v)?;

        // Normalize weight
        let w_normalized = weight.div(&sigma)?;

        // Update stored vectors
        self.u = u;
        self.v = v;

        Ok(w_normalized)
    }
}

Report Compiled By: Research Agent 15 Status: RESEARCH COMPLETE Next Step: DECISION REQUIRED (implement vs alternative)