Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,530 @@
|
||||
//! 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;
|
||||
Reference in New Issue
Block a user