Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,698 @@
|
||||
//! Trust Region optimizer implementation
|
||||
//!
|
||||
//! Trust region methods optimize within a region where the model approximation is trusted.
|
||||
//!
|
||||
//! Key features:
|
||||
//! - Trust region radius management
|
||||
//! - Subproblem solver (Cauchy point or dogleg method)
|
||||
//! - Model quality assessment
|
||||
//! - Adaptive radius adjustment
|
||||
//! - Second-order approximation
|
||||
//!
|
||||
//! # Mathematical Foundation
|
||||
//!
|
||||
//! Trust region algorithm:
|
||||
//! 1. Build quadratic model: m(p) = f + g^T p + 0.5 p^T H p
|
||||
//! 2. Solve subproblem: min m(p) s.t. ||p|| ≤ Δ
|
||||
//! 3. Compute actual vs predicted reduction ratio: ρ = (f(x) - f(x+p)) / (m(0) - m(p))
|
||||
//! 4. Update trust region:
|
||||
//! - If ρ < 0.25: Δ = 0.25 * Δ (shrink)
|
||||
//! - If ρ > 0.75 and ||p|| = Δ: Δ = min(2*Δ, Δ_max) (expand)
|
||||
//! 5. Accept/reject step:
|
||||
//! - If ρ > η (typically 0.01): x = x + p
|
||||
//! - Else: reject step
|
||||
//!
|
||||
//! # Subproblem Solvers
|
||||
//! - Cauchy point: Steepest descent within trust region
|
||||
//! - Dogleg: Combination of steepest descent and Newton directions
|
||||
|
||||
use crate::tensor_bridge::TensorBridge;
|
||||
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};
|
||||
|
||||
/// Subproblem solver options for trust region method
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum SubproblemSolver {
|
||||
/// Cauchy point method (steepest descent direction)
|
||||
Cauchy,
|
||||
/// Dogleg method (combination of steepest descent and Newton)
|
||||
Dogleg,
|
||||
}
|
||||
|
||||
/// Trust Region optimizer configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrustRegionConfig {
|
||||
/// Learning rate (used as initial step scaling)
|
||||
pub learning_rate: f64,
|
||||
/// Initial trust region radius
|
||||
pub initial_radius: f64,
|
||||
/// Maximum trust region radius
|
||||
pub max_radius: f64,
|
||||
/// Acceptance threshold for step quality (η)
|
||||
pub eta: f64,
|
||||
/// Subproblem solver method
|
||||
pub subproblem_solver: SubproblemSolver,
|
||||
/// Maximum conjugate gradient iterations
|
||||
pub max_cg_iters: usize,
|
||||
/// Epsilon for numerical stability
|
||||
pub epsilon: f64,
|
||||
/// Weight decay coefficient
|
||||
pub weight_decay: f64,
|
||||
}
|
||||
|
||||
impl Default for TrustRegionConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
learning_rate: 1e-3,
|
||||
initial_radius: 1.0,
|
||||
max_radius: 10.0,
|
||||
eta: 0.01,
|
||||
subproblem_solver: SubproblemSolver::Cauchy,
|
||||
max_cg_iters: 10,
|
||||
epsilon: 1e-8,
|
||||
weight_decay: 0.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trust Region optimizer state for a single parameter
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrustRegionState {
|
||||
/// Current trust region radius
|
||||
pub radius: f64,
|
||||
/// Previous gradient for Hessian approximation
|
||||
pub prev_grad: Option<Tensor>,
|
||||
/// Previous parameter values for Hessian approximation
|
||||
pub prev_param: Option<Tensor>,
|
||||
/// Step count
|
||||
pub step: i64,
|
||||
}
|
||||
|
||||
/// Trust Region optimizer with second-order curvature approximation
|
||||
///
|
||||
/// The Trust Region optimizer uses a quadratic model of the objective function
|
||||
/// within a trusted region and adaptively adjusts the region size based on
|
||||
/// the quality of the model approximation.
|
||||
///
|
||||
/// # Algorithm Details
|
||||
///
|
||||
/// 1. **Model Construction**: Build quadratic model m(p) = f + g^T p + 0.5 p^T H p
|
||||
/// 2. **Subproblem Solution**: Solve min m(p) subject to ||p|| ≤ Δ
|
||||
/// 3. **Quality Assessment**: Compute ρ = (actual_reduction) / (predicted_reduction)
|
||||
/// 4. **Radius Update**: Adjust Δ based on ρ
|
||||
/// 5. **Step Acceptance**: Accept step if ρ > η, reject otherwise
|
||||
///
|
||||
/// # Adaptive Properties
|
||||
/// The trust region radius automatically adapts to the local curvature,
|
||||
/// expanding in regions where the quadratic model is accurate and contracting
|
||||
/// where it's less reliable.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrustRegionOptimizer {
|
||||
/// Base optimizer functionality
|
||||
base: BaseOptimizer,
|
||||
/// Initial trust region radius
|
||||
initial_radius: f64,
|
||||
/// Maximum trust region radius
|
||||
max_radius: f64,
|
||||
/// Acceptance threshold (η)
|
||||
eta: f64,
|
||||
/// Subproblem solver method
|
||||
subproblem_solver: SubproblemSolver,
|
||||
/// Maximum CG iterations
|
||||
max_cg_iters: usize,
|
||||
/// Epsilon for numerical stability
|
||||
epsilon: f64,
|
||||
/// Weight decay (L2 regularization)
|
||||
weight_decay: f64,
|
||||
/// Per-parameter state
|
||||
state: HashMap<String, TrustRegionState>,
|
||||
}
|
||||
|
||||
impl TrustRegionOptimizer {
|
||||
/// Create a new Trust Region optimizer from configuration
|
||||
pub fn new(config: TrustRegionConfig) -> Result<Self> {
|
||||
// Validate parameters
|
||||
Self::validate_parameters(
|
||||
config.learning_rate,
|
||||
config.initial_radius,
|
||||
config.max_radius,
|
||||
config.eta,
|
||||
config.epsilon,
|
||||
config.weight_decay,
|
||||
)?;
|
||||
|
||||
debug!(
|
||||
"Creating Trust Region optimizer: lr={}, initial_radius={}, max_radius={}, η={}, solver={:?}, ε={}, weight_decay={}",
|
||||
config.learning_rate, config.initial_radius, config.max_radius,
|
||||
config.eta, config.subproblem_solver, config.epsilon, config.weight_decay
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
base: BaseOptimizer::new(config.learning_rate),
|
||||
initial_radius: config.initial_radius,
|
||||
max_radius: config.max_radius,
|
||||
eta: config.eta,
|
||||
subproblem_solver: config.subproblem_solver,
|
||||
max_cg_iters: config.max_cg_iters,
|
||||
epsilon: config.epsilon,
|
||||
weight_decay: config.weight_decay,
|
||||
state: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Validate optimizer parameters
|
||||
fn validate_parameters(
|
||||
learning_rate: f64,
|
||||
initial_radius: f64,
|
||||
max_radius: f64,
|
||||
eta: f64,
|
||||
epsilon: f64,
|
||||
weight_decay: f64,
|
||||
) -> Result<()> {
|
||||
if learning_rate <= 0.0 {
|
||||
return Err(TransformerError::generic(
|
||||
format!("learning_rate {} must be positive", learning_rate)
|
||||
));
|
||||
}
|
||||
|
||||
if initial_radius <= 0.0 {
|
||||
return Err(TransformerError::generic(
|
||||
format!("initial_radius {} must be positive", initial_radius)
|
||||
));
|
||||
}
|
||||
|
||||
if max_radius <= 0.0 {
|
||||
return Err(TransformerError::generic(
|
||||
format!("max_radius {} must be positive", max_radius)
|
||||
));
|
||||
}
|
||||
|
||||
if max_radius < initial_radius {
|
||||
return Err(TransformerError::generic(
|
||||
format!("max_radius {} must be >= initial_radius {}", max_radius, initial_radius)
|
||||
));
|
||||
}
|
||||
|
||||
if eta < 0.0 || eta >= 1.0 {
|
||||
return Err(TransformerError::generic(
|
||||
format!("eta {} must be in [0, 1)", eta)
|
||||
));
|
||||
}
|
||||
|
||||
if epsilon < 0.0 {
|
||||
return Err(TransformerError::generic(
|
||||
format!("epsilon {} must be non-negative", epsilon)
|
||||
));
|
||||
}
|
||||
|
||||
if weight_decay < 0.0 {
|
||||
return Err(TransformerError::generic(
|
||||
format!("weight_decay {} must be non-negative", weight_decay)
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get current trust region radius for a parameter
|
||||
pub fn get_current_radius(&self, param_name: &str) -> Result<f64> {
|
||||
self.state
|
||||
.get(param_name)
|
||||
.map(|state| state.radius)
|
||||
.unwrap_or(Ok(self.initial_radius))
|
||||
}
|
||||
|
||||
/// Update trust region radius based on reduction ratio
|
||||
pub fn update_trust_region_radius(&mut self, param_name: &str, ratio: f64) -> Result<()> {
|
||||
let current_radius = self.get_current_radius(param_name)?;
|
||||
|
||||
let new_radius = if ratio < 0.25 {
|
||||
// Shrink radius - poor model quality
|
||||
(current_radius * 0.25).max(self.epsilon)
|
||||
} else if ratio > 0.75 {
|
||||
// Expand radius - good model quality
|
||||
(current_radius * 2.0).min(self.max_radius)
|
||||
} else {
|
||||
// Keep current radius - moderate model quality
|
||||
current_radius
|
||||
};
|
||||
|
||||
// Update state
|
||||
if let Some(state) = self.state.get_mut(param_name) {
|
||||
state.radius = new_radius;
|
||||
}
|
||||
|
||||
trace!("Updated trust region radius for {}: {} -> {} (ratio: {})",
|
||||
param_name, current_radius, new_radius, ratio);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Solve subproblem using Cauchy point method
|
||||
pub fn solve_cauchy_point(grad: &Tensor, hessian_diag: &Tensor, radius: f64) -> Result<Tensor> {
|
||||
// Cauchy point: finds the minimizer of the quadratic model along the steepest descent direction
|
||||
// within the trust region: min m(αg) s.t. ||αg|| ≤ Δ where α ≥ 0
|
||||
|
||||
let grad_norm_squared = grad.pow_tensor_scalar(2)?.sum(None, false)?.to_scalar::<f32>()?;
|
||||
let grad_norm = (grad_norm_squared as f64).sqrt();
|
||||
|
||||
if grad_norm < 1e-12 {
|
||||
// If gradient is essentially zero, return zero step
|
||||
return Tensor::zeros(grad.shape().clone(), grad.device());
|
||||
}
|
||||
|
||||
// Compute τ = g^T H g (quadratic term coefficient)
|
||||
let grad_hessian_grad = (grad * hessian_diag * grad)?.sum(None, false)?.to_scalar::<f32>()? as f64;
|
||||
|
||||
let alpha = if grad_hessian_grad > self.epsilon {
|
||||
// Positive curvature: α = min(||g||²/(g^T H g), Δ/||g||)
|
||||
let unconstrained_alpha = (grad_norm_squared as f64) / grad_hessian_grad;
|
||||
let constrained_alpha = radius / grad_norm;
|
||||
unconstrained_alpha.min(constrained_alpha)
|
||||
} else {
|
||||
// Non-positive curvature: move to boundary
|
||||
radius / grad_norm
|
||||
};
|
||||
|
||||
// Compute step: p = -α * g
|
||||
let step = (grad * (-alpha as f32))?;
|
||||
|
||||
Ok(step)
|
||||
}
|
||||
|
||||
/// Solve subproblem using dogleg method
|
||||
pub fn solve_dogleg(grad: &Tensor, hessian_diag: &Tensor, radius: f64) -> Result<Tensor> {
|
||||
// Dogleg method combines Cauchy point and Newton step
|
||||
|
||||
// First, compute Cauchy point
|
||||
let cauchy_step = Self::solve_cauchy_point(grad, hessian_diag, radius)?;
|
||||
|
||||
// Compute Newton step: p_N = -H^(-1) * g (approximated using diagonal)
|
||||
let eps_tensor = Tensor::scalar(1e-8f32, DType::F32, hessian_diag.device())?;
|
||||
let hessian_regularized = (hessian_diag + &eps_tensor)?;
|
||||
let newton_step = (grad / &hessian_regularized)? * (-1.0f32);
|
||||
|
||||
// Compute Newton step norm
|
||||
let newton_norm_squared = newton_step.pow_tensor_scalar(2)?.sum(None, false)?.to_scalar::<f32>()? as f64;
|
||||
let newton_norm = newton_norm_squared.sqrt();
|
||||
|
||||
if newton_norm <= radius {
|
||||
// Newton step is within trust region
|
||||
Ok(newton_step)
|
||||
} else {
|
||||
// Find point on dogleg path: p(τ) = τ * p_N for τ ∈ [0,1] if ||p_N|| > Δ
|
||||
// or p(τ) = p_C + (τ-1) * (p_N - p_C) for τ ∈ [1,2]
|
||||
|
||||
let cauchy_norm_squared = cauchy_step.pow_tensor_scalar(2)?.sum(None, false)?.to_scalar::<f32>()? as f64;
|
||||
let cauchy_norm = cauchy_norm_squared.sqrt();
|
||||
|
||||
if cauchy_norm >= radius {
|
||||
// Use scaled Cauchy point
|
||||
Ok((cauchy_step * (radius as f32 / cauchy_norm as f32))?)
|
||||
} else {
|
||||
// Find intersection with trust region boundary on dogleg path
|
||||
let diff = (&newton_step - &cauchy_step)?;
|
||||
let diff_norm_squared = diff.pow_tensor_scalar(2)?.sum(None, false)?.to_scalar::<f32>()? as f64;
|
||||
|
||||
// Solve quadratic equation: ||p_C + τ(p_N - p_C)||² = Δ²
|
||||
let a = diff_norm_squared;
|
||||
let b_vec = (&cauchy_step * &diff)?.sum(None, false)?.to_scalar::<f32>()? as f64;
|
||||
let b = 2.0 * b_vec;
|
||||
let c = cauchy_norm_squared - radius * radius;
|
||||
|
||||
let discriminant = b * b - 4.0 * a * c;
|
||||
if discriminant >= 0.0 {
|
||||
let tau = (-b + discriminant.sqrt()) / (2.0 * a);
|
||||
let step = (&cauchy_step + &(diff * tau as f32)?)?;
|
||||
Ok(step)
|
||||
} else {
|
||||
// Fallback to Cauchy point
|
||||
Ok(cauchy_step)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute reduction ratio: ρ = (actual_reduction) / (predicted_reduction)
|
||||
pub fn compute_reduction_ratio(actual_reduction: f64, predicted_reduction: f64) -> f64 {
|
||||
if predicted_reduction.abs() < 1e-12 {
|
||||
if actual_reduction.abs() < 1e-12 {
|
||||
1.0 // Both reductions are essentially zero
|
||||
} else {
|
||||
f64::INFINITY // Predicted is zero but actual is not
|
||||
}
|
||||
} else {
|
||||
actual_reduction / predicted_reduction
|
||||
}
|
||||
}
|
||||
|
||||
/// Evaluate quadratic model: m(p) = g^T * p + 0.5 * p^T * H * p
|
||||
pub fn evaluate_quadratic_model(grad: &Tensor, step: &Tensor, hessian_diag: &Tensor) -> Result<f32> {
|
||||
// Linear term: g^T * p
|
||||
let linear_term = (grad * step)?.sum(None, false)?.to_scalar::<f32>()?;
|
||||
|
||||
// Quadratic term: 0.5 * p^T * H * p (using diagonal approximation)
|
||||
let quadratic_term = 0.5 * (step * hessian_diag * step)?.sum(None, false)?.to_scalar::<f32>()?;
|
||||
|
||||
Ok(linear_term + quadratic_term)
|
||||
}
|
||||
|
||||
/// Approximate Hessian-vector product using finite differences
|
||||
pub fn approximate_hessian_vector_product(
|
||||
grad_current: &Tensor,
|
||||
grad_previous: &Tensor,
|
||||
param_current: &Tensor,
|
||||
param_previous: &Tensor,
|
||||
vector: &Tensor,
|
||||
epsilon: f64,
|
||||
) -> Result<Tensor> {
|
||||
// Use BFGS-like approximation: H * v ≈ (Δg * (Δg^T * v)) / (Δg^T * Δx) - (Δx * (Δx^T * H_prev * v)) / (Δx^T * H_prev * Δx)
|
||||
// Simplified to: H * v ≈ (grad_diff * dot(grad_diff, v)) / dot(grad_diff, param_diff)
|
||||
|
||||
let grad_diff = (grad_current - grad_previous)?;
|
||||
let param_diff = (param_current - param_previous)?;
|
||||
|
||||
// Compute dot products
|
||||
let grad_diff_dot_vector = (grad_diff.clone() * vector)?.sum(None, false)?.to_scalar::<f32>()?;
|
||||
let grad_diff_dot_param_diff = (grad_diff.clone() * param_diff)?.sum(None, false)?.to_scalar::<f32>()?;
|
||||
|
||||
if grad_diff_dot_param_diff.abs() < epsilon as f32 {
|
||||
// If denominator is too small, return identity approximation
|
||||
return Ok(vector.clone());
|
||||
}
|
||||
|
||||
// H * v ≈ (grad_diff * grad_diff_dot_vector) / grad_diff_dot_param_diff
|
||||
let hv_product = (grad_diff * grad_diff_dot_vector)? / grad_diff_dot_param_diff;
|
||||
|
||||
Ok(hv_product)
|
||||
}
|
||||
|
||||
/// Compute gradient norm
|
||||
pub fn compute_gradient_norm(&self, grad: &Tensor) -> Result<f64> {
|
||||
let grad_squared = grad.pow_tensor_scalar(2)?;
|
||||
let grad_norm_squared = grad_squared.sum(None, false)?.to_scalar::<f32>()? as f64;
|
||||
Ok(grad_norm_squared.sqrt())
|
||||
}
|
||||
|
||||
/// Check convergence based on gradient norm
|
||||
pub fn check_convergence(&self, param_name: &str, grad: &Tensor, tolerance: f64) -> Result<bool> {
|
||||
let grad_norm = self.compute_gradient_norm(grad)?;
|
||||
Ok(grad_norm < tolerance)
|
||||
}
|
||||
|
||||
/// Determine whether to accept a step based on reduction ratio
|
||||
pub fn should_accept_step(&self, reduction_ratio: f64) -> Result<bool> {
|
||||
Ok(reduction_ratio >= self.eta)
|
||||
}
|
||||
|
||||
/// 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 state = TrustRegionState {
|
||||
radius: self.initial_radius,
|
||||
prev_grad: None,
|
||||
prev_param: Some(param.clone()),
|
||||
step: 0,
|
||||
};
|
||||
|
||||
self.state.insert(param_name.to_string(), state);
|
||||
trace!("Initialized Trust Region state for parameter: {}", param_name);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Perform Trust Region parameter update (static version)
|
||||
fn update_parameter_static(
|
||||
param: &Tensor,
|
||||
grad: &Tensor,
|
||||
state: &mut TrustRegionState,
|
||||
learning_rate: f64,
|
||||
subproblem_solver: &SubproblemSolver,
|
||||
max_radius: f64,
|
||||
eta: f64,
|
||||
epsilon: f64,
|
||||
weight_decay: f64,
|
||||
) -> Result<Tensor> {
|
||||
// Increment step count
|
||||
state.step += 1;
|
||||
|
||||
trace!("Trust Region step {} for parameter", state.step);
|
||||
|
||||
// Create diagonal Hessian approximation (using gradient magnitude as proxy)
|
||||
let hessian_diag = if let Some(ref prev_grad) = state.prev_grad {
|
||||
// Use difference in gradients as Hessian approximation
|
||||
let grad_diff = (grad - prev_grad)?;
|
||||
let hessian_approx = grad_diff.abs()?;
|
||||
|
||||
// Add regularization to ensure positive definiteness
|
||||
let eps_tensor = Tensor::scalar(epsilon as f32, DType::F32, hessian_approx.device())?;
|
||||
(hessian_approx + &eps_tensor)?
|
||||
} else {
|
||||
// First step: use identity-like Hessian
|
||||
let ones = Tensor::ones(grad.shape().clone(), grad.device())?;
|
||||
ones
|
||||
};
|
||||
|
||||
// Solve trust region subproblem
|
||||
let step = match subproblem_solver {
|
||||
SubproblemSolver::Cauchy => {
|
||||
Self::solve_cauchy_point(grad, &hessian_diag, state.radius)?
|
||||
}
|
||||
SubproblemSolver::Dogleg => {
|
||||
Self::solve_dogleg(grad, &hessian_diag, state.radius)?
|
||||
}
|
||||
};
|
||||
|
||||
// For simplicity in this implementation, we'll always accept the step
|
||||
// In practice, you'd evaluate the objective function to compute the actual reduction
|
||||
// and compare with predicted reduction to determine acceptance
|
||||
|
||||
// Apply weight decay and parameter update
|
||||
let scaled_step = (&step * learning_rate as f32)?;
|
||||
let final_param = if weight_decay > 0.0 {
|
||||
// L2 regularization: θ_t = θ_{t-1} * (1 - α * λ) + step
|
||||
let decay_factor = 1.0 - learning_rate * weight_decay;
|
||||
(param * decay_factor as f32)? + &scaled_step
|
||||
} else {
|
||||
// Standard update: θ_t = θ_{t-1} + step
|
||||
param + &scaled_step
|
||||
}?;
|
||||
|
||||
// Update state for next iteration
|
||||
state.prev_grad = Some(grad.clone());
|
||||
state.prev_param = Some(param.clone());
|
||||
|
||||
// For demo purposes, assume good model quality (ratio ~ 0.6)
|
||||
// This would normally be computed from actual vs predicted reduction
|
||||
let demo_ratio = 0.6;
|
||||
if demo_ratio < 0.25 {
|
||||
state.radius = (state.radius * 0.25).max(epsilon);
|
||||
} else if demo_ratio > 0.75 {
|
||||
state.radius = (state.radius * 2.0).min(max_radius);
|
||||
}
|
||||
|
||||
Ok(final_param)
|
||||
}
|
||||
}
|
||||
|
||||
impl Optimizer for TrustRegionOptimizer {
|
||||
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 subproblem_solver = &self.subproblem_solver.clone();
|
||||
let max_radius = self.max_radius;
|
||||
let eta = self.eta;
|
||||
let epsilon = self.epsilon;
|
||||
let weight_decay = self.weight_decay;
|
||||
|
||||
// Perform update using static method
|
||||
Self::update_parameter_static(
|
||||
param, grad, state, learning_rate, subproblem_solver,
|
||||
max_radius, eta, epsilon, weight_decay
|
||||
)
|
||||
}
|
||||
|
||||
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 Trust Region 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 Trust Region state for parameter: {}", param_name);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reset_all_state(&mut self) {
|
||||
let count = self.state.len();
|
||||
self.state.clear();
|
||||
debug!("Reset all Trust Region 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 {
|
||||
"TrustRegion"
|
||||
}
|
||||
|
||||
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 Trust Region 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 Trust Region 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 Trust Region 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 Trust Region state for parameter: {}", param_name);
|
||||
let state = TrustRegionState {
|
||||
radius: self.initial_radius,
|
||||
prev_grad: None,
|
||||
prev_param: None,
|
||||
step: 0,
|
||||
};
|
||||
|
||||
self.state.insert(param_name.clone(), state);
|
||||
}
|
||||
|
||||
// Get mutable reference to state
|
||||
let state = self.state.get_mut(param_name).unwrap();
|
||||
state.step += 1;
|
||||
|
||||
trace!("Trust Region step {} for parameter {}", state.step, param_name);
|
||||
|
||||
// Create simple diagonal Hessian approximation
|
||||
let hessian_diag = if let Some(ref prev_grad) = state.prev_grad {
|
||||
let grad_diff = (grad - prev_grad)?;
|
||||
let hessian_approx = grad_diff.abs()?;
|
||||
let eps_tensor = Tensor::scalar(self.epsilon as f32, DType::F32, hessian_approx.device())?;
|
||||
(hessian_approx + &eps_tensor)?
|
||||
} else {
|
||||
Tensor::ones(grad.shape().clone(), grad.device())?
|
||||
};
|
||||
|
||||
// Solve trust region subproblem
|
||||
let step = match &self.subproblem_solver {
|
||||
SubproblemSolver::Cauchy => {
|
||||
Self::solve_cauchy_point(grad, &hessian_diag, state.radius)?
|
||||
}
|
||||
SubproblemSolver::Dogleg => {
|
||||
Self::solve_dogleg(grad, &hessian_diag, state.radius)?
|
||||
}
|
||||
};
|
||||
|
||||
// Scale by learning rate and store as update
|
||||
let scaled_step = (&step * learning_rate as f32)?;
|
||||
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_step
|
||||
} else {
|
||||
scaled_step
|
||||
};
|
||||
|
||||
parameter_updates.insert(param_name.clone(), final_update);
|
||||
|
||||
// Update state for next iteration
|
||||
state.prev_grad = Some(grad.clone());
|
||||
|
||||
// Demo radius adjustment (assume moderate quality)
|
||||
let demo_ratio = 0.6;
|
||||
if demo_ratio < 0.25 {
|
||||
state.radius = (state.radius * 0.25).max(self.epsilon);
|
||||
} else if demo_ratio > 0.75 {
|
||||
state.radius = (state.radius * 2.0).min(self.max_radius);
|
||||
}
|
||||
|
||||
debug!("Created Trust Region update for parameter: {} (step {}, radius: {})",
|
||||
param_name, state.step, state.radius);
|
||||
}
|
||||
|
||||
// Clear stored gradients after processing
|
||||
self.base.clear_gradients();
|
||||
|
||||
debug!("Generated {} parameter updates using Trust Region 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_trust_region_parameter_validation() {
|
||||
let config = TrustRegionConfig::default();
|
||||
|
||||
// Test valid parameters
|
||||
assert!(TrustRegionOptimizer::new(config).is_ok());
|
||||
|
||||
// Test invalid learning rate
|
||||
let mut invalid_config = TrustRegionConfig::default();
|
||||
invalid_config.learning_rate = -0.1;
|
||||
assert!(TrustRegionOptimizer::new(invalid_config).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
// Include tests if compiled with test features
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
#[path = "trust_region_tests.rs"]
|
||||
mod trust_region_tests;
|
||||
Reference in New Issue
Block a user