trustformers/auto/optimizers/mod.rs
1//! # Optimizer System for TrustFormeRS
2//!
3//! This module provides automatic optimizer selection and configuration for various
4//! machine learning tasks and model architectures. It follows the design patterns
5//! established by HuggingFace Transformers, providing intelligent defaults while
6//! allowing for fine-grained control when needed.
7//!
8//! ## Key Components
9//!
10//! - **AutoOptimizer**: Main entry point for automatic optimizer creation
11//! - **Optimizer trait**: Base interface that all optimizers must implement
12//! - **OptimizerGradients/OptimizerUpdate**: Data structures for gradient-based optimization
13//! - **LearningRateSchedule**: Various learning rate scheduling strategies
14//! - **Concrete Optimizers**: AdamW, Adam, and scheduled optimizer implementations
15//!
16//! ## Usage Examples
17//!
18//! ### Automatic Optimizer Selection
19//!
20//! ```rust,ignore
21//! use trustformers::auto::optimizers::AutoOptimizer;
22//!
23//! // Create optimizer from model configuration
24//! let optimizer = AutoOptimizer::from_pretrained("bert-base-uncased")?;
25//!
26//! // Create optimizer for specific task
27//! let task_optimizer = AutoOptimizer::for_task("text-classification", &config)?;
28//! ```
29//!
30//! ### Manual Optimizer Configuration
31//!
32//! ```rust,ignore
33//! use trustformers::auto::optimizers::{AdamWOptimizer, AdamWConfig};
34//!
35//! let config = AdamWConfig {
36//! learning_rate: 2e-5,
37//! beta1: 0.9,
38//! beta2: 0.999,
39//! weight_decay: 0.01,
40//! eps: 1e-8,
41//! amsgrad: false,
42//! };
43//! let optimizer = AdamWOptimizer::new(config);
44//! ```
45//!
46//! ### Learning Rate Scheduling
47//!
48//! ```rust,ignore
49//! use trustformers::auto::optimizers::{AutoOptimizer, LearningRateSchedule};
50//!
51//! let base_optimizer = AutoOptimizer::from_config(&config)?;
52//! let schedule = LearningRateSchedule::LinearWarmup {
53//! warmup_steps: 1000,
54//! max_lr: 5e-5,
55//! };
56//! let scheduled_optimizer = AutoOptimizer::with_schedule(base_optimizer, schedule);
57//! ```
58
59use crate::error::Result;
60use std::collections::HashMap;
61
62// =============================================================================
63// AutoOptimizer - Main Entry Point
64// =============================================================================
65
66/// Automatically create optimizers based on model and training configuration
67///
68/// The AutoOptimizer provides intelligent defaults for different model architectures
69/// and tasks, while supporting custom configurations when needed. It follows the
70/// principle of "smart defaults, flexible overrides" to minimize configuration
71/// overhead while maintaining full control when required.
72#[derive(Debug, Clone)]
73pub struct AutoOptimizer;
74
75impl AutoOptimizer {
76 /// Create an optimizer from model configuration loaded from Hub
77 ///
78 /// This method loads model configuration from the HuggingFace Hub and selects
79 /// an appropriate optimizer based on model characteristics such as parameter
80 /// count and architecture type.
81 ///
82 /// # Arguments
83 ///
84 /// * `model_name_or_path` - Model identifier from Hub or local path
85 ///
86 /// # Examples
87 ///
88 /// ```rust,ignore
89 /// /// let optimizer = AutoOptimizer::from_pretrained("bert-base-uncased")?;
90
91 pub fn from_pretrained(model_name_or_path: &str) -> Result<Box<dyn Optimizer>> {
92 let config = crate::hub::load_config_from_hub(model_name_or_path, None)?;
93 Self::from_config(&config)
94 }
95
96 /// Create an optimizer from configuration object
97 ///
98 /// Analyzes the model configuration to estimate parameter count and choose
99 /// appropriate optimizer settings. Larger models typically benefit from
100 /// AdamW with higher weight decay, while smaller models work well with
101 /// standard Adam optimization.
102 ///
103 /// # Parameter Selection Logic
104 ///
105 /// - **> 1B parameters**: AdamW with lr=1e-5, weight_decay=0.1, beta2=0.95
106 /// - **> 100M parameters**: AdamW with lr=2e-5, weight_decay=0.01, beta2=0.999
107 /// - **< 100M parameters**: Adam with lr=5e-5, no weight decay
108 ///
109 /// # Arguments
110 ///
111 /// * `config` - Model configuration as JSON value
112 pub fn from_config(config: &serde_json::Value) -> Result<Box<dyn Optimizer>> {
113 // Selection below is purely a function of the estimated parameter
114 // count (see the doc comment above); `model_type` is not read here
115 // because it does not currently affect the choice of optimizer.
116
117 // Choose optimizer based on model characteristics
118 let hidden_size =
119 config.get("hidden_size").and_then(|v| v.as_u64()).unwrap_or(768) as usize;
120 let num_layers =
121 config.get("num_hidden_layers").and_then(|v| v.as_u64()).unwrap_or(12) as usize;
122
123 // Estimate parameter count (rough approximation)
124 let estimated_params = hidden_size * hidden_size * num_layers * 4;
125
126 if estimated_params > 1_000_000_000 {
127 // > 1B parameters - Use conservative settings for large models
128 Ok(Box::new(AdamWOptimizer::new(AdamWConfig {
129 learning_rate: 1e-5,
130 beta1: 0.9,
131 beta2: 0.95, // Lower beta2 for more stable training
132 weight_decay: 0.1,
133 eps: 1e-8,
134 amsgrad: false,
135 })))
136 } else if estimated_params > 100_000_000 {
137 // > 100M parameters - Standard settings for medium models
138 Ok(Box::new(AdamWOptimizer::new(AdamWConfig {
139 learning_rate: 2e-5,
140 beta1: 0.9,
141 beta2: 0.999,
142 weight_decay: 0.01,
143 eps: 1e-8,
144 amsgrad: false,
145 })))
146 } else {
147 // < 100M parameters - Higher learning rate for smaller models
148 Ok(Box::new(AdamOptimizer::new(AdamConfig {
149 learning_rate: 5e-5,
150 beta1: 0.9,
151 beta2: 0.999,
152 eps: 1e-8,
153 amsgrad: false,
154 })))
155 }
156 }
157
158 /// Create an optimizer optimized for a specific task
159 ///
160 /// Different tasks benefit from different optimization strategies based on
161 /// their specific requirements and characteristics.
162 ///
163 /// # Task-Specific Configurations
164 ///
165 /// - **Text Generation**: AdamW with beta2=0.95 for stable generation
166 /// - **Classification**: Adam with standard settings for faster convergence
167 /// - **Question Answering**: AdamW with moderate weight decay for generalization
168 ///
169 /// # Arguments
170 ///
171 /// * `task` - Task identifier (e.g., "text-generation", "text-classification")
172 /// * `model_config` - Model configuration for fallback parameter estimation
173 pub fn for_task(task: &str, model_config: &serde_json::Value) -> Result<Box<dyn Optimizer>> {
174 match task {
175 "text-generation" | "causal-lm" => {
176 // For generation tasks, use AdamW with specific settings
177 // Lower beta2 helps with stability during generation
178 Ok(Box::new(AdamWOptimizer::new(AdamWConfig {
179 learning_rate: 2e-5,
180 beta1: 0.9,
181 beta2: 0.95,
182 weight_decay: 0.1,
183 eps: 1e-8,
184 amsgrad: false,
185 })))
186 },
187 "text-classification" | "sentiment-analysis" => {
188 // For classification, standard Adam often works well
189 // Higher learning rate for faster convergence on classification heads
190 Ok(Box::new(AdamOptimizer::new(AdamConfig {
191 learning_rate: 2e-5,
192 beta1: 0.9,
193 beta2: 0.999,
194 eps: 1e-8,
195 amsgrad: false,
196 })))
197 },
198 "question-answering" => {
199 // QA benefits from AdamW with moderate weight decay
200 // Balances memorization and generalization
201 Ok(Box::new(AdamWOptimizer::new(AdamWConfig {
202 learning_rate: 3e-5,
203 beta1: 0.9,
204 beta2: 0.999,
205 weight_decay: 0.01,
206 eps: 1e-8,
207 amsgrad: false,
208 })))
209 },
210 _ => Self::from_config(model_config),
211 }
212 }
213
214 /// Create an optimizer with learning rate scheduling
215 ///
216 /// Wraps any base optimizer with a learning rate schedule for improved
217 /// training dynamics. Common schedules include warmup, cosine annealing,
218 /// and step decay.
219 ///
220 /// # Arguments
221 ///
222 /// * `base_optimizer` - Base optimizer to wrap with scheduling
223 /// * `schedule` - Learning rate schedule configuration
224 ///
225 /// # Examples
226 ///
227 /// ```rust,ignore
228 /// /// let base = AutoOptimizer::from_config(&config)?;
229 /// let schedule = LearningRateSchedule::LinearWarmup {
230 /// warmup_steps: 1000,
231 /// max_lr: 5e-5,
232 /// };
233 /// let scheduled = AutoOptimizer::with_schedule(base, schedule);
234
235 pub fn with_schedule(
236 base_optimizer: Box<dyn Optimizer>,
237 schedule: LearningRateSchedule,
238 ) -> ScheduledOptimizer {
239 ScheduledOptimizer::new(base_optimizer, schedule)
240 }
241}
242
243// =============================================================================
244// Base Optimizer Traits and Types
245// =============================================================================
246
247/// Core trait that all optimizers must implement
248///
249/// This trait defines the essential interface for gradient-based optimization,
250/// providing methods for parameter updates, state management, and learning
251/// rate control. All concrete optimizer implementations must provide these
252/// methods to ensure consistent behavior across the framework.
253pub trait Optimizer: Send + Sync + std::fmt::Debug {
254 /// Take an optimization step using provided gradients
255 ///
256 /// This is the core method that performs parameter updates based on
257 /// computed gradients. Implementations should update internal state
258 /// (momentum, variance estimates, etc.) and return parameter updates.
259 ///
260 /// Any gradients previously handed to [`Optimizer::accumulate_gradients`]
261 /// since the last [`Optimizer::zero_grad`] are folded in elementwise with
262 /// whatever is passed here, matching the usual "accumulate across
263 /// micro-batches, then step" training loop pattern.
264 ///
265 /// # Arguments
266 ///
267 /// * `gradients` - Gradients for all parameters to be updated
268 ///
269 /// # Returns
270 ///
271 /// Parameter updates that should be applied to model weights
272 ///
273 /// # Errors
274 ///
275 /// Returns an error if a gradient's length disagrees with the length of
276 /// the accumulated gradient (or restored moment state) for the same
277 /// parameter, rather than indexing out of bounds or silently truncating.
278 fn step(&mut self, gradients: &OptimizerGradients) -> Result<OptimizerUpdate>;
279
280 /// Accumulate gradients into an internal per-parameter buffer without
281 /// taking an optimization step.
282 ///
283 /// Calling this multiple times sums the gradients elementwise (the same
284 /// semantics as calling `.backward()` repeatedly without an intervening
285 /// `zero_grad()` in a typical autodiff-based trainer): a gradient
286 /// accumulation loop can call this once per micro-batch and then call
287 /// [`Optimizer::step`] once per effective batch.
288 ///
289 /// # Errors
290 ///
291 /// Returns an error if a parameter is accumulated at one length and then
292 /// accumulated again at a different length (e.g. the model shape
293 /// changed without an intervening [`Optimizer::zero_grad`]).
294 fn accumulate_gradients(&mut self, gradients: &OptimizerGradients) -> Result<()>;
295
296 /// Zero accumulated gradients
297 ///
298 /// Clears the buffer built up by [`Optimizer::accumulate_gradients`].
299 /// After this call, [`Optimizer::step`] uses only the gradients passed
300 /// to it directly, with nothing carried over from prior accumulation.
301 fn zero_grad(&mut self);
302
303 /// Get current learning rate
304 ///
305 /// Returns the current learning rate being used by the optimizer.
306 /// This may change over time when using learning rate schedules.
307 fn get_lr(&self) -> f64;
308
309 /// Set learning rate
310 ///
311 /// Updates the optimizer's learning rate. This is typically called
312 /// by learning rate schedulers or for manual learning rate adjustments.
313 ///
314 /// # Arguments
315 ///
316 /// * `lr` - New learning rate value
317 fn set_lr(&mut self, lr: f64);
318
319 /// Get optimizer state for serialization
320 ///
321 /// Returns a serializable representation of the optimizer's internal
322 /// state, including the first/second moment estimates (for optimizers
323 /// that have them) and step count. This enables saving and loading
324 /// optimizer state for training resumption.
325 ///
326 /// # Errors
327 ///
328 /// Returns an error rather than silently emitting JSON `null` when a
329 /// moment estimate holds a non-finite (`NaN`/`Infinity`) value — `null`
330 /// would round-trip back as `0.0`, hiding that the optimizer had
331 /// diverged.
332 fn state_dict(&self) -> Result<HashMap<String, serde_json::Value>>;
333
334 /// Load optimizer state from serialized data
335 ///
336 /// Restores the optimizer's internal state from previously saved data.
337 /// This is essential for resuming training from checkpoints.
338 ///
339 /// # Arguments
340 ///
341 /// * `state` - Serialized optimizer state
342 ///
343 /// # Errors
344 ///
345 /// Returns an error if a moment-estimate entry is present but is not a
346 /// JSON object of parameter name -> array-of-numbers, or contains a
347 /// non-numeric entry (including `null`, which a state dict produced by
348 /// an unguarded serializer could contain in place of a diverged value).
349 fn load_state_dict(&mut self, state: HashMap<String, serde_json::Value>) -> Result<()>;
350}
351
352/// Serialize a per-parameter moment-estimate map (`m` or `v`) to JSON,
353/// rejecting any non-finite value instead of letting `serde_json` silently
354/// turn it into `null`.
355///
356/// `serde_json::Number::from_f64` returns `None` for `NaN`/`Infinity`, and
357/// `serde_json::to_value` on such an `f32` therefore serializes it as JSON
358/// `null` with no error. A `null` in a moment estimate would round-trip back
359/// through [`moment_map_from_json`] as an error (good), but silently
360/// *skipping* the value at serialize time would be worse: it would make a
361/// diverged optimizer's checkpoint look like a healthy all-zero one. This
362/// helper fails loudly instead.
363fn moment_map_to_json(
364 moments: &HashMap<String, Vec<f32>>,
365 which: &str,
366) -> Result<serde_json::Value> {
367 let mut object = serde_json::Map::with_capacity(moments.len());
368 for (name, values) in moments {
369 let mut array = Vec::with_capacity(values.len());
370 for &value in values {
371 let number = serde_json::Number::from_f64(value as f64).ok_or_else(|| {
372 crate::error::TrustformersError::runtime_error(format!(
373 "optimizer state_dict: non-finite value in `{which}` moment estimate for \
374 parameter `{name}` (value = {value}); cannot serialize a diverged \
375 optimizer's state losslessly"
376 ))
377 })?;
378 array.push(serde_json::Value::Number(number));
379 }
380 object.insert(name.clone(), serde_json::Value::Array(array));
381 }
382 Ok(serde_json::Value::Object(object))
383}
384
385/// Inverse of [`moment_map_to_json`]. Returns an empty map when `value` is
386/// `None` (the key was absent from the state dict, e.g. a checkpoint saved
387/// before this field existed), and a structured error for anything present
388/// but malformed rather than silently defaulting missing entries to zero.
389fn moment_map_from_json(
390 value: Option<&serde_json::Value>,
391 which: &str,
392) -> Result<HashMap<String, Vec<f32>>> {
393 let mut result = HashMap::new();
394 let Some(value) = value else {
395 return Ok(result);
396 };
397 let object = value.as_object().ok_or_else(|| {
398 crate::error::TrustformersError::runtime_error(format!(
399 "optimizer load_state_dict: `{which}` must be a JSON object mapping parameter names \
400 to arrays of floats, got {value}"
401 ))
402 })?;
403 for (name, array_value) in object {
404 let array = array_value.as_array().ok_or_else(|| {
405 crate::error::TrustformersError::runtime_error(format!(
406 "optimizer load_state_dict: `{which}.{name}` must be a JSON array of floats, got \
407 {array_value}"
408 ))
409 })?;
410 let mut values = Vec::with_capacity(array.len());
411 for entry in array {
412 let f = entry.as_f64().ok_or_else(|| {
413 crate::error::TrustformersError::runtime_error(format!(
414 "optimizer load_state_dict: `{which}.{name}` contains a non-numeric entry \
415 ({entry}) -- a JSON `null` here usually means the checkpoint was written by \
416 a serializer that silently dropped a non-finite (NaN/Infinity) value"
417 ))
418 })?;
419 values.push(f as f32);
420 }
421 result.insert(name.clone(), values);
422 }
423 Ok(result)
424}
425
426/// Compute the effective per-parameter gradient for a `step()` call: the
427/// gradient passed to `step` plus whatever was accumulated via
428/// `accumulate_gradients` for the same parameter (elementwise sum), in
429/// stable order (parameters present only in `gradients`, then parameters
430/// present only in `accumulated`).
431///
432/// # Errors
433///
434/// Returns an error if a parameter appears in both maps with different
435/// lengths.
436fn merge_accumulated_gradients(
437 gradients: &OptimizerGradients,
438 accumulated: &HashMap<String, Vec<f32>>,
439) -> Result<HashMap<String, Vec<f32>>> {
440 let mut effective = HashMap::with_capacity(gradients.parameters.len().max(accumulated.len()));
441
442 for (name, passed) in &gradients.parameters {
443 match accumulated.get(name) {
444 Some(acc) => {
445 if acc.len() != passed.len() {
446 return Err(crate::error::TrustformersError::runtime_error(format!(
447 "optimizer step: accumulated gradient for `{name}` has {} values but the \
448 step's gradient has {} -- shapes must match (call zero_grad() if the \
449 model shape changed)",
450 acc.len(),
451 passed.len()
452 )));
453 }
454 effective.insert(
455 name.clone(),
456 passed.iter().zip(acc.iter()).map(|(g, a)| g + a).collect(),
457 );
458 },
459 None => {
460 effective.insert(name.clone(), passed.clone());
461 },
462 }
463 }
464 for (name, acc) in accumulated {
465 effective.entry(name.clone()).or_insert_with(|| acc.clone());
466 }
467
468 Ok(effective)
469}
470
471/// Accumulate `gradients` elementwise into `accumulated`, in place.
472///
473/// # Errors
474///
475/// Returns an error if a parameter was previously accumulated at a
476/// different length than the newly supplied gradient.
477fn accumulate_into(
478 accumulated: &mut HashMap<String, Vec<f32>>,
479 gradients: &OptimizerGradients,
480) -> Result<()> {
481 for (name, grad) in &gradients.parameters {
482 match accumulated.get_mut(name) {
483 Some(existing) => {
484 if existing.len() != grad.len() {
485 return Err(crate::error::TrustformersError::runtime_error(format!(
486 "optimizer accumulate_gradients: `{name}` was previously accumulated at \
487 {} values, new gradient has {} -- call zero_grad() before changing \
488 parameter shape",
489 existing.len(),
490 grad.len()
491 )));
492 }
493 for (acc, g) in existing.iter_mut().zip(grad.iter()) {
494 *acc += g;
495 }
496 },
497 None => {
498 accumulated.insert(name.clone(), grad.clone());
499 },
500 }
501 }
502 Ok(())
503}
504
505/// Guard against a restored (or freshly initialized) moment-estimate vector
506/// whose length disagrees with the current effective gradient -- indexing
507/// into a mismatched vector would otherwise panic instead of erroring.
508fn ensure_moment_len(
509 moment: &[f32],
510 expected_len: usize,
511 which: &str,
512 param_name: &str,
513) -> Result<()> {
514 if moment.len() != expected_len {
515 return Err(crate::error::TrustformersError::runtime_error(format!(
516 "optimizer step: restored `{which}` state for `{param_name}` has {} entries but the \
517 gradient has {expected_len} -- this checkpoint does not match the current model shape",
518 moment.len()
519 )));
520 }
521 Ok(())
522}
523
524/// Container for gradients during optimization
525///
526/// This structure holds gradients for all model parameters along with
527/// their shapes, enabling efficient gradient-based optimization across
528/// parameters of different dimensions.
529#[derive(Debug, Clone)]
530pub struct OptimizerGradients {
531 /// Flattened gradients for each named parameter
532 pub parameters: HashMap<String, Vec<f32>>,
533 /// Original shapes of parameters for reconstruction
534 pub parameter_shapes: HashMap<String, Vec<usize>>,
535}
536
537/// Container for parameter updates from optimization step
538///
539/// This structure contains the computed parameter updates along with
540/// metadata about the optimization step, such as the effective learning
541/// rate and step count.
542#[derive(Debug, Clone)]
543pub struct OptimizerUpdate {
544 /// Parameter updates to be applied to model weights
545 pub parameter_updates: HashMap<String, Vec<f32>>,
546 /// Learning rate used for this step
547 pub learning_rate: f64,
548 /// Current step count for tracking training progress
549 pub step_count: usize,
550}
551
552/// Learning rate scheduling strategies
553///
554/// Different learning rate schedules can significantly impact training
555/// dynamics and final model performance. This enum provides common
556/// scheduling strategies used in modern deep learning.
557#[derive(Debug, Clone)]
558pub enum LearningRateSchedule {
559 /// Constant learning rate throughout training
560 Constant,
561
562 /// Linear warmup to a maximum learning rate
563 ///
564 /// Gradually increases learning rate from initial value to max_lr
565 /// over warmup_steps, then maintains max_lr
566 LinearWarmup { warmup_steps: usize, max_lr: f64 },
567
568 /// Cosine annealing schedule
569 ///
570 /// Follows a cosine curve from initial learning rate down to eta_min
571 /// over t_max steps, providing smooth learning rate decay
572 CosineAnnealing { t_max: usize, eta_min: f64 },
573
574 /// Step-wise learning rate decay
575 ///
576 /// Multiplies learning rate by gamma every step_size steps,
577 /// providing periodic learning rate reductions
578 StepLR { step_size: usize, gamma: f64 },
579
580 /// Polynomial learning rate decay
581 ///
582 /// Smoothly decays learning rate from initial value to end_lr
583 /// following a polynomial curve with specified power
584 PolynomialDecay {
585 power: f64,
586 end_lr: f64,
587 total_steps: usize,
588 },
589}
590
591// =============================================================================
592// Concrete Optimizer Implementations
593// =============================================================================
594//
595// NOTE: These implementations are currently included in this module for
596// completeness, but should be refactored into separate files as the
597// optimizer system grows:
598//
599// - adamw.rs: AdamW optimizer implementation
600// - adam.rs: Adam optimizer implementation
601// - sgd.rs: SGD with momentum implementation
602// - scheduled.rs: Learning rate scheduling wrapper
603// - lamb.rs: LAMB optimizer for large batch training
604// - adafactor.rs: Memory-efficient Adafactor optimizer
605//
606// This modular structure will improve maintainability and allow for
607// easier testing and documentation of individual optimizers.
608
609/// AdamW optimizer implementation
610///
611/// AdamW (Adam with decoupled Weight decay) is a variant of Adam that
612/// separates weight decay from gradient-based optimization, leading to
613/// better generalization in many scenarios, especially for transformer models.
614///
615/// The key difference from Adam is that weight decay is applied directly
616/// to parameters rather than being included in the gradient computation,
617/// which provides more consistent regularization behavior.
618#[derive(Debug, Clone)]
619pub struct AdamWOptimizer {
620 config: AdamWConfig,
621 step_count: usize,
622 m: HashMap<String, Vec<f32>>, // First moment estimates
623 v: HashMap<String, Vec<f32>>, // Second moment estimates
624 /// Per-parameter running maximum of `v`, used only when
625 /// `config.amsgrad` is set (Reddi et al., 2018). Kept separate from `v`
626 /// so state_dict() can omit it for the common non-AMSGrad case.
627 v_max: HashMap<String, Vec<f32>>,
628 /// Gradients accumulated via [`Optimizer::accumulate_gradients`] since
629 /// the last [`Optimizer::zero_grad`]; folded into the next [`Optimizer::step`].
630 accumulated_gradients: HashMap<String, Vec<f32>>,
631}
632
633/// Configuration for AdamW optimizer
634#[derive(Debug, Clone)]
635pub struct AdamWConfig {
636 /// Learning rate (alpha)
637 pub learning_rate: f64,
638 /// Exponential decay rate for first moment estimates
639 pub beta1: f64,
640 /// Exponential decay rate for second moment estimates
641 pub beta2: f64,
642 /// Weight decay coefficient for regularization
643 pub weight_decay: f64,
644 /// Small constant for numerical stability
645 pub eps: f64,
646 /// Whether to use AMSGrad variant
647 pub amsgrad: bool,
648}
649
650impl AdamWOptimizer {
651 /// Create new AdamW optimizer with given configuration
652 pub fn new(config: AdamWConfig) -> Self {
653 Self {
654 config,
655 step_count: 0,
656 m: HashMap::new(),
657 v: HashMap::new(),
658 v_max: HashMap::new(),
659 accumulated_gradients: HashMap::new(),
660 }
661 }
662}
663
664impl Optimizer for AdamWOptimizer {
665 fn step(&mut self, gradients: &OptimizerGradients) -> Result<OptimizerUpdate> {
666 let effective_gradients =
667 merge_accumulated_gradients(gradients, &self.accumulated_gradients)?;
668 self.step_count += 1;
669 let mut parameter_updates = HashMap::new();
670
671 for (param_name, grad) in &effective_gradients {
672 // Initialize moment estimates if needed (entry API avoids a fallible lookup)
673 let m = self.m.entry(param_name.clone()).or_insert_with(|| vec![0.0; grad.len()]);
674 ensure_moment_len(m, grad.len(), "m", param_name)?;
675 let v = self.v.entry(param_name.clone()).or_insert_with(|| vec![0.0; grad.len()]);
676 ensure_moment_len(v, grad.len(), "v", param_name)?;
677 let v_max = if self.config.amsgrad {
678 let entry =
679 self.v_max.entry(param_name.clone()).or_insert_with(|| vec![0.0; grad.len()]);
680 ensure_moment_len(entry, grad.len(), "v_max", param_name)?;
681 Some(entry)
682 } else {
683 None
684 };
685
686 let mut updates = Vec::with_capacity(grad.len());
687 let mut v_max = v_max;
688
689 for i in 0..grad.len() {
690 // Update biased first moment estimate
691 m[i] = self.config.beta1 as f32 * m[i] + (1.0 - self.config.beta1 as f32) * grad[i];
692
693 // Update biased second raw moment estimate
694 v[i] = self.config.beta2 as f32 * v[i]
695 + (1.0 - self.config.beta2 as f32) * grad[i] * grad[i];
696
697 // Compute bias-corrected first moment estimate
698 let m_hat = m[i] / (1.0 - (self.config.beta1 as f32).powi(self.step_count as i32));
699
700 // Compute bias-corrected second raw moment estimate. Under
701 // AMSGrad (Reddi et al., 2018) the denominator uses the
702 // running *maximum* of `v_hat`'s numerator instead of the
703 // current `v[i]`, which prevents the effective learning rate
704 // from increasing late in training and fixes Adam's
705 // non-convergence counterexample.
706 let v_for_denom = if let Some(v_max) = v_max.as_deref_mut() {
707 v_max[i] = v_max[i].max(v[i]);
708 v_max[i]
709 } else {
710 v[i]
711 };
712 let v_hat =
713 v_for_denom / (1.0 - (self.config.beta2 as f32).powi(self.step_count as i32));
714
715 // AdamW-style decoupled weight decay: `weight_decay` shrinks
716 // the parameter directly (scaled by the learning rate, as in
717 // Loshchilov & Hutter, 2019), rather than being folded into
718 // the gradient the way plain L2 regularization would be. It
719 // is therefore added on top of the raw Adam update rather
720 // than mixed into `grad[i]` above.
721 let adam_update = -self.config.learning_rate as f32 * m_hat
722 / (v_hat.sqrt() + self.config.eps as f32);
723 let decay_update =
724 -self.config.learning_rate as f32 * self.config.weight_decay as f32;
725 updates.push(adam_update + decay_update);
726 }
727
728 parameter_updates.insert(param_name.clone(), updates);
729 }
730
731 Ok(OptimizerUpdate {
732 parameter_updates,
733 learning_rate: self.config.learning_rate,
734 step_count: self.step_count,
735 })
736 }
737
738 fn accumulate_gradients(&mut self, gradients: &OptimizerGradients) -> Result<()> {
739 accumulate_into(&mut self.accumulated_gradients, gradients)
740 }
741
742 fn zero_grad(&mut self) {
743 self.accumulated_gradients.clear();
744 }
745
746 fn get_lr(&self) -> f64 {
747 self.config.learning_rate
748 }
749
750 fn set_lr(&mut self, lr: f64) {
751 self.config.learning_rate = lr;
752 }
753
754 fn state_dict(&self) -> Result<HashMap<String, serde_json::Value>> {
755 let mut state = HashMap::new();
756 state.insert(
757 "step_count".to_string(),
758 serde_json::Value::Number(self.step_count.into()),
759 );
760 state.insert(
761 "learning_rate".to_string(),
762 serde_json::Number::from_f64(self.config.learning_rate)
763 .map(serde_json::Value::Number)
764 .unwrap_or_else(|| {
765 serde_json::Value::String(format!("{}", self.config.learning_rate))
766 }),
767 );
768 state.insert("m".to_string(), moment_map_to_json(&self.m, "m")?);
769 state.insert("v".to_string(), moment_map_to_json(&self.v, "v")?);
770 if self.config.amsgrad {
771 state.insert(
772 "v_max".to_string(),
773 moment_map_to_json(&self.v_max, "v_max")?,
774 );
775 }
776 Ok(state)
777 }
778
779 fn load_state_dict(&mut self, state: HashMap<String, serde_json::Value>) -> Result<()> {
780 if let Some(step_count) = state.get("step_count").and_then(|v| v.as_u64()) {
781 self.step_count = step_count as usize;
782 }
783 if let Some(lr) = state.get("learning_rate").and_then(|v| v.as_f64()) {
784 self.config.learning_rate = lr;
785 }
786 self.m = moment_map_from_json(state.get("m"), "m")?;
787 self.v = moment_map_from_json(state.get("v"), "v")?;
788 self.v_max = moment_map_from_json(state.get("v_max"), "v_max")?;
789 Ok(())
790 }
791}
792
793/// Adam optimizer implementation
794///
795/// The classic Adam (Adaptive Moment Estimation) optimizer that adapts
796/// learning rates for each parameter based on first and second moment
797/// estimates of gradients. Works well for many tasks but can sometimes
798/// suffer from poor generalization compared to AdamW.
799#[derive(Debug, Clone)]
800pub struct AdamOptimizer {
801 config: AdamConfig,
802 step_count: usize,
803 m: HashMap<String, Vec<f32>>, // First moment estimates
804 v: HashMap<String, Vec<f32>>, // Second moment estimates
805 /// Per-parameter running maximum of `v`, used only when
806 /// `config.amsgrad` is set (Reddi et al., 2018).
807 v_max: HashMap<String, Vec<f32>>,
808 /// Gradients accumulated via [`Optimizer::accumulate_gradients`] since
809 /// the last [`Optimizer::zero_grad`]; folded into the next [`Optimizer::step`].
810 accumulated_gradients: HashMap<String, Vec<f32>>,
811}
812
813/// Configuration for Adam optimizer
814#[derive(Debug, Clone)]
815pub struct AdamConfig {
816 /// Learning rate (alpha)
817 pub learning_rate: f64,
818 /// Exponential decay rate for first moment estimates
819 pub beta1: f64,
820 /// Exponential decay rate for second moment estimates
821 pub beta2: f64,
822 /// Small constant for numerical stability
823 pub eps: f64,
824 /// Whether to use AMSGrad variant
825 pub amsgrad: bool,
826}
827
828impl AdamOptimizer {
829 /// Create new Adam optimizer with given configuration
830 pub fn new(config: AdamConfig) -> Self {
831 Self {
832 config,
833 step_count: 0,
834 m: HashMap::new(),
835 v: HashMap::new(),
836 v_max: HashMap::new(),
837 accumulated_gradients: HashMap::new(),
838 }
839 }
840}
841
842impl Optimizer for AdamOptimizer {
843 fn step(&mut self, gradients: &OptimizerGradients) -> Result<OptimizerUpdate> {
844 // Similar to AdamW but without weight decay
845 let effective_gradients =
846 merge_accumulated_gradients(gradients, &self.accumulated_gradients)?;
847 self.step_count += 1;
848 let mut parameter_updates = HashMap::new();
849
850 for (param_name, grad) in &effective_gradients {
851 // Initialize moment estimates if needed (entry API avoids a fallible lookup)
852 let m = self.m.entry(param_name.clone()).or_insert_with(|| vec![0.0; grad.len()]);
853 ensure_moment_len(m, grad.len(), "m", param_name)?;
854 let v = self.v.entry(param_name.clone()).or_insert_with(|| vec![0.0; grad.len()]);
855 ensure_moment_len(v, grad.len(), "v", param_name)?;
856 let v_max = if self.config.amsgrad {
857 let entry =
858 self.v_max.entry(param_name.clone()).or_insert_with(|| vec![0.0; grad.len()]);
859 ensure_moment_len(entry, grad.len(), "v_max", param_name)?;
860 Some(entry)
861 } else {
862 None
863 };
864
865 let mut updates = Vec::with_capacity(grad.len());
866 let mut v_max = v_max;
867
868 for i in 0..grad.len() {
869 m[i] = self.config.beta1 as f32 * m[i] + (1.0 - self.config.beta1 as f32) * grad[i];
870 v[i] = self.config.beta2 as f32 * v[i]
871 + (1.0 - self.config.beta2 as f32) * grad[i] * grad[i];
872
873 let m_hat = m[i] / (1.0 - (self.config.beta1 as f32).powi(self.step_count as i32));
874 // See `AdamWOptimizer::step` for why AMSGrad uses a running
875 // maximum of `v` in the denominator instead of `v` itself.
876 let v_for_denom = if let Some(v_max) = v_max.as_deref_mut() {
877 v_max[i] = v_max[i].max(v[i]);
878 v_max[i]
879 } else {
880 v[i]
881 };
882 let v_hat =
883 v_for_denom / (1.0 - (self.config.beta2 as f32).powi(self.step_count as i32));
884
885 let update = -self.config.learning_rate as f32 * m_hat
886 / (v_hat.sqrt() + self.config.eps as f32);
887 updates.push(update);
888 }
889
890 parameter_updates.insert(param_name.clone(), updates);
891 }
892
893 Ok(OptimizerUpdate {
894 parameter_updates,
895 learning_rate: self.config.learning_rate,
896 step_count: self.step_count,
897 })
898 }
899
900 fn accumulate_gradients(&mut self, gradients: &OptimizerGradients) -> Result<()> {
901 accumulate_into(&mut self.accumulated_gradients, gradients)
902 }
903
904 fn zero_grad(&mut self) {
905 self.accumulated_gradients.clear();
906 }
907
908 fn get_lr(&self) -> f64 {
909 self.config.learning_rate
910 }
911
912 fn set_lr(&mut self, lr: f64) {
913 self.config.learning_rate = lr;
914 }
915
916 fn state_dict(&self) -> Result<HashMap<String, serde_json::Value>> {
917 let mut state = HashMap::new();
918 state.insert(
919 "step_count".to_string(),
920 serde_json::Value::Number(self.step_count.into()),
921 );
922 state.insert(
923 "learning_rate".to_string(),
924 serde_json::Number::from_f64(self.config.learning_rate)
925 .map(serde_json::Value::Number)
926 .unwrap_or_else(|| {
927 serde_json::Value::String(format!("{}", self.config.learning_rate))
928 }),
929 );
930 state.insert("m".to_string(), moment_map_to_json(&self.m, "m")?);
931 state.insert("v".to_string(), moment_map_to_json(&self.v, "v")?);
932 if self.config.amsgrad {
933 state.insert(
934 "v_max".to_string(),
935 moment_map_to_json(&self.v_max, "v_max")?,
936 );
937 }
938 Ok(state)
939 }
940
941 fn load_state_dict(&mut self, state: HashMap<String, serde_json::Value>) -> Result<()> {
942 if let Some(step_count) = state.get("step_count").and_then(|v| v.as_u64()) {
943 self.step_count = step_count as usize;
944 }
945 if let Some(lr) = state.get("learning_rate").and_then(|v| v.as_f64()) {
946 self.config.learning_rate = lr;
947 }
948 self.m = moment_map_from_json(state.get("m"), "m")?;
949 self.v = moment_map_from_json(state.get("v"), "v")?;
950 self.v_max = moment_map_from_json(state.get("v_max"), "v_max")?;
951 Ok(())
952 }
953}
954
955/// Optimizer wrapper that applies learning rate scheduling
956///
957/// This wrapper can be applied to any base optimizer to provide dynamic
958/// learning rate adjustment during training. Different schedules can
959/// significantly impact convergence speed and final model quality.
960#[derive(Debug)]
961pub struct ScheduledOptimizer {
962 optimizer: Box<dyn Optimizer>,
963 schedule: LearningRateSchedule,
964 initial_lr: f64,
965 current_step: usize,
966}
967
968impl ScheduledOptimizer {
969 /// Create new scheduled optimizer
970 ///
971 /// # Arguments
972 ///
973 /// * `optimizer` - Base optimizer to wrap with scheduling
974 /// * `schedule` - Learning rate schedule to apply
975 pub fn new(optimizer: Box<dyn Optimizer>, schedule: LearningRateSchedule) -> Self {
976 let initial_lr = optimizer.get_lr();
977 Self {
978 optimizer,
979 schedule,
980 initial_lr,
981 current_step: 0,
982 }
983 }
984
985 /// Update learning rate based on current step and schedule
986 fn update_learning_rate(&mut self) {
987 let new_lr = match &self.schedule {
988 LearningRateSchedule::Constant => self.initial_lr,
989 LearningRateSchedule::LinearWarmup {
990 warmup_steps,
991 max_lr,
992 } => {
993 if self.current_step < *warmup_steps {
994 self.initial_lr
995 + (max_lr - self.initial_lr)
996 * (self.current_step as f64 / *warmup_steps as f64)
997 } else {
998 *max_lr
999 }
1000 },
1001 LearningRateSchedule::CosineAnnealing { t_max, eta_min } => {
1002 eta_min
1003 + (self.initial_lr - eta_min)
1004 * (1.0
1005 + (std::f64::consts::PI * self.current_step as f64 / *t_max as f64)
1006 .cos())
1007 / 2.0
1008 },
1009 LearningRateSchedule::StepLR { step_size, gamma } => {
1010 self.initial_lr * gamma.powi((self.current_step / step_size) as i32)
1011 },
1012 LearningRateSchedule::PolynomialDecay {
1013 power,
1014 end_lr,
1015 total_steps,
1016 } => {
1017 if self.current_step >= *total_steps {
1018 *end_lr
1019 } else {
1020 let decay_factor =
1021 (1.0 - self.current_step as f64 / *total_steps as f64).powf(*power);
1022 end_lr + (self.initial_lr - end_lr) * decay_factor
1023 }
1024 },
1025 };
1026
1027 self.optimizer.set_lr(new_lr);
1028 }
1029}
1030
1031impl Optimizer for ScheduledOptimizer {
1032 fn step(&mut self, gradients: &OptimizerGradients) -> Result<OptimizerUpdate> {
1033 self.current_step += 1;
1034 self.update_learning_rate();
1035 self.optimizer.step(gradients)
1036 }
1037
1038 fn accumulate_gradients(&mut self, gradients: &OptimizerGradients) -> Result<()> {
1039 self.optimizer.accumulate_gradients(gradients)
1040 }
1041
1042 fn zero_grad(&mut self) {
1043 self.optimizer.zero_grad();
1044 }
1045
1046 fn get_lr(&self) -> f64 {
1047 self.optimizer.get_lr()
1048 }
1049
1050 fn set_lr(&mut self, lr: f64) {
1051 self.initial_lr = lr;
1052 self.optimizer.set_lr(lr);
1053 }
1054
1055 fn state_dict(&self) -> Result<HashMap<String, serde_json::Value>> {
1056 let mut state = self.optimizer.state_dict()?;
1057 state.insert(
1058 "current_step".to_string(),
1059 serde_json::Value::Number(self.current_step.into()),
1060 );
1061 state.insert(
1062 "initial_lr".to_string(),
1063 serde_json::Number::from_f64(self.initial_lr)
1064 .map(serde_json::Value::Number)
1065 .unwrap_or_else(|| serde_json::Value::String(format!("{}", self.initial_lr))),
1066 );
1067 Ok(state)
1068 }
1069
1070 fn load_state_dict(&mut self, mut state: HashMap<String, serde_json::Value>) -> Result<()> {
1071 if let Some(step) = state.remove("current_step").and_then(|v| v.as_u64()) {
1072 self.current_step = step as usize;
1073 }
1074 if let Some(lr) = state.remove("initial_lr").and_then(|v| v.as_f64()) {
1075 self.initial_lr = lr;
1076 }
1077 self.optimizer.load_state_dict(state)
1078 }
1079}
1080
1081// =============================================================================
1082// Public API
1083// =============================================================================
1084
1085// All main components are already public and available for import:
1086// - AutoOptimizer: Main entry point for automatic optimizer creation
1087// - Optimizer: Base trait for all optimizers
1088// - OptimizerGradients/OptimizerUpdate: Data structures for optimization
1089// - LearningRateSchedule: Learning rate scheduling strategies
1090// - AdamWOptimizer/AdamOptimizer: Concrete optimizer implementations
1091// - ScheduledOptimizer: Optimizer wrapper with learning rate scheduling
1092
1093#[cfg(test)]
1094mod tests;