530 lines
18 KiB
Rust
530 lines
18 KiB
Rust
//! Sophia optimizer implementation
|
||
//!
|
||
//! Sophia: Second-order Clipped Stochastic Optimization (Stanford, 2023)
|
||
//!
|
||
//! Key features:
|
||
//! - Second-order information via Hessian diagonal estimation
|
||
//! - Adaptive learning with gradient clipping
|
||
//! - Parameters: lr, betas, rho (Hessian EMA), weight_decay
|
||
//! - Update rule uses Hessian for per-parameter learning rates
|
||
//! - Numerical stability with clipping
|
||
//!
|
||
//! # Mathematical Foundation
|
||
//!
|
||
//! Sophia approximates the diagonal Hessian and uses it for adaptive learning:
|
||
//! - m_t = β₁ * m_{t-1} + (1 - β₁) * g_t (momentum)
|
||
//! - h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t) (Hessian diagonal approximation)
|
||
//! - u_t = m_t / max(h_t, ε) (adaptive update)
|
||
//! - u_t = clip(u_t, γ) (gradient clipping)
|
||
//! - θ_t = θ_{t-1} - η * u_t (parameter update)
|
||
//!
|
||
//! # Performance Characteristics
|
||
//! - Uses second-order curvature information for better conditioning
|
||
//! - Adaptive per-parameter learning rates
|
||
//! - Numerically stable with clipping mechanisms
|
||
|
||
use crate::{Result, TransformerError};
|
||
use crate::optimizers::{Optimizer, BaseOptimizer};
|
||
use rtx_tensor::{Tensor, TensorError, DType};
|
||
use std::collections::HashMap;
|
||
use serde::{Deserialize, Serialize};
|
||
use tracing::{debug, trace};
|
||
|
||
/// Sophia optimizer configuration
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SophiaConfig {
|
||
/// Learning rate
|
||
pub learning_rate: f64,
|
||
/// Beta1 parameter (momentum coefficient)
|
||
pub beta1: f64,
|
||
/// Beta2 parameter (second moment coefficient)
|
||
pub beta2: f64,
|
||
/// Rho parameter (Hessian EMA coefficient)
|
||
pub rho: f64,
|
||
/// Weight decay coefficient
|
||
pub weight_decay: f64,
|
||
/// Clipping threshold for gradient clipping
|
||
pub clip_threshold: f64,
|
||
/// Epsilon for numerical stability
|
||
pub eps: f64,
|
||
}
|
||
|
||
impl Default for SophiaConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
learning_rate: 1e-3,
|
||
beta1: 0.965,
|
||
beta2: 0.99,
|
||
rho: 0.04,
|
||
weight_decay: 0.0,
|
||
clip_threshold: 1.0,
|
||
eps: 1e-8,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Sophia optimizer state for a single parameter
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SophiaState {
|
||
/// Exponential moving average of gradient (momentum)
|
||
pub momentum: Tensor,
|
||
/// Exponential moving average of Hessian diagonal approximation
|
||
pub hessian_diagonal: Tensor,
|
||
/// Step count
|
||
pub step: i64,
|
||
}
|
||
|
||
/// Sophia optimizer with second-order curvature information
|
||
///
|
||
/// Sophia uses an approximation of the diagonal Hessian to provide adaptive
|
||
/// per-parameter learning rates, leading to better convergence properties
|
||
/// especially for ill-conditioned optimization problems.
|
||
///
|
||
/// # Mathematical Details
|
||
///
|
||
/// The Sophia update rule consists of several phases:
|
||
/// 1. Update momentum: m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
|
||
/// 2. Update Hessian diagonal: h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)
|
||
/// 3. Compute adaptive update: u_t = m_t / max(h_t, ε)
|
||
/// 4. Apply clipping: u_t = clip(u_t, γ)
|
||
/// 5. Update parameter: θ_t = θ_{t-1} - η * u_t
|
||
///
|
||
/// # Adaptive Learning
|
||
/// The Hessian diagonal approximation provides per-parameter scaling,
|
||
/// allowing the optimizer to adapt to the local curvature of the loss surface.
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SophiaOptimizer {
|
||
/// Base optimizer functionality
|
||
base: BaseOptimizer,
|
||
/// Beta1 parameter (momentum coefficient)
|
||
beta1: f64,
|
||
/// Beta2 parameter (second moment coefficient)
|
||
beta2: f64,
|
||
/// Rho parameter (Hessian EMA coefficient)
|
||
rho: f64,
|
||
/// Weight decay (L2 regularization)
|
||
weight_decay: f64,
|
||
/// Clipping threshold for gradient clipping
|
||
clip_threshold: f64,
|
||
/// Epsilon for numerical stability
|
||
eps: f64,
|
||
/// Per-parameter state
|
||
state: HashMap<String, SophiaState>,
|
||
}
|
||
|
||
impl SophiaOptimizer {
|
||
/// Create a new Sophia optimizer from configuration
|
||
pub fn new(config: SophiaConfig) -> Result<Self> {
|
||
// Validate parameters
|
||
Self::validate_parameters(
|
||
config.learning_rate,
|
||
config.beta1,
|
||
config.beta2,
|
||
config.rho,
|
||
config.weight_decay,
|
||
config.clip_threshold,
|
||
config.eps,
|
||
)?;
|
||
|
||
debug!(
|
||
"Creating Sophia optimizer: lr={}, β₁={}, β₂={}, ρ={}, weight_decay={}, clip={}, ε={}",
|
||
config.learning_rate, config.beta1, config.beta2, config.rho,
|
||
config.weight_decay, config.clip_threshold, config.eps
|
||
);
|
||
|
||
Ok(Self {
|
||
base: BaseOptimizer::new(config.learning_rate),
|
||
beta1: config.beta1,
|
||
beta2: config.beta2,
|
||
rho: config.rho,
|
||
weight_decay: config.weight_decay,
|
||
clip_threshold: config.clip_threshold,
|
||
eps: config.eps,
|
||
state: HashMap::new(),
|
||
})
|
||
}
|
||
|
||
/// Validate optimizer parameters
|
||
fn validate_parameters(
|
||
learning_rate: f64,
|
||
beta1: f64,
|
||
beta2: f64,
|
||
rho: f64,
|
||
weight_decay: f64,
|
||
clip_threshold: f64,
|
||
eps: f64,
|
||
) -> Result<()> {
|
||
if learning_rate <= 0.0 {
|
||
return Err(TransformerError::generic(
|
||
format!("learning_rate {} must be positive", learning_rate)
|
||
));
|
||
}
|
||
|
||
if !(0.0..1.0).contains(&beta1) {
|
||
return Err(TransformerError::generic(
|
||
format!("beta1 {} must be in [0, 1)", beta1)
|
||
));
|
||
}
|
||
|
||
if !(0.0..1.0).contains(&beta2) {
|
||
return Err(TransformerError::generic(
|
||
format!("beta2 {} must be in [0, 1)", beta2)
|
||
));
|
||
}
|
||
|
||
if !(0.0..=1.0).contains(&rho) {
|
||
return Err(TransformerError::generic(
|
||
format!("rho {} must be in [0, 1]", rho)
|
||
));
|
||
}
|
||
|
||
if weight_decay < 0.0 {
|
||
return Err(TransformerError::generic(
|
||
format!("weight_decay {} must be non-negative", weight_decay)
|
||
));
|
||
}
|
||
|
||
if clip_threshold <= 0.0 {
|
||
return Err(TransformerError::generic(
|
||
format!("clip_threshold {} must be positive", clip_threshold)
|
||
));
|
||
}
|
||
|
||
if eps < 0.0 {
|
||
return Err(TransformerError::generic(
|
||
format!("eps {} must be non-negative", eps)
|
||
));
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Get beta1 parameter
|
||
pub fn beta1(&self) -> f64 {
|
||
self.beta1
|
||
}
|
||
|
||
/// Get beta2 parameter
|
||
pub fn beta2(&self) -> f64 {
|
||
self.beta2
|
||
}
|
||
|
||
/// Get rho parameter
|
||
pub fn rho(&self) -> f64 {
|
||
self.rho
|
||
}
|
||
|
||
/// Get weight decay parameter
|
||
pub fn weight_decay(&self) -> f64 {
|
||
self.weight_decay
|
||
}
|
||
|
||
/// Get clip threshold parameter
|
||
pub fn clip_threshold(&self) -> f64 {
|
||
self.clip_threshold
|
||
}
|
||
|
||
/// Get epsilon parameter
|
||
pub fn eps(&self) -> f64 {
|
||
self.eps
|
||
}
|
||
|
||
/// Initialize state for a parameter if it doesn't exist
|
||
fn ensure_state(&mut self, param_name: &str, param: &Tensor) -> Result<()> {
|
||
if !self.state.contains_key(param_name) {
|
||
let shape = param.shape();
|
||
let device = param.device();
|
||
|
||
let momentum = Tensor::zeros(shape.clone(), device)?;
|
||
let hessian_diagonal = Tensor::zeros(shape.clone(), device)?;
|
||
|
||
let state = SophiaState {
|
||
momentum,
|
||
hessian_diagonal,
|
||
step: 0,
|
||
};
|
||
|
||
self.state.insert(param_name.to_string(), state);
|
||
trace!("Initialized Sophia state for parameter: {}", param_name);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Apply gradient clipping element-wise
|
||
fn apply_clipping(update: &Tensor, threshold: f64) -> Result<Tensor> {
|
||
// Element-wise clipping: clip(x, threshold) = max(-threshold, min(threshold, x))
|
||
let threshold_tensor = Tensor::scalar(threshold as f32, DType::F32, update.device())?;
|
||
let neg_threshold_tensor = Tensor::scalar(-threshold as f32, DType::F32, update.device())?;
|
||
|
||
// Clamp to [-threshold, threshold]
|
||
let clipped = update.clamp(&neg_threshold_tensor, &threshold_tensor)?;
|
||
Ok(clipped)
|
||
}
|
||
|
||
/// Perform Sophia parameter update (static version)
|
||
fn update_parameter_static(
|
||
param: &Tensor,
|
||
grad: &Tensor,
|
||
state: &mut SophiaState,
|
||
learning_rate: f64,
|
||
beta1: f64,
|
||
beta2: f64,
|
||
rho: f64,
|
||
weight_decay: f64,
|
||
clip_threshold: f64,
|
||
eps: f64,
|
||
) -> Result<Tensor> {
|
||
// Increment step count
|
||
state.step += 1;
|
||
|
||
trace!("Sophia step {} for parameter", state.step);
|
||
|
||
// Step 1: Update momentum m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
|
||
let momentum_term = (&state.momentum * beta1)?;
|
||
let grad_term = (grad * (1.0 - beta1))?;
|
||
state.momentum = (momentum_term + grad_term)?;
|
||
|
||
// Step 2: Update Hessian diagonal approximation h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)
|
||
let hessian_term = (&state.hessian_diagonal * rho)?;
|
||
let grad_squared = (grad * grad)?;
|
||
let hessian_update = (&grad_squared * (1.0 - rho))?;
|
||
state.hessian_diagonal = (hessian_term + hessian_update)?;
|
||
|
||
// Step 3: Compute adaptive update u_t = m_t / max(h_t, ε)
|
||
let eps_tensor = Tensor::scalar(eps as f32, DType::F32, state.hessian_diagonal.device())?;
|
||
let hessian_with_eps = state.hessian_diagonal.maximum(&eps_tensor)?;
|
||
let hessian_sqrt = hessian_with_eps.sqrt()?;
|
||
let adaptive_update = state.momentum.div(&hessian_sqrt)?;
|
||
|
||
// Step 4: Apply gradient clipping u_t = clip(u_t, γ)
|
||
let clipped_update = Self::apply_clipping(&adaptive_update, clip_threshold)?;
|
||
|
||
// Step 5: Apply weight decay and parameter update
|
||
let scaled_update = (&clipped_update * learning_rate)?;
|
||
let final_param = if weight_decay > 0.0 {
|
||
// L2 regularization: θ_t = θ_{t-1} - α * λ * θ_{t-1} - update
|
||
(param * (1.0 - learning_rate * weight_decay))? - &scaled_update
|
||
} else {
|
||
// Standard update: θ_t = θ_{t-1} - update
|
||
param - &scaled_update
|
||
}?;
|
||
|
||
Ok(final_param)
|
||
}
|
||
}
|
||
|
||
impl Optimizer for SophiaOptimizer {
|
||
fn step_param(&mut self, param_name: &str, param: &Tensor, grad: &Tensor) -> Result<Tensor> {
|
||
// Validate input shapes match
|
||
if param.shape() != grad.shape() {
|
||
return Err(TransformerError::shape_mismatch(
|
||
format!("Parameter shape {:?} doesn't match gradient shape {:?}",
|
||
param.shape().dims().to_vec(),
|
||
grad.shape().dims().to_vec())
|
||
));
|
||
}
|
||
|
||
// Ensure state exists
|
||
self.ensure_state(param_name, param)?;
|
||
|
||
// Get the state and perform update
|
||
let state = self.state.get_mut(param_name).unwrap();
|
||
let learning_rate = self.base.learning_rate;
|
||
let beta1 = self.beta1;
|
||
let beta2 = self.beta2;
|
||
let rho = self.rho;
|
||
let weight_decay = self.weight_decay;
|
||
let clip_threshold = self.clip_threshold;
|
||
let eps = self.eps;
|
||
|
||
// Perform update using static method
|
||
Self::update_parameter_static(
|
||
param, grad, state, learning_rate, beta1, beta2,
|
||
rho, weight_decay, clip_threshold, eps
|
||
)
|
||
}
|
||
|
||
fn learning_rate(&self) -> f64 {
|
||
self.base.learning_rate
|
||
}
|
||
|
||
fn set_learning_rate(&mut self, lr: f64) -> Result<()> {
|
||
if lr <= 0.0 {
|
||
return Err(TransformerError::generic(
|
||
format!("learning_rate {} must be positive", lr)
|
||
));
|
||
}
|
||
|
||
self.base.set_learning_rate(lr);
|
||
debug!("Updated Sophia learning rate to: {}", lr);
|
||
Ok(())
|
||
}
|
||
|
||
fn has_state(&self, param_name: &str) -> bool {
|
||
self.state.contains_key(param_name)
|
||
}
|
||
|
||
fn reset_state(&mut self, param_name: &str) -> Result<()> {
|
||
if self.state.remove(param_name).is_some() {
|
||
debug!("Reset Sophia state for parameter: {}", param_name);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn reset_all_state(&mut self) {
|
||
let count = self.state.len();
|
||
self.state.clear();
|
||
debug!("Reset all Sophia state ({} parameters)", count);
|
||
}
|
||
|
||
fn get_step_count(&self, param_name: &str) -> Result<i64> {
|
||
self.state
|
||
.get(param_name)
|
||
.map(|state| state.step)
|
||
.ok_or_else(|| {
|
||
TransformerError::optimizer(format!("No state found for parameter: {}", param_name))
|
||
})
|
||
}
|
||
|
||
fn optimizer_type(&self) -> &'static str {
|
||
"Sophia"
|
||
}
|
||
|
||
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
|
||
self
|
||
}
|
||
|
||
fn store_gradients_internal(&mut self, gradients: HashMap<String, Tensor>) -> Result<()> {
|
||
debug!("Storing {} gradients for Sophia optimizer", gradients.len());
|
||
self.base.store_gradients(gradients);
|
||
Ok(())
|
||
}
|
||
|
||
fn process_stored_gradients(&mut self, learning_rate: f64) -> Result<HashMap<String, Tensor>> {
|
||
debug!("Processing {} stored gradients with lr={} using Sophia algorithm",
|
||
self.base.stored_gradients.len(), learning_rate);
|
||
|
||
// Update learning rate if different
|
||
if (self.learning_rate() - learning_rate).abs() > f64::EPSILON {
|
||
self.set_learning_rate(learning_rate)?;
|
||
}
|
||
|
||
let mut parameter_updates = HashMap::new();
|
||
|
||
// Process each stored gradient with Sophia algorithm
|
||
for (param_name, grad) in &self.base.stored_gradients {
|
||
// Initialize state if it doesn't exist
|
||
if !self.state.contains_key(param_name) {
|
||
debug!("Creating Sophia state for parameter: {}", param_name);
|
||
let shape = grad.shape();
|
||
let device = grad.device();
|
||
|
||
let momentum = Tensor::zeros(shape.clone(), device)?;
|
||
let hessian_diagonal = Tensor::zeros(shape.clone(), device)?;
|
||
|
||
let state = SophiaState {
|
||
momentum,
|
||
hessian_diagonal,
|
||
step: 0,
|
||
};
|
||
|
||
self.state.insert(param_name.clone(), state);
|
||
}
|
||
|
||
// Get mutable reference to state
|
||
let state = self.state.get_mut(param_name).unwrap();
|
||
|
||
// Increment step count
|
||
state.step += 1;
|
||
|
||
trace!("Sophia step {} for parameter {}", state.step, param_name);
|
||
|
||
// Step 1: Update momentum m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
|
||
let momentum_term = (&state.momentum * self.beta1)?;
|
||
let grad_term = (grad * (1.0 - self.beta1))?;
|
||
state.momentum = (momentum_term + grad_term)?;
|
||
|
||
// Step 2: Update Hessian diagonal approximation h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)
|
||
let hessian_term = (&state.hessian_diagonal * self.rho)?;
|
||
let grad_squared = (grad * grad)?;
|
||
let hessian_update = (&grad_squared * (1.0 - self.rho))?;
|
||
state.hessian_diagonal = (hessian_term + hessian_update)?;
|
||
|
||
// Step 3: Compute adaptive update u_t = m_t / max(h_t, ε)
|
||
let eps_tensor = Tensor::scalar(self.eps as f32, DType::F32, state.hessian_diagonal.device())?;
|
||
let hessian_with_eps = state.hessian_diagonal.maximum(&eps_tensor)?;
|
||
let hessian_sqrt = hessian_with_eps.sqrt()?;
|
||
let adaptive_update = state.momentum.div(&hessian_sqrt)?;
|
||
|
||
// Step 4: Apply gradient clipping u_t = clip(u_t, γ)
|
||
let clipped_update = Self::apply_clipping(&adaptive_update, self.clip_threshold)?;
|
||
|
||
// Step 5: Scale by learning rate and store negative update
|
||
let scaled_update = (&clipped_update * learning_rate)?;
|
||
let final_update = if self.weight_decay > 0.0 {
|
||
// Note: Weight decay requires parameter values, which we don't have here
|
||
debug!("Weight decay configured but parameter values not available in process_stored_gradients");
|
||
(scaled_update * -1.0)? // Negative because we want to subtract the update
|
||
} else {
|
||
(scaled_update * -1.0)? // Negative because we want to subtract the update
|
||
};
|
||
|
||
parameter_updates.insert(param_name.clone(), final_update);
|
||
debug!("Created Sophia update for parameter: {} (step {})", param_name, state.step);
|
||
}
|
||
|
||
// Clear stored gradients after processing
|
||
self.base.clear_gradients();
|
||
|
||
debug!("Generated {} parameter updates using Sophia algorithm", parameter_updates.len());
|
||
Ok(parameter_updates)
|
||
}
|
||
}
|
||
|
||
#[cfg(all(test, feature = "disabled_tests"))]
|
||
mod tests {
|
||
use super::*;
|
||
use rtx_tensor::{Device, Shape};
|
||
|
||
#[test]
|
||
fn test_sophia_parameter_validation() {
|
||
let config = SophiaConfig::default();
|
||
|
||
// Test valid parameters
|
||
assert!(SophiaOptimizer::new(config).is_ok());
|
||
|
||
// Test invalid learning rate
|
||
let mut invalid_config = SophiaConfig::default();
|
||
invalid_config.learning_rate = -0.1;
|
||
assert!(SophiaOptimizer::new(invalid_config).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_sophia_basic_functionality() -> Result<()> {
|
||
let config = SophiaConfig::default();
|
||
let mut optimizer = SophiaOptimizer::new(config)?;
|
||
let device = Device::cpu();
|
||
|
||
let param = Tensor::from_data(vec![1.0, 2.0], &[2], &device)?;
|
||
let grad = Tensor::from_data(vec![0.1, 0.2], &[2], &device)?;
|
||
|
||
// First step
|
||
let updated_param = optimizer.step_param("test_param", ¶m, &grad)?;
|
||
|
||
assert!(optimizer.has_state("test_param"));
|
||
assert_eq!(optimizer.get_step_count("test_param")?, 1);
|
||
|
||
// Verify parameter changed
|
||
let param_data = param.to_cpu()?;
|
||
let updated_data = updated_param.to_cpu()?;
|
||
assert_ne!(param_data, updated_data);
|
||
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
// Include tests if compiled with test features
|
||
#[cfg(all(test, feature = "disabled_tests"))]
|
||
#[path = "sophia_tests.rs"]
|
||
mod sophia_tests; |