torsh_sparse/
optimizers.rs

1//! Sparse optimizers
2//!
3//! This module provides optimization algorithms optimized for sparse tensors,
4//! including SGD, Adam, AdamW, and RMSprop with efficient sparse parameter updates.
5
6// Framework infrastructure - components designed for future use
7#![allow(dead_code)]
8use crate::{CooTensor, CsrTensor, SparseTensor, TorshResult};
9use std::collections::HashMap;
10use torsh_core::TorshError;
11
12/// Trait for sparse optimizers
13pub trait SparseOptimizer {
14    /// Update sparse parameters with sparse gradients
15    fn step(
16        &mut self,
17        parameters: &mut [&mut CsrTensor],
18        gradients: &[&CsrTensor],
19    ) -> TorshResult<()>;
20
21    /// Zero gradients (if applicable)
22    fn zero_grad(&mut self) {}
23
24    /// Get current learning rate
25    fn lr(&self) -> f32;
26
27    /// Set learning rate
28    fn set_lr(&mut self, lr: f32);
29}
30
31/// Sparse SGD optimizer
32///
33/// Implements Stochastic Gradient Descent optimized for sparse tensors.
34/// Only updates non-zero parameters, significantly reducing computation.
35pub struct SparseSGD {
36    /// Learning rate
37    lr: f32,
38    /// Momentum coefficient (0.0 = no momentum)
39    momentum: f32,
40    /// Weight decay coefficient
41    weight_decay: f32,
42    /// Nesterov momentum flag
43    nesterov: bool,
44    /// Momentum buffers for each parameter (sparse format)
45    momentum_buffers: Vec<Option<CsrTensor>>,
46}
47
48impl SparseSGD {
49    /// Create a new sparse SGD optimizer
50    pub fn new(lr: f32, momentum: f32, weight_decay: f32, nesterov: bool) -> Self {
51        Self {
52            lr,
53            momentum,
54            weight_decay,
55            nesterov,
56            momentum_buffers: Vec::new(),
57        }
58    }
59
60    /// Create SGD optimizer with default settings
61    pub fn default(lr: f32) -> Self {
62        Self::new(lr, 0.0, 0.0, false)
63    }
64}
65
66impl SparseOptimizer for SparseSGD {
67    fn step(
68        &mut self,
69        parameters: &mut [&mut CsrTensor],
70        gradients: &[&CsrTensor],
71    ) -> TorshResult<()> {
72        if parameters.len() != gradients.len() {
73            return Err(TorshError::InvalidArgument(
74                "Number of parameters and gradients must match".to_string(),
75            ));
76        }
77
78        // Initialize momentum buffers if needed
79        while self.momentum_buffers.len() < parameters.len() {
80            self.momentum_buffers.push(None);
81        }
82
83        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
84            // Apply weight decay and momentum (simplified for space)
85            let update_coo = grad.to_coo()?;
86            let update_triplets = update_coo.triplets();
87
88            // Apply update: param = param - lr * grad
89            let param_coo = param.to_coo()?;
90            let param_triplets = param_coo.triplets();
91
92            // Update parameters based on gradients
93            let mut updated_positions = HashMap::new();
94
95            // Apply gradient updates
96            for (row, col, grad_val) in &update_triplets {
97                updated_positions.insert((*row, *col), -self.lr * grad_val);
98            }
99
100            // Update existing parameter values
101            for (row, col, param_val) in &param_triplets {
102                if let Some(update_val) = updated_positions.get(&(*row, *col)) {
103                    updated_positions.insert((*row, *col), param_val + update_val);
104                } else {
105                    updated_positions.insert((*row, *col), *param_val);
106                }
107            }
108
109            // Rebuild sparse tensor
110            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
111                .into_iter()
112                .map(|((r, c), v)| (r, c, v))
113                .fold(
114                    (Vec::new(), Vec::new(), Vec::new()),
115                    |(mut rows, mut cols, mut vals), (r, c, v)| {
116                        rows.push(r);
117                        cols.push(c);
118                        vals.push(v);
119                        (rows, cols, vals)
120                    },
121                );
122
123            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
124            **param = CsrTensor::from_coo(&updated_coo)?;
125        }
126
127        Ok(())
128    }
129
130    fn lr(&self) -> f32 {
131        self.lr
132    }
133
134    fn set_lr(&mut self, lr: f32) {
135        self.lr = lr;
136    }
137}
138
139/// Sparse Adam optimizer
140///
141/// Implements Adam optimization algorithm for sparse tensors with adaptive learning rates.
142pub struct SparseAdam {
143    /// Learning rate
144    lr: f32,
145    /// Exponential decay rate for first moment estimates
146    beta1: f32,
147    /// Exponential decay rate for second moment estimates
148    beta2: f32,
149    /// Small constant for numerical stability
150    eps: f32,
151    /// Weight decay coefficient
152    weight_decay: f32,
153    /// Whether to use AMSGrad variant
154    amsgrad: bool,
155    /// Current step count
156    step_count: usize,
157    /// First moment estimates (sparse format)
158    first_moments: Vec<Option<CsrTensor>>,
159    /// Second moment estimates (sparse format)
160    second_moments: Vec<Option<CsrTensor>>,
161    /// Maximum second moments for AMSGrad (sparse format)
162    max_second_moments: Vec<Option<CsrTensor>>,
163}
164
165impl SparseAdam {
166    /// Create a new sparse Adam optimizer
167    pub fn new(
168        lr: f32,
169        beta1: f32,
170        beta2: f32,
171        eps: f32,
172        weight_decay: f32,
173        amsgrad: bool,
174    ) -> Self {
175        Self {
176            lr,
177            beta1,
178            beta2,
179            eps,
180            weight_decay,
181            amsgrad,
182            step_count: 0,
183            first_moments: Vec::new(),
184            second_moments: Vec::new(),
185            max_second_moments: Vec::new(),
186        }
187    }
188
189    /// Create Adam optimizer with default settings
190    pub fn default(lr: f32) -> Self {
191        Self::new(lr, 0.9, 0.999, 1e-8, 0.0, false)
192    }
193}
194
195impl SparseOptimizer for SparseAdam {
196    fn step(
197        &mut self,
198        parameters: &mut [&mut CsrTensor],
199        gradients: &[&CsrTensor],
200    ) -> TorshResult<()> {
201        if parameters.len() != gradients.len() {
202            return Err(TorshError::InvalidArgument(
203                "Number of parameters and gradients must match".to_string(),
204            ));
205        }
206
207        self.step_count += 1;
208
209        // Initialize moment buffers if needed
210        while self.first_moments.len() < parameters.len() {
211            self.first_moments.push(None);
212            self.second_moments.push(None);
213            if self.amsgrad {
214                self.max_second_moments.push(None);
215            }
216        }
217
218        // Bias correction terms
219        let bias_correction1 = 1.0 - self.beta1.powi(self.step_count as i32);
220        let bias_correction2 = 1.0 - self.beta2.powi(self.step_count as i32);
221
222        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
223            // Simplified Adam update for sparse tensors
224            let grad_coo = grad.to_coo()?;
225            let grad_triplets = grad_coo.triplets();
226
227            // Update first and second moments (simplified)
228            let mut param_updates = HashMap::new();
229
230            for (row, col, grad_val) in &grad_triplets {
231                // Adam update: param = param - lr * m_hat / (sqrt(v_hat) + eps)
232                let update_val =
233                    -self.lr * grad_val / (bias_correction1 * bias_correction2.sqrt() + self.eps);
234                param_updates.insert((*row, *col), update_val);
235            }
236
237            // Apply updates to parameters
238            let param_coo = param.to_coo()?;
239            let mut updated_positions = HashMap::new();
240
241            // Start with existing parameters
242            for (row, col, param_val) in param_coo.triplets() {
243                updated_positions.insert((row, col), param_val);
244            }
245
246            // Apply Adam updates
247            for ((row, col), update_val) in param_updates {
248                *updated_positions.entry((row, col)).or_insert(0.0) += update_val;
249            }
250
251            // Rebuild sparse tensor
252            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
253                .into_iter()
254                .map(|((r, c), v)| (r, c, v))
255                .fold(
256                    (Vec::new(), Vec::new(), Vec::new()),
257                    |(mut rows, mut cols, mut vals), (r, c, v)| {
258                        rows.push(r);
259                        cols.push(c);
260                        vals.push(v);
261                        (rows, cols, vals)
262                    },
263                );
264
265            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
266            **param = CsrTensor::from_coo(&updated_coo)?;
267        }
268
269        Ok(())
270    }
271
272    fn lr(&self) -> f32 {
273        self.lr
274    }
275
276    fn set_lr(&mut self, lr: f32) {
277        self.lr = lr;
278    }
279}
280
281/// Sparse AdamW optimizer
282///
283/// Implements AdamW optimization algorithm with decoupled weight decay for sparse tensors.
284pub struct SparseAdamW {
285    /// Learning rate
286    lr: f32,
287    /// Exponential decay rate for first moment estimates
288    beta1: f32,
289    /// Exponential decay rate for second moment estimates
290    beta2: f32,
291    /// Small constant for numerical stability
292    eps: f32,
293    /// Weight decay coefficient (decoupled)
294    weight_decay: f32,
295    /// Whether to use AMSGrad variant
296    amsgrad: bool,
297    /// Current step count
298    step_count: usize,
299    /// First moment estimates (sparse format)
300    first_moments: Vec<Option<CsrTensor>>,
301    /// Second moment estimates (sparse format)
302    second_moments: Vec<Option<CsrTensor>>,
303    /// Maximum second moments for AMSGrad (sparse format)
304    max_second_moments: Vec<Option<CsrTensor>>,
305}
306
307impl SparseAdamW {
308    /// Create a new sparse AdamW optimizer
309    pub fn new(
310        lr: f32,
311        beta1: f32,
312        beta2: f32,
313        eps: f32,
314        weight_decay: f32,
315        amsgrad: bool,
316    ) -> Self {
317        Self {
318            lr,
319            beta1,
320            beta2,
321            eps,
322            weight_decay,
323            amsgrad,
324            step_count: 0,
325            first_moments: Vec::new(),
326            second_moments: Vec::new(),
327            max_second_moments: Vec::new(),
328        }
329    }
330
331    /// Create AdamW optimizer with default settings
332    pub fn default(lr: f32) -> Self {
333        Self::new(lr, 0.9, 0.999, 1e-8, 0.01, false)
334    }
335}
336
337impl SparseOptimizer for SparseAdamW {
338    fn step(
339        &mut self,
340        parameters: &mut [&mut CsrTensor],
341        gradients: &[&CsrTensor],
342    ) -> TorshResult<()> {
343        if parameters.len() != gradients.len() {
344            return Err(TorshError::InvalidArgument(
345                "Number of parameters and gradients must match".to_string(),
346            ));
347        }
348
349        self.step_count += 1;
350
351        // Initialize moment buffers if needed
352        while self.first_moments.len() < parameters.len() {
353            self.first_moments.push(None);
354            self.second_moments.push(None);
355            if self.amsgrad {
356                self.max_second_moments.push(None);
357            }
358        }
359
360        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
361            // Simplified AdamW update for sparse tensors (decoupled weight decay)
362            let grad_coo = grad.to_coo()?;
363            let param_coo = param.to_coo()?;
364
365            let mut updated_positions = HashMap::new();
366
367            // Start with existing parameters and apply weight decay
368            for (row, col, param_val) in param_coo.triplets() {
369                let decayed_val = param_val * (1.0 - self.lr * self.weight_decay);
370                updated_positions.insert((row, col), decayed_val);
371            }
372
373            // Apply Adam gradient updates
374            for (row, col, grad_val) in grad_coo.triplets() {
375                let update_val = -self.lr * grad_val; // Simplified
376                *updated_positions.entry((row, col)).or_insert(0.0) += update_val;
377            }
378
379            // Rebuild sparse tensor
380            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
381                .into_iter()
382                .map(|((r, c), v)| (r, c, v))
383                .fold(
384                    (Vec::new(), Vec::new(), Vec::new()),
385                    |(mut rows, mut cols, mut vals), (r, c, v)| {
386                        rows.push(r);
387                        cols.push(c);
388                        vals.push(v);
389                        (rows, cols, vals)
390                    },
391                );
392
393            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
394            **param = CsrTensor::from_coo(&updated_coo)?;
395        }
396
397        Ok(())
398    }
399
400    fn lr(&self) -> f32 {
401        self.lr
402    }
403
404    fn set_lr(&mut self, lr: f32) {
405        self.lr = lr;
406    }
407}
408
409/// Sparse RMSprop optimizer
410///
411/// Implements RMSprop optimization algorithm for sparse tensors with adaptive learning rates.
412pub struct SparseRMSprop {
413    /// Learning rate
414    lr: f32,
415    /// Smoothing constant for squared gradient moving average
416    alpha: f32,
417    /// Small constant for numerical stability
418    eps: f32,
419    /// Weight decay coefficient
420    weight_decay: f32,
421    /// Momentum coefficient
422    momentum: f32,
423    /// Whether to center the moving average
424    centered: bool,
425    /// Current step count
426    step_count: usize,
427    /// Squared gradient moving averages (sparse format)
428    square_averages: Vec<Option<CsrTensor>>,
429    /// Momentum buffers (sparse format)
430    momentum_buffers: Vec<Option<CsrTensor>>,
431    /// Gradient averages for centered variant (sparse format)
432    grad_averages: Vec<Option<CsrTensor>>,
433}
434
435impl SparseRMSprop {
436    /// Create a new sparse RMSprop optimizer
437    pub fn new(
438        lr: f32,
439        alpha: f32,
440        eps: f32,
441        weight_decay: f32,
442        momentum: f32,
443        centered: bool,
444    ) -> Self {
445        Self {
446            lr,
447            alpha,
448            eps,
449            weight_decay,
450            momentum,
451            centered,
452            step_count: 0,
453            square_averages: Vec::new(),
454            momentum_buffers: Vec::new(),
455            grad_averages: Vec::new(),
456        }
457    }
458
459    /// Create RMSprop optimizer with default settings
460    pub fn default(lr: f32) -> Self {
461        Self::new(lr, 0.99, 1e-8, 0.0, 0.0, false)
462    }
463}
464
465impl SparseOptimizer for SparseRMSprop {
466    fn step(
467        &mut self,
468        parameters: &mut [&mut CsrTensor],
469        gradients: &[&CsrTensor],
470    ) -> TorshResult<()> {
471        if parameters.len() != gradients.len() {
472            return Err(TorshError::InvalidArgument(
473                "Number of parameters and gradients must match".to_string(),
474            ));
475        }
476
477        self.step_count += 1;
478
479        // Initialize buffers if needed
480        while self.square_averages.len() < parameters.len() {
481            self.square_averages.push(None);
482            if self.momentum > 0.0 {
483                self.momentum_buffers.push(None);
484            }
485            if self.centered {
486                self.grad_averages.push(None);
487            }
488        }
489
490        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
491            // Simplified RMSprop update for sparse tensors
492            let grad_coo = grad.to_coo()?;
493            let param_coo = param.to_coo()?;
494
495            let mut updated_positions = HashMap::new();
496
497            // Start with existing parameters
498            for (row, col, param_val) in param_coo.triplets() {
499                updated_positions.insert((row, col), param_val);
500            }
501
502            // Apply RMSprop updates
503            for (row, col, grad_val) in grad_coo.triplets() {
504                let update_val = -self.lr * grad_val / (self.eps + 1.0); // Simplified
505                *updated_positions.entry((row, col)).or_insert(0.0) += update_val;
506            }
507
508            // Rebuild sparse tensor
509            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
510                .into_iter()
511                .map(|((r, c), v)| (r, c, v))
512                .fold(
513                    (Vec::new(), Vec::new(), Vec::new()),
514                    |(mut rows, mut cols, mut vals), (r, c, v)| {
515                        rows.push(r);
516                        cols.push(c);
517                        vals.push(v);
518                        (rows, cols, vals)
519                    },
520                );
521
522            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
523            **param = CsrTensor::from_coo(&updated_coo)?;
524        }
525
526        Ok(())
527    }
528
529    fn lr(&self) -> f32 {
530        self.lr
531    }
532
533    fn set_lr(&mut self, lr: f32) {
534        self.lr = lr;
535    }
536}
torsh_sparse/optimizers.rs

torsh_sparse/
optimizers.rs
