optirs_learned/transformer_based_optimizer/
config.rs

1// Configuration structures for transformer-based optimizer
2
3use super::positional_encoding::PositionalEncodingType;
4use scirs2_core::numeric::Float;
5use serde::{Deserialize, Serialize};
6use std::fmt::Debug;
7
8/// Configuration for transformer-based optimizer
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct TransformerBasedOptimizerConfig<
11    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
12> {
13    /// Model dimension (embedding size)
14    pub model_dimension: usize,
15
16    /// Number of transformer layers
17    pub num_transformer_layers: usize,
18
19    /// Number of attention heads
20    pub num_attention_heads: usize,
21
22    /// Dimension of each attention head
23    pub attention_head_dimension: usize,
24
25    /// Feed-forward network dimension
26    pub feedforward_dimension: usize,
27
28    /// Maximum sequence length
29    pub sequence_length: usize,
30
31    /// Dropout rate
32    pub dropout_rate: f64,
33
34    /// Learning rate
35    pub learning_rate: T,
36
37    /// Batch size for training
38    pub batch_size: usize,
39
40    /// Number of training epochs
41    pub num_epochs: usize,
42
43    /// Activation function type
44    pub activation_function: ActivationFunction,
45
46    /// Positional encoding type
47    pub positional_encoding_type: PositionalEncodingType,
48
49    /// Memory management configuration
50    pub memory_config: MemoryConfig,
51
52    /// Meta-learning configuration
53    pub meta_learning_config: MetaLearningConfig<T>,
54
55    /// Performance tracking configuration
56    pub performance_config: PerformanceConfig,
57
58    /// Enable gradient clipping
59    pub enable_gradient_clipping: bool,
60
61    /// Gradient clipping threshold
62    pub gradient_clip_value: T,
63
64    /// Weight decay factor
65    pub weight_decay: T,
66
67    /// Warmup steps for learning rate schedule
68    pub warmup_steps: usize,
69
70    /// Enable layer normalization
71    pub enable_layer_norm: bool,
72
73    /// Pre-norm vs post-norm
74    pub use_pre_norm: bool,
75
76    /// Enable residual connections
77    pub enable_residual_connections: bool,
78}
79
80impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static> Default
81    for TransformerBasedOptimizerConfig<T>
82{
83    fn default() -> Self {
84        Self {
85            model_dimension: 512,
86            num_transformer_layers: 6,
87            num_attention_heads: 8,
88            attention_head_dimension: 64,
89            feedforward_dimension: 2048,
90            sequence_length: 128,
91            dropout_rate: 0.1,
92            learning_rate: scirs2_core::numeric::NumCast::from(1e-4).unwrap_or_else(|| T::zero()),
93            batch_size: 32,
94            num_epochs: 100,
95            activation_function: ActivationFunction::ReLU,
96            positional_encoding_type: PositionalEncodingType::Sinusoidal,
97            memory_config: MemoryConfig::default(),
98            meta_learning_config: MetaLearningConfig::default(),
99            performance_config: PerformanceConfig::default(),
100            enable_gradient_clipping: true,
101            gradient_clip_value: scirs2_core::numeric::NumCast::from(1.0)
102                .unwrap_or_else(|| T::zero()),
103            weight_decay: scirs2_core::numeric::NumCast::from(1e-5).unwrap_or_else(|| T::zero()),
104            warmup_steps: 1000,
105            enable_layer_norm: true,
106            use_pre_norm: true,
107            enable_residual_connections: true,
108        }
109    }
110}
111
112impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
113    TransformerBasedOptimizerConfig<T>
114{
115    /// Create configuration for small models
116    pub fn small() -> Self {
117        Self {
118            model_dimension: 256,
119            num_transformer_layers: 4,
120            num_attention_heads: 4,
121            attention_head_dimension: 64,
122            feedforward_dimension: 1024,
123            sequence_length: 64,
124            ..Self::default()
125        }
126    }
127
128    /// Create configuration for large models
129    pub fn large() -> Self {
130        Self {
131            model_dimension: 1024,
132            num_transformer_layers: 12,
133            num_attention_heads: 16,
134            attention_head_dimension: 64,
135            feedforward_dimension: 4096,
136            sequence_length: 256,
137            ..Self::default()
138        }
139    }
140
141    /// Create configuration optimized for training
142    pub fn for_training() -> Self {
143        Self {
144            batch_size: 64,
145            num_epochs: 200,
146            learning_rate: scirs2_core::numeric::NumCast::from(2e-4).unwrap_or_else(|| T::zero()),
147            warmup_steps: 2000,
148            enable_gradient_clipping: true,
149            weight_decay: scirs2_core::numeric::NumCast::from(1e-4).unwrap_or_else(|| T::zero()),
150            ..Self::default()
151        }
152    }
153
154    /// Create configuration optimized for inference
155    pub fn for_inference() -> Self {
156        Self {
157            batch_size: 1,
158            dropout_rate: 0.0,
159            enable_gradient_clipping: false,
160            ..Self::default()
161        }
162    }
163
164    /// Validate configuration
165    pub fn validate(&self) -> Result<(), String> {
166        if self.model_dimension == 0 {
167            return Err("model_dimension must be greater than 0".to_string());
168        }
169
170        if self.num_transformer_layers == 0 {
171            return Err("num_transformer_layers must be greater than 0".to_string());
172        }
173
174        if self.num_attention_heads == 0 {
175            return Err("num_attention_heads must be greater than 0".to_string());
176        }
177
178        if !self
179            .model_dimension
180            .is_multiple_of(self.num_attention_heads)
181        {
182            return Err("model_dimension must be divisible by num_attention_heads".to_string());
183        }
184
185        if self.attention_head_dimension * self.num_attention_heads != self.model_dimension {
186            return Err(
187                "attention_head_dimension * num_attention_heads must equal model_dimension"
188                    .to_string(),
189            );
190        }
191
192        if self.sequence_length == 0 {
193            return Err("sequence_length must be greater than 0".to_string());
194        }
195
196        if self.dropout_rate < 0.0 || self.dropout_rate > 1.0 {
197            return Err("dropout_rate must be between 0.0 and 1.0".to_string());
198        }
199
200        if self.learning_rate <= T::zero() {
201            return Err("learning_rate must be positive".to_string());
202        }
203
204        if self.batch_size == 0 {
205            return Err("batch_size must be greater than 0".to_string());
206        }
207
208        self.memory_config.validate()?;
209        self.meta_learning_config.validate()?;
210        self.performance_config.validate()?;
211
212        Ok(())
213    }
214
215    /// Calculate total parameters estimate
216    pub fn estimate_parameter_count(&self) -> usize {
217        let embedding_params = self.model_dimension * self.model_dimension; // Input embedding
218        let positional_params = self.sequence_length * self.model_dimension;
219
220        let attention_params_per_layer = 4 * self.model_dimension * self.model_dimension; // Q, K, V, O projections
221        let ffn_params_per_layer = 2 * self.model_dimension * self.feedforward_dimension; // Up and down projections
222        let norm_params_per_layer = 2 * self.model_dimension; // Layer norm parameters
223
224        let layer_params =
225            attention_params_per_layer + ffn_params_per_layer + norm_params_per_layer;
226        let total_layer_params = layer_params * self.num_transformer_layers;
227
228        let output_params = self.model_dimension * self.model_dimension; // Output projection
229
230        embedding_params + positional_params + total_layer_params + output_params
231    }
232
233    /// Calculate memory requirements (in MB)
234    pub fn estimate_memory_usage(&self) -> f64 {
235        let param_count = self.estimate_parameter_count();
236        let bytes_per_param = if std::mem::size_of::<T>() == 4 {
237            4.0
238        } else {
239            8.0
240        };
241
242        let model_memory = param_count as f64 * bytes_per_param;
243        let activation_memory = self.batch_size as f64
244            * self.sequence_length as f64
245            * self.model_dimension as f64
246            * bytes_per_param;
247        let gradient_memory = model_memory; // Assume same as model for gradients
248
249        let total_bytes = model_memory + activation_memory + gradient_memory;
250        total_bytes / (1024.0 * 1024.0) // Convert to MB
251    }
252}
253
254/// Transformer architecture configuration
255#[derive(Debug, Clone)]
256pub struct TransformerArchConfig {
257    pub model_dimension: usize,
258    pub num_layers: usize,
259    pub num_attention_heads: usize,
260    pub feedforward_dimension: usize,
261    pub dropout_rate: f64,
262    pub use_pre_norm: bool,
263    pub enable_residual_connections: bool,
264}
265
266impl TransformerArchConfig {
267    pub fn from_optimizer_config<
268        T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
269    >(
270        config: &TransformerBasedOptimizerConfig<T>,
271    ) -> Self {
272        Self {
273            model_dimension: config.model_dimension,
274            num_layers: config.num_transformer_layers,
275            num_attention_heads: config.num_attention_heads,
276            feedforward_dimension: config.feedforward_dimension,
277            dropout_rate: config.dropout_rate,
278            use_pre_norm: config.use_pre_norm,
279            enable_residual_connections: config.enable_residual_connections,
280        }
281    }
282}
283
284/// Activation function types
285#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
286pub enum ActivationFunction {
287    ReLU,
288    GELU,
289    Swish,
290    Tanh,
291    Sigmoid,
292    LeakyReLU,
293}
294
295/// Memory management configuration
296#[derive(Debug, Clone, Serialize, Deserialize)]
297pub struct MemoryConfig {
298    /// Maximum memory cache size
299    pub max_cache_size: usize,
300    /// Enable memory compression
301    pub enable_compression: bool,
302    /// Cache eviction strategy
303    pub eviction_strategy: CacheEvictionStrategy,
304    /// Memory allocation block size
305    pub allocation_block_size: usize,
306}
307
308impl Default for MemoryConfig {
309    fn default() -> Self {
310        Self {
311            max_cache_size: 1024 * 1024 * 1024, // 1GB
312            enable_compression: false,
313            eviction_strategy: CacheEvictionStrategy::LRU,
314            allocation_block_size: 4096,
315        }
316    }
317}
318
319impl MemoryConfig {
320    pub fn validate(&self) -> Result<(), String> {
321        if self.max_cache_size == 0 {
322            return Err("max_cache_size must be greater than 0".to_string());
323        }
324
325        if self.allocation_block_size == 0 {
326            return Err("allocation_block_size must be greater than 0".to_string());
327        }
328
329        Ok(())
330    }
331}
332
333/// Cache eviction strategies
334#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
335pub enum CacheEvictionStrategy {
336    LRU,
337    LFU,
338    FIFO,
339    Random,
340}
341
342/// Meta-learning configuration
343#[derive(Debug, Clone, Serialize, Deserialize)]
344pub struct MetaLearningConfig<
345    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
346> {
347    /// Meta-learning rate
348    pub meta_learning_rate: T,
349    /// Number of inner optimization steps
350    pub inner_steps: usize,
351    /// Inner learning rate
352    pub inner_learning_rate: T,
353    /// Enable first-order approximation
354    pub first_order: bool,
355    /// Number of support examples
356    pub num_support: usize,
357    /// Number of query examples
358    pub num_query: usize,
359    /// Heavy-ball momentum coefficient for the **outer** (meta) loop.
360    ///
361    /// `None` (the default) means plain SGD. Setting it activates the momentum
362    /// branch of `MetaOptimizer::update`, which was previously unreachable
363    /// because the momentum field was hardcoded to `None`.
364    pub meta_momentum: Option<T>,
365}
366
367impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static> Default
368    for MetaLearningConfig<T>
369{
370    fn default() -> Self {
371        Self {
372            meta_learning_rate: scirs2_core::numeric::NumCast::from(1e-3)
373                .unwrap_or_else(|| T::zero()),
374            inner_steps: 5,
375            inner_learning_rate: scirs2_core::numeric::NumCast::from(1e-2)
376                .unwrap_or_else(|| T::zero()),
377            first_order: false,
378            num_support: 5,
379            num_query: 15,
380            meta_momentum: None,
381        }
382    }
383}
384
385impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
386    MetaLearningConfig<T>
387{
388    pub fn validate(&self) -> Result<(), String> {
389        if self.meta_learning_rate <= T::zero() {
390            return Err("meta_learning_rate must be positive".to_string());
391        }
392
393        if self.inner_learning_rate <= T::zero() {
394            return Err("inner_learning_rate must be positive".to_string());
395        }
396
397        if self.inner_steps == 0 {
398            return Err("inner_steps must be greater than 0".to_string());
399        }
400
401        if self.num_support == 0 {
402            return Err("num_support must be greater than 0".to_string());
403        }
404
405        if self.num_query == 0 {
406            return Err("num_query must be greater than 0".to_string());
407        }
408
409        Ok(())
410    }
411}
412
413/// Performance tracking configuration
414#[derive(Debug, Clone, Serialize, Deserialize)]
415pub struct PerformanceConfig {
416    /// Enable detailed performance tracking
417    pub enable_detailed_tracking: bool,
418    /// Performance metrics collection interval
419    pub metrics_interval: usize,
420    /// Maximum history size for metrics
421    pub max_history_size: usize,
422    /// Enable memory usage tracking
423    pub track_memory_usage: bool,
424    /// Enable timing analysis
425    pub enable_timing_analysis: bool,
426}
427
428impl Default for PerformanceConfig {
429    fn default() -> Self {
430        Self {
431            enable_detailed_tracking: true,
432            metrics_interval: 10,
433            max_history_size: 10000,
434            track_memory_usage: true,
435            enable_timing_analysis: true,
436        }
437    }
438}
439
440impl PerformanceConfig {
441    pub fn validate(&self) -> Result<(), String> {
442        if self.metrics_interval == 0 {
443            return Err("metrics_interval must be greater than 0".to_string());
444        }
445
446        if self.max_history_size == 0 {
447            return Err("max_history_size must be greater than 0".to_string());
448        }
449
450        Ok(())
451    }
452}
453
454#[cfg(test)]
455mod tests {
456    use super::*;
457
458    #[test]
459    fn test_config_default() {
460        let config = TransformerBasedOptimizerConfig::<f32>::default();
461        assert!(config.validate().is_ok());
462        assert_eq!(config.model_dimension, 512);
463        assert_eq!(config.num_transformer_layers, 6);
464    }
465
466    #[test]
467    fn test_config_validation() {
468        // Test invalid model dimension
469        let config = TransformerBasedOptimizerConfig::<f32> {
470            model_dimension: 0,
471            ..Default::default()
472        };
473        assert!(config.validate().is_err());
474
475        // Test mismatched attention dimensions
476        let config = TransformerBasedOptimizerConfig::<f32> {
477            model_dimension: 512,
478            num_attention_heads: 7, // 512 is not divisible by 7
479            ..Default::default()
480        };
481        assert!(config.validate().is_err());
482    }
483
484    #[test]
485    fn test_parameter_estimation() {
486        let config = TransformerBasedOptimizerConfig::<f32>::small();
487        let param_count = config.estimate_parameter_count();
488        assert!(param_count > 0);
489
490        let memory_usage = config.estimate_memory_usage();
491        assert!(memory_usage > 0.0);
492    }
493
494    #[test]
495    fn test_preset_configs() {
496        let small_config = TransformerBasedOptimizerConfig::<f32>::small();
497        assert!(small_config.validate().is_ok());
498        assert_eq!(small_config.model_dimension, 256);
499
500        let large_config = TransformerBasedOptimizerConfig::<f32>::large();
501        assert!(large_config.validate().is_ok());
502        assert_eq!(large_config.model_dimension, 1024);
503
504        let training_config = TransformerBasedOptimizerConfig::<f32>::for_training();
505        assert!(training_config.validate().is_ok());
506        assert_eq!(training_config.batch_size, 64);
507    }
508}
optirs_learned/transformer_based_optimizer/config.rs

optirs_learned/transformer_based_optimizer/
config.rs
