1use crate::{
7 Optimizer, OptimizerError, OptimizerResult, OptimizerState, ParamGroup, ParamGroupState,
8};
9use parking_lot::RwLock;
10use std::collections::HashMap;
11use std::ops::Add;
12use std::sync::Arc;
13use torsh_core::error::{Result, TorshError};
14use torsh_tensor::Tensor;
15
16#[derive(Clone, Copy, Debug)]
18pub enum TrustRegionStrategy {
19 Standard,
21 Adaptive,
23 Conservative,
25 Aggressive,
27}
28
29#[derive(Clone)]
31pub struct TrustRegionConfig {
32 pub initial_radius: f32,
34 pub max_radius: f32,
36 pub min_radius: f32,
38 pub eta1: f32,
40 pub eta2: f32,
42 pub gamma1: f32,
44 pub gamma2: f32,
46 pub strategy: TrustRegionStrategy,
48 pub max_iter: usize,
50 pub tolerance_grad: f32,
52 pub tolerance_step: f32,
54}
55
56impl Default for TrustRegionConfig {
57 fn default() -> Self {
58 Self {
59 initial_radius: 1.0,
60 max_radius: 100.0,
61 min_radius: 1e-6,
62 eta1: 0.25,
63 eta2: 0.75,
64 gamma1: 0.25,
65 gamma2: 2.0,
66 strategy: TrustRegionStrategy::Standard,
67 max_iter: 100,
68 tolerance_grad: 1e-6,
69 tolerance_step: 1e-8,
70 }
71 }
72}
73
74#[derive(Clone, Copy, Debug)]
76pub enum SubproblemSolver {
77 CauchyPoint,
79 Dogleg,
81 ConjugateGradient,
83 TwoDSubspace,
85}
86
87pub type ObjectiveFn = Arc<dyn Fn(&Tensor) -> Result<f32> + Send + Sync>;
94
95pub struct TrustRegionMethod {
104 param_groups: Vec<ParamGroup>,
105 state: HashMap<String, HashMap<String, Tensor>>,
106 step_count: usize,
107 config: TrustRegionConfig,
108 solver: SubproblemSolver,
109 objective: Option<ObjectiveFn>,
110}
111
112impl TrustRegionMethod {
113 pub fn new(
114 params: Vec<Arc<RwLock<Tensor>>>,
115 lr: Option<f32>,
116 config: Option<TrustRegionConfig>,
117 solver: Option<SubproblemSolver>,
118 ) -> Self {
119 let lr = lr.unwrap_or(1.0);
120 let param_group = ParamGroup::new(params, lr);
121
122 Self {
123 param_groups: vec![param_group],
124 state: HashMap::new(),
125 step_count: 0,
126 config: config.unwrap_or_default(),
127 solver: solver.unwrap_or(SubproblemSolver::Dogleg),
128 objective: None,
129 }
130 }
131
132 pub fn set_objective<F>(&mut self, objective: F)
138 where
139 F: Fn(&Tensor) -> Result<f32> + Send + Sync + 'static,
140 {
141 self.objective = Some(Arc::new(objective));
142 }
143
144 pub fn clear_objective(&mut self) {
146 self.objective = None;
147 }
148
149 pub fn builder() -> TrustRegionBuilder {
150 TrustRegionBuilder::new()
151 }
152
153 fn get_param_id(param: &Arc<RwLock<Tensor>>) -> String {
154 format!("{:p}", Arc::as_ptr(param))
155 }
156
157 fn flatten_params(&self) -> Result<Tensor> {
159 let mut flattened = Vec::new();
160
161 for group in &self.param_groups {
162 for param in &group.params {
163 let param_read = param.read();
164 let param_flat = param_read.flatten()?;
165 let param_data = param_flat.data()?;
166 flattened.extend_from_slice(¶m_data);
167 }
168 }
169
170 let len = flattened.len();
171 Ok(Tensor::from_data(
172 flattened,
173 vec![len],
174 torsh_core::device::DeviceType::Cpu,
175 )?)
176 }
177
178 fn flatten_grads(&self) -> Result<Tensor> {
180 let mut flattened = Vec::new();
181
182 for group in &self.param_groups {
183 for param in &group.params {
184 let param_read = param.read();
185 let grad = param_read.grad().ok_or_else(|| {
186 TorshError::invalid_argument_with_context(
187 "Parameter has no gradient",
188 "trust_region_step",
189 )
190 })?;
191 let grad_flat = grad.flatten()?;
192 let grad_data = grad_flat.data()?;
193 flattened.extend_from_slice(&grad_data);
194 }
195 }
196
197 let len = flattened.len();
198 Ok(Tensor::from_data(
199 flattened,
200 vec![len],
201 torsh_core::device::DeviceType::Cpu,
202 )?)
203 }
204
205 fn update_params_from_flat(&self, flat_params: &Tensor) -> Result<()> {
207 let flat_data = flat_params.data()?;
208 let mut offset = 0;
209
210 for group in &self.param_groups {
211 for param in &group.params {
212 let mut param_write = param.write();
213 let param_shape = param_write.shape();
214 let param_size = param_shape.numel();
215
216 let param_data = &flat_data[offset..offset + param_size];
217 let new_values = Tensor::from_data(
218 param_data.to_vec(),
219 param_shape.dims().to_vec(),
220 param_write.device(),
221 )?;
222 crate::param_update::assign(&mut param_write, &new_values)?;
223
224 offset += param_size;
225 }
226 }
227
228 Ok(())
229 }
230
231 fn cauchy_point(&self, grad: &Tensor, radius: f32) -> Result<Tensor> {
233 let grad_norm = grad.norm()?.item()?;
234
235 if grad_norm < self.config.tolerance_grad {
236 return Ok(Tensor::zeros(grad.shape().dims(), grad.device())?);
237 }
238
239 let tau = (radius / grad_norm).min(1.0);
241 Ok(grad.mul_scalar(-tau * radius / grad_norm)?)
242 }
243
244 fn approximate_hessian(&self, grad: &Tensor) -> Result<Tensor> {
246 let grad_norm = grad.norm()?.item()?;
249 let scale = if grad_norm > 1e-8 { grad_norm } else { 1.0 };
250
251 let n = grad.shape().numel();
253 Ok(Tensor::from_data(vec![scale; n], vec![n], grad.device())?)
254 }
255
256 fn dogleg(&self, grad: &Tensor, hessian_diag: &Tensor, radius: f32) -> Result<Tensor> {
258 let grad_norm = grad.norm()?.item()?;
259
260 if grad_norm < self.config.tolerance_grad {
261 return Ok(Tensor::zeros(grad.shape().dims(), grad.device())?);
262 }
263
264 let cauchy_step = self.cauchy_point(grad, radius)?;
266 let cauchy_norm = cauchy_step.norm()?.item()?;
267
268 if cauchy_norm >= radius - self.config.tolerance_step {
270 return Ok(cauchy_step);
271 }
272
273 let hessian_data = hessian_diag.data()?;
276 let grad_data = grad.data()?;
277 let newton_data: Vec<f32> = grad_data
278 .iter()
279 .zip(hessian_data.iter())
280 .map(|(g, h)| if h.abs() > 1e-12 { -g / h } else { -g })
281 .collect();
282
283 let newton_step =
284 Tensor::from_data(newton_data, grad.shape().dims().to_vec(), grad.device())?;
285 let newton_norm = newton_step.norm()?.item()?;
286
287 if newton_norm <= radius {
289 return Ok(newton_step);
290 }
291
292 let diff = newton_step.sub(&cauchy_step)?;
297 let a = diff.dot(&diff)?.item()?;
298 let b = 2.0 * cauchy_step.dot(&diff)?.item()?;
299 let c = cauchy_norm * cauchy_norm - radius * radius;
300
301 let discriminant = b * b - 4.0 * a * c;
302 if discriminant < 0.0 {
303 return Ok(cauchy_step);
304 }
305
306 let tau = (-b + discriminant.sqrt()) / (2.0 * a);
307 let tau = tau.clamp(0.0, 1.0);
308
309 let result = cauchy_step
311 .mul_scalar(tau)?
312 .add(&newton_step.mul_scalar(1.0 - tau)?)?;
313 Ok(result)
314 }
315
316 fn conjugate_gradient_tr(
318 &self,
319 grad: &Tensor,
320 hessian_diag: &Tensor,
321 radius: f32,
322 ) -> Result<Tensor> {
323 let n = grad.shape().numel();
324 let tolerance = self.config.tolerance_grad;
325 let max_iter = n.min(50); let mut x = Tensor::zeros(&[n], grad.device())?;
329 let mut r = grad.neg()?; let mut p = r.clone();
331 let mut rsold = r.dot(&r)?.item()?;
332
333 for _i in 0..max_iter {
334 if rsold.sqrt() < tolerance {
335 break;
336 }
337
338 let hp_data: Vec<f32> = {
340 let hessian_data = hessian_diag.data()?;
341 let p_data = p.data()?;
342 p_data
343 .iter()
344 .zip(hessian_data.iter())
345 .map(|(p_val, h_val)| p_val * h_val)
346 .collect()
347 };
348 let hp = Tensor::from_data(hp_data, grad.shape().dims().to_vec(), grad.device())?;
349
350 let pap = p.dot(&hp)?.item()?;
351
352 if pap <= 0.0 {
354 let x_norm_sq = x.dot(&x)?.item()?;
356 let xp = x.dot(&p)?.item()?;
357 let p_norm_sq = p.dot(&p)?.item()?;
358
359 let discriminant = xp * xp + p_norm_sq * (radius * radius - x_norm_sq);
360 if discriminant >= 0.0 {
361 let alpha = (-xp + discriminant.sqrt()) / p_norm_sq;
362 return Ok(x.add(&p.mul_scalar(alpha)?)?);
363 } else {
364 return Ok(x);
365 }
366 }
367
368 let alpha = rsold / pap;
369 let x_new = x.add(&p.mul_scalar(alpha)?)?;
370
371 let x_norm = x_new.norm()?.item()?;
373 if x_norm >= radius {
374 let x_norm_old_sq = x.dot(&x)?.item()?;
376 let xp = x.dot(&p)?.item()?;
377 let p_norm_sq = p.dot(&p)?.item()?;
378
379 let discriminant = xp * xp + p_norm_sq * (radius * radius - x_norm_old_sq);
380 if discriminant >= 0.0 {
381 let tau = (-xp + discriminant.sqrt()) / p_norm_sq;
382 return Ok(x.add(&p.mul_scalar(tau)?)?);
383 } else {
384 return Ok(x);
385 }
386 }
387
388 x = x_new;
389 let r_new = r.sub(&hp.mul_scalar(alpha)?)?;
390 let rsnew = r_new.dot(&r_new)?.item()?;
391
392 if rsnew.sqrt() < tolerance {
393 break;
394 }
395
396 let beta = rsnew / rsold;
397 p = r_new.add(&p.mul_scalar(beta)?)?;
398 r = r_new;
399 rsold = rsnew;
400 }
401
402 Ok(x)
403 }
404
405 fn solve_subproblem(&self, grad: &Tensor, radius: f32) -> Result<Tensor> {
407 match self.solver {
408 SubproblemSolver::CauchyPoint => self.cauchy_point(grad, radius),
409 SubproblemSolver::Dogleg => {
410 let hessian_diag = self.approximate_hessian(grad)?;
411 self.dogleg(grad, &hessian_diag, radius)
412 }
413 SubproblemSolver::ConjugateGradient => {
414 let hessian_diag = self.approximate_hessian(grad)?;
415 self.conjugate_gradient_tr(grad, &hessian_diag, radius)
416 }
417 SubproblemSolver::TwoDSubspace => {
418 let hessian_diag = self.approximate_hessian(grad)?;
420 self.dogleg(grad, &hessian_diag, radius)
421 }
422 }
423 }
424
425 fn model_decrease(&self, grad: &Tensor, step: &Tensor, hessian_diag: &Tensor) -> Result<f32> {
431 let linear_term = grad.dot(step)?.item()?;
432 let hessian_step = step.mul_op(hessian_diag)?;
433 let quadratic_term = step.dot(&hessian_step)?.item()?;
434 let decrease = -linear_term - 0.5 * quadratic_term;
435 Ok(decrease.max(0.0))
436 }
437
438 fn actual_decrease(&self, old_params: &Tensor, new_params: &Tensor) -> Result<f32> {
446 let objective = self.objective.as_ref().ok_or_else(|| {
447 TorshError::InvalidArgument(
448 "TrustRegionMethod requires an objective function to measure the actual \
449 decrease; register one with `set_objective` before calling `step`"
450 .to_string(),
451 )
452 })?;
453 Ok(objective(old_params)? - objective(new_params)?)
454 }
455
456 fn update_radius(&self, current_radius: f32, reduction_ratio: f32) -> f32 {
458 let config = &self.config;
459
460 match config.strategy {
461 TrustRegionStrategy::Standard => {
462 if reduction_ratio < config.eta1 {
463 (current_radius * config.gamma1).max(config.min_radius)
464 } else if reduction_ratio > config.eta2 {
465 (current_radius * config.gamma2).min(config.max_radius)
466 } else {
467 current_radius
468 }
469 }
470 TrustRegionStrategy::Adaptive => {
471 if reduction_ratio < 0.1 {
473 (current_radius * 0.1).max(config.min_radius)
474 } else if reduction_ratio > 0.9 {
475 (current_radius * 3.0).min(config.max_radius)
476 } else {
477 current_radius * (0.5 + reduction_ratio)
478 }
479 }
480 TrustRegionStrategy::Conservative => {
481 if reduction_ratio < config.eta1 {
483 (current_radius * 0.5).max(config.min_radius)
484 } else if reduction_ratio > config.eta2 {
485 (current_radius * 1.5).min(config.max_radius)
486 } else {
487 current_radius
488 }
489 }
490 TrustRegionStrategy::Aggressive => {
491 if reduction_ratio < config.eta1 {
493 (current_radius * 0.1).max(config.min_radius)
494 } else if reduction_ratio > 0.5 {
495 (current_radius * 4.0).min(config.max_radius)
496 } else {
497 current_radius
498 }
499 }
500 }
501 }
502}
503
504impl Optimizer for TrustRegionMethod {
505 fn step(&mut self) -> OptimizerResult<()> {
506 self.step_count += 1;
507
508 let current_params = self.flatten_params()?;
510 let current_grad = self.flatten_grads()?;
511
512 let grad_norm = current_grad.norm()?.item()?;
514 if grad_norm < self.config.tolerance_grad {
515 return Ok(());
516 }
517
518 let state_id = "trust_region_state".to_string();
520 let mut radius = {
521 let param_state = self.state.entry(state_id.clone()).or_default();
522 if let Some(radius_tensor) = param_state.get("radius") {
523 radius_tensor.item()?
524 } else {
525 let initial_radius = self.config.initial_radius;
526 param_state.insert("radius".to_string(), Tensor::scalar(initial_radius)?);
527 initial_radius
528 }
529 };
530
531 let hessian_diag = self.approximate_hessian(¤t_grad)?;
534
535 for _iter in 0..self.config.max_iter {
537 let step = self.solve_subproblem(¤t_grad, radius)?;
539 let step_norm = step.norm()?.item()?;
540
541 if step_norm < self.config.tolerance_step {
543 break;
544 }
545
546 let new_params = current_params.add(&step)?;
548
549 let model_dec = self.model_decrease(¤t_grad, &step, &hessian_diag)?;
551 let actual_dec = self.actual_decrease(¤t_params, &new_params)?;
552
553 let reduction_ratio = if model_dec > 1e-12 {
554 actual_dec / model_dec
555 } else {
556 0.0
557 };
558
559 if reduction_ratio > self.config.eta1 {
561 self.update_params_from_flat(&new_params)?;
563 }
564
565 radius = self.update_radius(radius, reduction_ratio);
567
568 let param_state = self
570 .state
571 .get_mut(&state_id)
572 .expect("state should exist for state_id");
573 param_state.insert("radius".to_string(), Tensor::scalar(radius)?);
574
575 if reduction_ratio > self.config.eta1 {
577 break;
578 }
579
580 if radius < self.config.min_radius {
582 break;
583 }
584 }
585
586 Ok(())
587 }
588
589 fn zero_grad(&mut self) {
590 for group in &self.param_groups {
591 for param in &group.params {
592 param.write().zero_grad();
593 }
594 }
595 }
596
597 fn get_lr(&self) -> Vec<f32> {
598 self.param_groups.iter().map(|g| g.lr).collect()
599 }
600
601 fn set_lr(&mut self, lr: f32) {
602 for group in &mut self.param_groups {
603 group.lr = lr;
604 }
605 }
606
607 fn add_param_group(&mut self, params: Vec<Arc<RwLock<Tensor>>>, options: HashMap<String, f32>) {
608 let lr = options.get("lr").copied().unwrap_or(1.0);
609 let group = ParamGroup::new(params, lr).with_options(options);
610 self.param_groups.push(group);
611 }
612
613 fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
614 crate::optimizer::collect_parameters(&self.param_groups)
615 }
616
617 fn state_dict(&self) -> OptimizerResult<OptimizerState> {
618 let param_groups = self
619 .param_groups
620 .iter()
621 .map(|g| ParamGroupState {
622 lr: g.lr,
623 options: g.options.clone(),
624 param_count: g.params.len(),
625 })
626 .collect();
627
628 Ok(OptimizerState {
629 optimizer_type: "TrustRegion".to_string(),
630 version: "0.1.0".to_string(),
631 param_groups,
632 state: self.state.clone(),
633 global_state: HashMap::new(),
634 })
635 }
636
637 fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
638 if state.param_groups.len() != self.param_groups.len() {
639 return Err(OptimizerError::TensorError(TorshError::InvalidArgument(
640 "Parameter group count mismatch".to_string(),
641 )));
642 }
643
644 for (i, group_state) in state.param_groups.iter().enumerate() {
645 self.param_groups[i].lr = group_state.lr;
646 self.param_groups[i].options = group_state.options.clone();
647 }
648
649 self.state = state.state;
650 Ok(())
651 }
652}
653
654pub struct TrustRegionBuilder {
656 lr: f32,
657 config: TrustRegionConfig,
658 solver: SubproblemSolver,
659}
660
661impl TrustRegionBuilder {
662 pub fn new() -> Self {
663 Self {
664 lr: 1.0,
665 config: TrustRegionConfig::default(),
666 solver: SubproblemSolver::Dogleg,
667 }
668 }
669
670 pub fn lr(mut self, lr: f32) -> Self {
671 self.lr = lr;
672 self
673 }
674
675 pub fn initial_radius(mut self, radius: f32) -> Self {
676 self.config.initial_radius = radius;
677 self
678 }
679
680 pub fn max_radius(mut self, radius: f32) -> Self {
681 self.config.max_radius = radius;
682 self
683 }
684
685 pub fn min_radius(mut self, radius: f32) -> Self {
686 self.config.min_radius = radius;
687 self
688 }
689
690 pub fn strategy(mut self, strategy: TrustRegionStrategy) -> Self {
691 self.config.strategy = strategy;
692 self
693 }
694
695 pub fn solver(mut self, solver: SubproblemSolver) -> Self {
696 self.solver = solver;
697 self
698 }
699
700 pub fn tolerance_grad(mut self, tol: f32) -> Self {
701 self.config.tolerance_grad = tol;
702 self
703 }
704
705 pub fn tolerance_step(mut self, tol: f32) -> Self {
706 self.config.tolerance_step = tol;
707 self
708 }
709
710 pub fn max_iter(mut self, max_iter: usize) -> Self {
711 self.config.max_iter = max_iter;
712 self
713 }
714
715 pub fn build(self, params: Vec<Arc<RwLock<Tensor>>>) -> TrustRegionMethod {
716 TrustRegionMethod::new(params, Some(self.lr), Some(self.config), Some(self.solver))
717 }
718}
719
720impl Default for TrustRegionBuilder {
721 fn default() -> Self {
722 Self::new()
723 }
724}
725
726#[cfg(test)]
727mod tests {
728 use super::*;
729 use torsh_tensor::creation::randn;
730
731 #[test]
732 fn test_trust_region_creation() -> OptimizerResult<()> {
733 let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
734 let optimizer = TrustRegionMethod::new(params, None, None, None);
735 assert_eq!(optimizer.get_lr()[0], 1.0);
736 Ok(())
737 }
738
739 #[test]
740 fn test_trust_region_builder() -> OptimizerResult<()> {
741 let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
742 let optimizer = TrustRegionMethod::builder()
743 .lr(0.1)
744 .initial_radius(0.5)
745 .strategy(TrustRegionStrategy::Adaptive)
746 .solver(SubproblemSolver::ConjugateGradient)
747 .build(params);
748
749 assert_eq!(optimizer.get_lr()[0], 0.1);
750 assert_eq!(optimizer.config.initial_radius, 0.5);
751 assert!(matches!(
752 optimizer.config.strategy,
753 TrustRegionStrategy::Adaptive
754 ));
755 assert!(matches!(
756 optimizer.solver,
757 SubproblemSolver::ConjugateGradient
758 ));
759 Ok(())
760 }
761
762 #[test]
763 fn test_cauchy_point() -> OptimizerResult<()> {
764 let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
765 let optimizer = TrustRegionMethod::new(params, None, None, None);
766
767 let grad = Tensor::from_data(
768 vec![1.0, 0.0, 0.0, 1.0],
769 vec![4],
770 torsh_core::device::DeviceType::Cpu,
771 )?;
772 let radius = 1.0;
773
774 let cauchy = optimizer.cauchy_point(&grad, radius)?;
775 let cauchy_norm = cauchy.norm()?.item()?;
776
777 assert!(cauchy_norm <= radius + 1e-6);
779 Ok(())
780 }
781
782 #[test]
783 fn test_trust_region_config() {
784 let config = TrustRegionConfig {
785 initial_radius: 2.0,
786 max_radius: 50.0,
787 min_radius: 1e-5,
788 eta1: 0.1,
789 eta2: 0.9,
790 gamma1: 0.1,
791 gamma2: 3.0,
792 strategy: TrustRegionStrategy::Aggressive,
793 max_iter: 50,
794 tolerance_grad: 1e-8,
795 tolerance_step: 1e-10,
796 };
797
798 assert_eq!(config.initial_radius, 2.0);
799 assert_eq!(config.eta1, 0.1);
800 assert_eq!(config.eta2, 0.9);
801 assert!(matches!(config.strategy, TrustRegionStrategy::Aggressive));
802 }
803}