torsh_optim/
online_learning.rs

1//! Online learning optimizers and variance-reduced methods
2//!
3//! This module provides optimizers designed for online learning scenarios,
4//! including variance-reduced methods like SVRG and SAGA.
5
6use crate::{Optimizer, OptimizerResult, OptimizerState, ParamGroupState};
7use parking_lot::RwLock;
8use std::collections::HashMap;
9use std::ops::Add;
10use std::sync::Arc;
11use torsh_core::error::{Result, TorshError};
12use torsh_tensor::Tensor;
13
14/// Online Gradient Descent (OGD) optimizer
15///
16/// This is a simple online learning algorithm that updates parameters
17/// immediately upon receiving each gradient.
18pub struct OnlineGradientDescent {
19    /// Parameters
20    params: Vec<Arc<RwLock<Tensor>>>,
21    /// Learning rate
22    lr: f32,
23    /// Regularization strength
24    regularization: f32,
25    /// Regret bound parameter
26    regret_bound: f32,
27    /// Number of steps taken
28    step_count: usize,
29}
30
31impl OnlineGradientDescent {
32    /// Create a new OnlineGradientDescent optimizer
33    pub fn new(
34        params: Vec<Arc<RwLock<Tensor>>>,
35        lr: f32,
36        regularization: Option<f32>,
37        regret_bound: Option<f32>,
38    ) -> Self {
39        Self {
40            params,
41            lr,
42            regularization: regularization.unwrap_or(0.0),
43            regret_bound: regret_bound.unwrap_or(1.0),
44            step_count: 0,
45        }
46    }
47
48    /// Create OGD with adaptive learning rate
49    pub fn new_adaptive(params: Vec<Arc<RwLock<Tensor>>>, regret_bound: f32) -> Self {
50        Self::new(params, 1.0, Some(0.01), Some(regret_bound))
51    }
52
53    /// Get adaptive learning rate based on step count
54    fn get_adaptive_lr(&self) -> f32 {
55        if self.step_count == 0 {
56            self.lr
57        } else {
58            self.regret_bound / (self.step_count as f32).sqrt()
59        }
60    }
61
62    /// Get regret bound
63    pub fn regret_bound(&self) -> f32 {
64        self.regret_bound
65    }
66}
67
68impl Optimizer for OnlineGradientDescent {
69    fn step(&mut self) -> OptimizerResult<()> {
70        self.step_count += 1;
71        let adaptive_lr = self.get_adaptive_lr();
72
73        for param_arc in &self.params {
74            let mut param = param_arc.write();
75            let grad = param
76                .grad()
77                .ok_or_else(|| TorshError::AutogradError("No gradient available".to_string()))?;
78
79            // Apply regularization
80            let mut effective_grad = grad.clone();
81            if self.regularization > 0.0 {
82                effective_grad = effective_grad.add(&param.mul_scalar(self.regularization)?)?;
83            }
84
85            // Update parameters
86            let update = effective_grad.mul_scalar(adaptive_lr)?;
87            crate::param_update::sub_assign(&mut param, &update)?;
88        }
89
90        Ok(())
91    }
92
93    fn zero_grad(&mut self) {
94        for param in &self.params {
95            param.write().zero_grad();
96        }
97    }
98
99    fn get_lr(&self) -> Vec<f32> {
100        vec![self.get_adaptive_lr()]
101    }
102
103    fn set_lr(&mut self, lr: f32) {
104        self.lr = lr;
105    }
106
107    fn add_param_group(
108        &mut self,
109        params: Vec<Arc<RwLock<Tensor>>>,
110        _options: HashMap<String, f32>,
111    ) {
112        self.params.extend(params);
113    }
114
115    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
116        self.params.clone()
117    }
118
119    fn state_dict(&self) -> OptimizerResult<OptimizerState> {
120        let param_group = ParamGroupState {
121            lr: self.lr,
122            options: [
123                ("regularization".to_string(), self.regularization),
124                ("regret_bound".to_string(), self.regret_bound),
125                ("step_count".to_string(), self.step_count as f32),
126            ]
127            .iter()
128            .cloned()
129            .collect(),
130            param_count: self.params.len(),
131        };
132
133        Ok(OptimizerState {
134            optimizer_type: "OnlineGradientDescent".to_string(),
135            version: "0.1.0".to_string(),
136            param_groups: vec![param_group],
137            state: HashMap::new(),
138            global_state: HashMap::new(),
139        })
140    }
141
142    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
143        if state.optimizer_type != "OnlineGradientDescent" {
144            return Err(crate::OptimizerError::InvalidParameter(format!(
145                "Expected OnlineGradientDescent, got {}",
146                state.optimizer_type
147            )));
148        }
149
150        // Load hyperparameters from param groups
151        if let Some(param_group) = state.param_groups.first() {
152            self.lr = param_group.lr;
153
154            if let Some(&regularization) = param_group.options.get("regularization") {
155                self.regularization = regularization;
156            }
157            if let Some(&regret_bound) = param_group.options.get("regret_bound") {
158                self.regret_bound = regret_bound;
159            }
160            if let Some(&step_count) = param_group.options.get("step_count") {
161                self.step_count = step_count as usize;
162            }
163        }
164
165        Ok(())
166    }
167}
168
169/// Stochastic Variance Reduced Gradient (SVRG) optimizer
170///
171/// SVRG reduces variance in stochastic gradients by periodically computing
172/// full gradients and using them as control variates.
173pub struct SVRG {
174    /// Parameters
175    params: Vec<Arc<RwLock<Tensor>>>,
176    /// Learning rate
177    lr: f32,
178    /// Epoch length (frequency of full gradient computation)
179    epoch_length: usize,
180    /// Current step in epoch
181    epoch_step: usize,
182    /// Full gradient from last epoch
183    full_gradients: Vec<Tensor>,
184    /// Parameters at start of epoch
185    epoch_params: Vec<Tensor>,
186    /// Whether full gradient is available
187    has_full_gradient: bool,
188}
189
190impl SVRG {
191    /// Create a new SVRG optimizer
192    pub fn new(params: Vec<Arc<RwLock<Tensor>>>, lr: f32, epoch_length: Option<usize>) -> Self {
193        let epoch_length = epoch_length.unwrap_or(100);
194
195        Self {
196            params,
197            lr,
198            epoch_length,
199            epoch_step: 0,
200            full_gradients: Vec::new(),
201            epoch_params: Vec::new(),
202            has_full_gradient: false,
203        }
204    }
205
206    /// Compute full gradient (to be called with full dataset)
207    pub fn compute_full_gradient(&mut self) -> Result<()> {
208        // Store current parameters as epoch parameters
209        self.epoch_params = self.params.iter().map(|p| p.read().clone()).collect();
210
211        // Store full gradients
212        self.full_gradients = self
213            .params
214            .iter()
215            .map(|p| {
216                let param = p.read();
217                param.grad().unwrap_or_else(|| {
218                    Tensor::zeros(param.shape().dims(), param.device())
219                        .expect("tensor creation should succeed")
220                })
221            })
222            .collect();
223
224        self.has_full_gradient = true;
225        self.epoch_step = 0;
226
227        Ok(())
228    }
229
230    /// Perform SVRG update with mini-batch gradient
231    pub fn svrg_step(
232        &mut self,
233        minibatch_grad_at_current: &[Tensor],
234        minibatch_grad_at_epoch: &[Tensor],
235    ) -> Result<()> {
236        if !self.has_full_gradient {
237            return Err(TorshError::AutogradError(
238                "Must compute full gradient before SVRG steps".to_string(),
239            ));
240        }
241
242        for (i, param_arc) in self.params.iter().enumerate() {
243            let mut param = param_arc.write();
244
245            // SVRG update: gradient = minibatch_grad_current - minibatch_grad_epoch + full_grad_epoch
246            let variance_reduced_grad = minibatch_grad_at_current[i]
247                .sub(&minibatch_grad_at_epoch[i])?
248                .add(&self.full_gradients[i])?;
249
250            // Update parameters
251            let update = variance_reduced_grad.mul_scalar(self.lr)?;
252            crate::param_update::sub_assign(&mut param, &update)?;
253        }
254
255        self.epoch_step += 1;
256
257        // Check if epoch is complete
258        if self.epoch_step >= self.epoch_length {
259            self.has_full_gradient = false; // Need to recompute full gradient
260        }
261
262        Ok(())
263    }
264
265    /// Check if new epoch is needed
266    pub fn needs_new_epoch(&self) -> bool {
267        !self.has_full_gradient || self.epoch_step >= self.epoch_length
268    }
269
270    /// Get epoch parameters
271    pub fn epoch_params(&self) -> &[Tensor] {
272        &self.epoch_params
273    }
274}
275
276impl Optimizer for SVRG {
277    fn step(&mut self) -> OptimizerResult<()> {
278        // Regular SGD step if no full gradient available
279        for param_arc in &self.params {
280            let mut param = param_arc.write();
281            let grad = param
282                .grad()
283                .ok_or_else(|| TorshError::AutogradError("No gradient available".to_string()))?;
284
285            let update = grad.mul_scalar(self.lr)?;
286            crate::param_update::sub_assign(&mut param, &update)?;
287        }
288
289        Ok(())
290    }
291
292    fn zero_grad(&mut self) {
293        for param in &self.params {
294            param.write().zero_grad();
295        }
296    }
297
298    fn get_lr(&self) -> Vec<f32> {
299        vec![self.lr]
300    }
301
302    fn set_lr(&mut self, lr: f32) {
303        self.lr = lr;
304    }
305
306    fn add_param_group(
307        &mut self,
308        params: Vec<Arc<RwLock<Tensor>>>,
309        _options: HashMap<String, f32>,
310    ) {
311        self.params.extend(params);
312    }
313
314    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
315        self.params.clone()
316    }
317
318    fn state_dict(&self) -> OptimizerResult<OptimizerState> {
319        let param_group = ParamGroupState {
320            lr: self.lr,
321            options: [
322                ("epoch_length".to_string(), self.epoch_length as f32),
323                ("epoch_step".to_string(), self.epoch_step as f32),
324                (
325                    "has_full_gradient".to_string(),
326                    if self.has_full_gradient { 1.0 } else { 0.0 },
327                ),
328            ]
329            .iter()
330            .cloned()
331            .collect(),
332            param_count: self.params.len(),
333        };
334
335        // NOTE: Full gradient and epoch parameter serialization not included
336        // Rationale: These are large tensors that should be recomputed fresh on the next epoch
337        // for variance reduction methods. This is consistent with standard SVRG practices
338        // where the full gradient is recalculated at each epoch start.
339        // See: Johnson & Zhang (2013) "Accelerating Stochastic Gradient Descent using Predictive Variance Reduction"
340
341        Ok(OptimizerState {
342            optimizer_type: "SVRG".to_string(),
343            version: "0.1.0".to_string(),
344            param_groups: vec![param_group],
345            state: HashMap::new(), // Full gradients excluded by design - recomputed on epoch start
346            global_state: HashMap::new(),
347        })
348    }
349
350    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
351        if state.optimizer_type != "SVRG" {
352            return Err(crate::OptimizerError::InvalidParameter(format!(
353                "Expected SVRG, got {}",
354                state.optimizer_type
355            )));
356        }
357
358        // Load hyperparameters from param groups
359        if let Some(param_group) = state.param_groups.first() {
360            self.lr = param_group.lr;
361
362            if let Some(&epoch_length) = param_group.options.get("epoch_length") {
363                self.epoch_length = epoch_length as usize;
364            }
365            if let Some(&epoch_step) = param_group.options.get("epoch_step") {
366                self.epoch_step = epoch_step as usize;
367            }
368            if let Some(&has_full_gradient) = param_group.options.get("has_full_gradient") {
369                self.has_full_gradient = has_full_gradient > 0.5;
370            }
371        }
372
373        // Note: Full gradients and epoch params are not restored as they should be
374        // recomputed on the next epoch. This is consistent with variance reduction
375        // methods which require fresh gradient computations.
376        if self.has_full_gradient {
377            // Reset to require fresh full gradient computation
378            self.has_full_gradient = false;
379        }
380
381        Ok(())
382    }
383}
384
385/// Stochastic Average Gradient Algorithm (SAGA) optimizer
386///
387/// SAGA maintains a table of gradients for each data point and uses
388/// the average as a control variate to reduce variance.
389pub struct SAGA {
390    /// Parameters
391    params: Vec<Arc<RwLock<Tensor>>>,
392    /// Learning rate
393    lr: f32,
394    /// Gradient table (one gradient per data point per parameter)
395    gradient_table: HashMap<usize, Vec<Tensor>>,
396    /// Sum of all gradients in table
397    gradient_sum: Vec<Tensor>,
398    /// Number of data points
399    num_data_points: usize,
400    /// Whether gradient table is initialized
401    is_initialized: bool,
402}
403
404impl SAGA {
405    /// Create a new SAGA optimizer
406    pub fn new(params: Vec<Arc<RwLock<Tensor>>>, lr: f32, num_data_points: usize) -> Self {
407        Self {
408            params,
409            lr,
410            gradient_table: HashMap::new(),
411            gradient_sum: Vec::new(),
412            num_data_points,
413            is_initialized: false,
414        }
415    }
416
417    /// Initialize gradient table with zeros
418    pub fn initialize(&mut self) -> Result<()> {
419        // Initialize gradient sum with zeros
420        self.gradient_sum = self
421            .params
422            .iter()
423            .map(|p| {
424                Tensor::zeros(p.read().shape().dims(), p.read().device())
425                    .expect("tensor creation should succeed")
426            })
427            .collect();
428
429        // Initialize gradient table with zeros for each data point
430        for data_idx in 0..self.num_data_points {
431            let grad_for_point: Vec<Tensor> = self
432                .params
433                .iter()
434                .map(|p| {
435                    Tensor::zeros(p.read().shape().dims(), p.read().device())
436                        .expect("tensor creation should succeed")
437                })
438                .collect();
439            self.gradient_table.insert(data_idx, grad_for_point);
440        }
441
442        self.is_initialized = true;
443        Ok(())
444    }
445
446    /// Perform SAGA update for a specific data point
447    pub fn saga_step(&mut self, data_index: usize, current_gradients: &[Tensor]) -> Result<()> {
448        if !self.is_initialized {
449            self.initialize()?;
450        }
451
452        if data_index >= self.num_data_points {
453            return Err(TorshError::InvalidArgument(format!(
454                "Data index {} exceeds number of data points {}",
455                data_index, self.num_data_points
456            )));
457        }
458
459        // Get old gradient for this data point
460        let old_gradients = self.gradient_table.get(&data_index).ok_or_else(|| {
461            TorshError::InvalidArgument("Gradient table not properly initialized".to_string())
462        })?;
463
464        // Update parameters using SAGA rule
465        for (i, param_arc) in self.params.iter().enumerate() {
466            let mut param = param_arc.write();
467
468            // SAGA update: gradient = current_grad - old_grad + average_grad
469            let average_grad = self.gradient_sum[i].div_scalar(self.num_data_points as f32)?;
470            let saga_grad = current_gradients[i]
471                .sub(&old_gradients[i])?
472                .add(&average_grad)?;
473
474            // Update parameters
475            let update = saga_grad.mul_scalar(self.lr)?;
476            crate::param_update::sub_assign(&mut param, &update)?;
477
478            // Update gradient sum
479            self.gradient_sum[i] = self.gradient_sum[i]
480                .sub(&old_gradients[i])?
481                .add(&current_gradients[i])?;
482        }
483
484        // Update gradient table
485        self.gradient_table
486            .insert(data_index, current_gradients.to_vec());
487
488        Ok(())
489    }
490
491    /// Get average gradient
492    pub fn average_gradient(&self) -> Result<Vec<Tensor>> {
493        if !self.is_initialized {
494            return Err(TorshError::AutogradError(
495                "SAGA not initialized".to_string(),
496            ));
497        }
498
499        self.gradient_sum
500            .iter()
501            .map(|grad| grad.div_scalar(self.num_data_points as f32))
502            .collect::<Result<Vec<_>>>()
503    }
504
505    /// Get number of data points
506    pub fn num_data_points(&self) -> usize {
507        self.num_data_points
508    }
509}
510
511impl Optimizer for SAGA {
512    fn step(&mut self) -> OptimizerResult<()> {
513        // Regular SGD step if not using SAGA-specific method
514        for param_arc in &self.params {
515            let mut param = param_arc.write();
516            let grad = param
517                .grad()
518                .ok_or_else(|| TorshError::AutogradError("No gradient available".to_string()))?;
519
520            let update = grad.mul_scalar(self.lr)?;
521            crate::param_update::sub_assign(&mut param, &update)?;
522        }
523
524        Ok(())
525    }
526
527    fn zero_grad(&mut self) {
528        for param in &self.params {
529            param.write().zero_grad();
530        }
531    }
532
533    fn get_lr(&self) -> Vec<f32> {
534        vec![self.lr]
535    }
536
537    fn set_lr(&mut self, lr: f32) {
538        self.lr = lr;
539    }
540
541    fn add_param_group(
542        &mut self,
543        params: Vec<Arc<RwLock<Tensor>>>,
544        _options: HashMap<String, f32>,
545    ) {
546        self.params.extend(params);
547    }
548
549    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
550        self.params.clone()
551    }
552
553    fn state_dict(&self) -> OptimizerResult<OptimizerState> {
554        let param_group = ParamGroupState {
555            lr: self.lr,
556            options: [
557                ("num_data_points".to_string(), self.num_data_points as f32),
558                (
559                    "is_initialized".to_string(),
560                    if self.is_initialized { 1.0 } else { 0.0 },
561                ),
562            ]
563            .iter()
564            .cloned()
565            .collect(),
566            param_count: self.params.len(),
567        };
568
569        // Serialize the full gradient table and the running gradient sum, so a
570        // restored optimizer resumes with the variance-reduction state it had
571        // rather than restarting from zeros.
572        //
573        // Layout: the per-data-point gradients live under
574        // `"data_<index>"` -> `"param_<slot>"`, and the running sum under
575        // `"gradient_sum"` -> `"param_<slot>"`.
576        let mut state: HashMap<String, HashMap<String, Tensor>> = HashMap::new();
577        for (data_index, gradients) in &self.gradient_table {
578            let entry = state.entry(format!("data_{data_index}")).or_default();
579            for (slot, gradient) in gradients.iter().enumerate() {
580                entry.insert(format!("param_{slot}"), gradient.clone());
581            }
582        }
583        if !self.gradient_sum.is_empty() {
584            let entry = state.entry("gradient_sum".to_string()).or_default();
585            for (slot, gradient) in self.gradient_sum.iter().enumerate() {
586                entry.insert(format!("param_{slot}"), gradient.clone());
587            }
588        }
589
590        Ok(OptimizerState {
591            optimizer_type: "SAGA".to_string(),
592            version: "0.1.0".to_string(),
593            param_groups: vec![param_group],
594            state,
595            global_state: HashMap::new(),
596        })
597    }
598
599    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
600        // Validate optimizer type
601        if state.optimizer_type != "SAGA" {
602            return Err(crate::OptimizerError::InvalidParameter(format!(
603                "Expected SAGA optimizer state, got {}",
604                state.optimizer_type
605            )));
606        }
607
608        // Load parameter group state
609        if let Some(param_group) = state.param_groups.first() {
610            self.lr = param_group.lr;
611
612            // Load SAGA-specific parameters from options
613            if let Some(&num_data_points) = param_group.options.get("num_data_points") {
614                self.num_data_points = num_data_points as usize;
615            }
616            if let Some(&is_initialized) = param_group.options.get("is_initialized") {
617                self.is_initialized = is_initialized > 0.0;
618            }
619
620            // Validate parameter count
621            if param_group.param_count != self.params.len() {
622                return Err(crate::OptimizerError::InvalidParameter(format!(
623                    "Parameter count mismatch: expected {}, got {}",
624                    self.params.len(),
625                    param_group.param_count
626                )));
627            }
628        } else {
629            return Err(crate::OptimizerError::InvalidParameter(
630                "No parameter groups found in state".to_string(),
631            ));
632        }
633
634        // Restore the gradient table and the running gradient sum written by
635        // `state_dict`. Entries are keyed `data_<index>` / `gradient_sum`, each
636        // holding one tensor per parameter slot (`param_<slot>`).
637        self.gradient_table.clear();
638        self.gradient_sum.clear();
639        let slot_count = self.params.len();
640        for (key, entry) in &state.state {
641            let mut gradients = Vec::with_capacity(slot_count);
642            for slot in 0..slot_count {
643                let tensor = entry.get(&format!("param_{slot}")).ok_or_else(|| {
644                    crate::OptimizerError::StateError(format!(
645                        "SAGA state entry `{key}` is missing parameter slot {slot}"
646                    ))
647                })?;
648                gradients.push(tensor.clone());
649            }
650
651            if key == "gradient_sum" {
652                self.gradient_sum = gradients;
653            } else if let Some(index) = key.strip_prefix("data_") {
654                let index: usize = index.parse().map_err(|_| {
655                    crate::OptimizerError::StateError(format!(
656                        "SAGA state entry `{key}` does not carry a numeric data index"
657                    ))
658                })?;
659                self.gradient_table.insert(index, gradients);
660            } else {
661                return Err(crate::OptimizerError::StateError(format!(
662                    "Unrecognized SAGA state entry `{key}`"
663                )));
664            }
665        }
666
667        Ok(())
668    }
669}
670
671/// Proximal Gradient Method optimizer
672///
673/// This optimizer handles non-smooth regularization terms using
674/// proximal operators, commonly used for L1 regularization.
675pub struct ProximalGradient {
676    /// Parameters
677    params: Vec<Arc<RwLock<Tensor>>>,
678    /// Learning rate
679    lr: f32,
680    /// L1 regularization strength
681    l1_reg: f32,
682    /// L2 regularization strength
683    l2_reg: f32,
684    /// Proximal operator type
685    prox_type: ProximalOperator,
686}
687
688/// Types of proximal operators
689#[derive(Debug, Clone)]
690pub enum ProximalOperator {
691    /// L1 regularization (soft thresholding)
692    L1,
693    /// L2 regularization (scaling)
694    L2,
695    /// Elastic net (L1 + L2)
696    ElasticNet,
697    /// Group LASSO
698    GroupLasso,
699}
700
701impl ProximalGradient {
702    /// Create a new ProximalGradient optimizer
703    pub fn new(
704        params: Vec<Arc<RwLock<Tensor>>>,
705        lr: f32,
706        l1_reg: Option<f32>,
707        l2_reg: Option<f32>,
708        prox_type: Option<ProximalOperator>,
709    ) -> Self {
710        Self {
711            params,
712            lr,
713            l1_reg: l1_reg.unwrap_or(0.0),
714            l2_reg: l2_reg.unwrap_or(0.0),
715            prox_type: prox_type.unwrap_or(ProximalOperator::L1),
716        }
717    }
718
719    /// Create ProximalGradient for LASSO (L1 regularization)
720    pub fn lasso(params: Vec<Arc<RwLock<Tensor>>>, lr: f32, l1_reg: f32) -> Self {
721        Self::new(params, lr, Some(l1_reg), None, Some(ProximalOperator::L1))
722    }
723
724    /// Create ProximalGradient for Elastic Net
725    pub fn elastic_net(
726        params: Vec<Arc<RwLock<Tensor>>>,
727        lr: f32,
728        l1_reg: f32,
729        l2_reg: f32,
730    ) -> Self {
731        Self::new(
732            params,
733            lr,
734            Some(l1_reg),
735            Some(l2_reg),
736            Some(ProximalOperator::ElasticNet),
737        )
738    }
739
740    /// Apply proximal operator
741    fn apply_proximal_operator(&self, param: &Tensor) -> Result<Tensor> {
742        match self.prox_type {
743            ProximalOperator::L1 => {
744                // Soft thresholding for L1
745                self.soft_threshold(param, self.lr * self.l1_reg)
746            }
747            ProximalOperator::L2 => {
748                // Scaling for L2
749                let scale = 1.0 / (1.0 + self.lr * self.l2_reg);
750                Ok(param.mul_scalar(scale)?)
751            }
752            ProximalOperator::ElasticNet => {
753                // L2 scaling followed by L1 soft thresholding
754                let l2_scale = 1.0 / (1.0 + self.lr * self.l2_reg);
755                let l2_result = param.mul_scalar(l2_scale)?;
756                self.soft_threshold(&l2_result, self.lr * self.l1_reg)
757            }
758            ProximalOperator::GroupLasso => {
759                // Group soft thresholding (simplified version)
760                let param_norm = param.norm()?.item()?;
761                let threshold = self.lr * self.l1_reg;
762
763                if param_norm <= threshold {
764                    Ok(Tensor::zeros(param.shape().dims(), param.device())?)
765                } else {
766                    let scale = (param_norm - threshold) / param_norm;
767                    Ok(param.mul_scalar(scale)?)
768                }
769            }
770        }
771    }
772
773    /// Soft thresholding operator for L1 regularization
774    fn soft_threshold(&self, param: &Tensor, threshold: f32) -> Result<Tensor> {
775        // Element-wise soft thresholding: sign(x) * max(|x| - threshold, 0)
776        let abs_param = param.abs()?;
777        let mask = abs_param.gt_scalar(threshold)?;
778        let threshold_vec = vec![threshold; abs_param.numel()];
779        let threshold_tensor = Tensor::from_vec(threshold_vec, &abs_param.shape().dims())?;
780        let thresholded = abs_param
781            .sub(&threshold_tensor)?
782            .maximum(&Tensor::zeros_like(param)?)?;
783        let result = param.sign()?.mul_op(&thresholded)?;
784        // Convert boolean mask to float values (1.0 for true, 0.0 for false)
785        let mask_data = mask.to_vec()?;
786        let mask_f32_data: Vec<f32> = mask_data
787            .iter()
788            .map(|&b| if b { 1.0 } else { 0.0 })
789            .collect();
790        let mask_f32 = Tensor::from_vec(mask_f32_data, &mask.shape().dims())?;
791        Ok(result.mul_op(&mask_f32)?) // Apply mask to zero out elements below threshold
792    }
793
794    /// Get regularization strengths
795    pub fn regularization(&self) -> (f32, f32) {
796        (self.l1_reg, self.l2_reg)
797    }
798
799    /// Set regularization strengths
800    pub fn set_regularization(&mut self, l1_reg: f32, l2_reg: f32) {
801        self.l1_reg = l1_reg;
802        self.l2_reg = l2_reg;
803    }
804}
805
806impl Optimizer for ProximalGradient {
807    fn step(&mut self) -> OptimizerResult<()> {
808        for param_arc in &self.params {
809            let mut param = param_arc.write();
810            let grad = param
811                .grad()
812                .ok_or_else(|| TorshError::AutogradError("No gradient available".to_string()))?;
813
814            // Gradient step
815            let grad_step = param.sub(&grad.mul_scalar(self.lr)?)?;
816
817            // Apply proximal operator
818            let proximal_result = self.apply_proximal_operator(&grad_step)?;
819            crate::param_update::assign(&mut param, &proximal_result)?;
820        }
821
822        Ok(())
823    }
824
825    fn zero_grad(&mut self) {
826        for param in &self.params {
827            param.write().zero_grad();
828        }
829    }
830
831    fn get_lr(&self) -> Vec<f32> {
832        vec![self.lr]
833    }
834
835    fn set_lr(&mut self, lr: f32) {
836        self.lr = lr;
837    }
838
839    fn add_param_group(
840        &mut self,
841        params: Vec<Arc<RwLock<Tensor>>>,
842        _options: HashMap<String, f32>,
843    ) {
844        self.params.extend(params);
845    }
846
847    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
848        self.params.clone()
849    }
850
851    fn state_dict(&self) -> OptimizerResult<OptimizerState> {
852        let param_group = ParamGroupState {
853            lr: self.lr,
854            options: [
855                ("l1_reg".to_string(), self.l1_reg),
856                ("l2_reg".to_string(), self.l2_reg),
857            ]
858            .iter()
859            .cloned()
860            .collect(),
861            param_count: self.params.len(),
862        };
863
864        Ok(OptimizerState {
865            optimizer_type: "ProximalGradient".to_string(),
866            version: "0.1.0".to_string(),
867            param_groups: vec![param_group],
868            state: HashMap::new(),
869            global_state: HashMap::new(),
870        })
871    }
872
873    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
874        // Validate optimizer type
875        if state.optimizer_type != "ProximalGradient" {
876            return Err(crate::OptimizerError::InvalidParameter(format!(
877                "Expected ProximalGradient optimizer state, got {}",
878                state.optimizer_type
879            )));
880        }
881
882        // Load parameter group state
883        if let Some(param_group) = state.param_groups.first() {
884            self.lr = param_group.lr;
885
886            // Load regularization parameters from options
887            if let Some(&l1_reg) = param_group.options.get("l1_reg") {
888                self.l1_reg = l1_reg;
889            }
890            if let Some(&l2_reg) = param_group.options.get("l2_reg") {
891                self.l2_reg = l2_reg;
892            }
893
894            // Validate parameter count
895            if param_group.param_count != self.params.len() {
896                return Err(crate::OptimizerError::InvalidParameter(format!(
897                    "Parameter count mismatch: expected {}, got {}",
898                    self.params.len(),
899                    param_group.param_count
900                )));
901            }
902        } else {
903            return Err(crate::OptimizerError::InvalidParameter(
904                "No parameter groups found in state".to_string(),
905            ));
906        }
907
908        // Note: ProximalGradient doesn't maintain per-parameter state like momentum,
909        // so we only need to restore the hyperparameters (lr, l1_reg, l2_reg)
910
911        Ok(())
912    }
913}
914
915#[cfg(test)]
916mod tests {
917    use super::*;
918    use torsh_tensor::creation::randn;
919
920    #[test]
921    fn test_online_gradient_descent() {
922        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[10, 10]).unwrap()))];
923        let optimizer = OnlineGradientDescent::new(params, 0.01, None, None);
924        assert_eq!(optimizer.get_lr()[0], 0.01);
925    }
926
927    #[test]
928    fn test_svrg_creation() {
929        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[10, 10]).unwrap()))];
930        let optimizer = SVRG::new(params, 0.01, Some(50));
931        assert!(optimizer.needs_new_epoch());
932    }
933
934    #[test]
935    fn test_saga_creation() {
936        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[10, 10]).unwrap()))];
937        let optimizer = SAGA::new(params, 0.01, 100);
938        assert_eq!(optimizer.num_data_points(), 100);
939    }
940
941    #[test]
942    fn test_proximal_gradient() {
943        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[10, 10]).unwrap()))];
944        let optimizer = ProximalGradient::lasso(params, 0.01, 0.1);
945        let (l1, l2) = optimizer.regularization();
946        assert_eq!(l1, 0.1);
947        assert_eq!(l2, 0.0);
948    }
949}
torsh_optim/online_learning.rs

torsh_optim/
online_learning.rs
