optirs_core/second_order/kfac/
config.rs

1// Configuration structures for K-FAC optimizer
2//
3// This module contains configuration types and data structures for the
4// K-FAC (Kronecker-Factored Approximate Curvature) second-order optimizer.
5
6use scirs2_core::numeric::Float;
7use std::fmt::Debug;
8
9/// K-FAC optimizer configuration
10#[derive(Debug, Clone)]
11pub struct KFACConfig<T: Float + Debug + Send + Sync + 'static> {
12    /// Learning rate
13    pub learning_rate: T,
14
15    /// Damping parameter for numerical stability
16    pub damping: T,
17
18    /// Weight decay (L2 regularization)
19    pub weight_decay: T,
20
21    /// Update frequency for covariance matrices
22    pub cov_update_freq: usize,
23
24    /// Update frequency for inverse covariance matrices
25    pub inv_update_freq: usize,
26
27    /// Exponential moving average decay for statistics
28    pub stat_decay: T,
29
30    /// Minimum eigenvalue for regularization
31    pub min_eigenvalue: T,
32
33    /// Maximum number of iterations for iterative inversion
34    pub max_inv_iterations: usize,
35
36    /// Tolerance for iterative inversion
37    pub inv_tolerance: T,
38
39    /// Use Tikhonov regularization
40    pub use_tikhonov: bool,
41
42    /// Enable automatic damping adjustment
43    pub auto_damping: bool,
44
45    /// Target acceptance ratio for damping adjustment
46    pub target_acceptance_ratio: T,
47}
48
49impl<T: Float + Debug + Send + Sync + 'static> Default for KFACConfig<T> {
50    fn default() -> Self {
51        Self {
52            learning_rate: T::from(0.001).unwrap_or_else(|| T::zero()),
53            damping: T::from(0.001).unwrap_or_else(|| T::zero()),
54            weight_decay: T::from(0.0).unwrap_or_else(|| T::zero()),
55            cov_update_freq: 10,
56            inv_update_freq: 100,
57            stat_decay: T::from(0.95).unwrap_or_else(|| T::zero()),
58            min_eigenvalue: T::from(1e-7).unwrap_or_else(|| T::zero()),
59            max_inv_iterations: 50,
60            inv_tolerance: T::from(1e-6).unwrap_or_else(|| T::zero()),
61            use_tikhonov: true,
62            auto_damping: true,
63            target_acceptance_ratio: T::from(0.75).unwrap_or_else(|| T::zero()),
64        }
65    }
66}
67
68/// Layer information for K-FAC
69#[derive(Debug, Clone)]
70pub struct LayerInfo {
71    /// Layer name/identifier
72    pub name: String,
73
74    /// Input dimension
75    pub input_dim: usize,
76
77    /// Output dimension
78    pub output_dim: usize,
79
80    /// Layer type
81    pub layer_type: LayerType,
82
83    /// Whether to include bias
84    pub has_bias: bool,
85}
86
87/// Types of layers supported by K-FAC
88#[derive(Debug, Clone, Copy, PartialEq)]
89pub enum LayerType {
90    /// Dense/Fully connected layer
91    Dense,
92
93    /// Convolutional layer
94    Convolution,
95
96    /// Convolutional layer with grouped/depthwise convolution
97    GroupedConvolution { groups: usize },
98
99    /// Embedding layer
100    Embedding,
101
102    /// Batch normalization layer
103    BatchNorm,
104}
105
106/// K-FAC performance statistics
107#[derive(Debug, Clone, Default)]
108pub struct KFACStats<T: Float + Debug + Send + Sync + 'static> {
109    /// Total number of optimization steps
110    pub total_steps: usize,
111
112    /// Number of covariance updates
113    pub cov_updates: usize,
114
115    /// Number of inverse updates
116    pub inv_updates: usize,
117
118    /// Average condition number of covariance matrices
119    pub avg_condition_number: T,
120
121    /// Time spent in different operations (in microseconds)
122    pub time_cov_update: u64,
123    pub time_inv_update: u64,
124    pub time_gradient_update: u64,
125
126    /// Memory usage estimate (in bytes)
127    pub memory_usage: usize,
128}
129
130impl<T: Float + Debug + Send + Sync + 'static> KFACConfig<T> {
131    /// Create configuration optimized for large models
132    pub fn for_large_model() -> Self {
133        Self {
134            cov_update_freq: 20,
135            inv_update_freq: 200,
136            stat_decay: T::from(0.99).unwrap_or_else(|| T::zero()),
137            damping: T::from(0.01).unwrap_or_else(|| T::zero()),
138            ..Default::default()
139        }
140    }
141
142    /// Create configuration optimized for small models with frequent updates
143    pub fn for_small_model() -> Self {
144        Self {
145            cov_update_freq: 5,
146            inv_update_freq: 50,
147            stat_decay: T::from(0.9).unwrap_or_else(|| T::zero()),
148            damping: T::from(0.001).unwrap_or_else(|| T::zero()),
149            ..Default::default()
150        }
151    }
152
153    /// Create configuration with conservative damping for stability
154    pub fn for_stability() -> Self {
155        Self {
156            damping: T::from(0.1).unwrap_or_else(|| T::zero()),
157            min_eigenvalue: T::from(1e-5).unwrap_or_else(|| T::zero()),
158            auto_damping: false,
159            use_tikhonov: true,
160            ..Default::default()
161        }
162    }
163
164    /// Validate configuration parameters
165    pub fn validate(&self) -> Result<(), String> {
166        if self.learning_rate <= T::zero() {
167            return Err("Learning rate must be positive".to_string());
168        }
169        if self.damping < T::zero() {
170            return Err("Damping must be non-negative".to_string());
171        }
172        if self.weight_decay < T::zero() {
173            return Err("Weight decay must be non-negative".to_string());
174        }
175        if self.cov_update_freq == 0 {
176            return Err("Covariance update frequency must be positive".to_string());
177        }
178        if self.inv_update_freq == 0 {
179            return Err("Inverse update frequency must be positive".to_string());
180        }
181        if self.stat_decay < T::zero() || self.stat_decay > T::one() {
182            return Err("Statistics decay must be between 0 and 1".to_string());
183        }
184        if self.min_eigenvalue <= T::zero() {
185            return Err("Minimum eigenvalue must be positive".to_string());
186        }
187        if self.inv_tolerance <= T::zero() {
188            return Err("Inverse tolerance must be positive".to_string());
189        }
190        if self.target_acceptance_ratio <= T::zero() || self.target_acceptance_ratio >= T::one() {
191            return Err("Target acceptance ratio must be between 0 and 1".to_string());
192        }
193        Ok(())
194    }
195}
196
197impl LayerInfo {
198    /// Create layer info for a dense layer
199    pub fn dense(name: String, input_dim: usize, output_dim: usize, has_bias: bool) -> Self {
200        Self {
201            name,
202            input_dim,
203            output_dim,
204            layer_type: LayerType::Dense,
205            has_bias,
206        }
207    }
208
209    /// Create layer info for a convolutional layer
210    pub fn convolution(name: String, input_dim: usize, output_dim: usize, has_bias: bool) -> Self {
211        Self {
212            name,
213            input_dim,
214            output_dim,
215            layer_type: LayerType::Convolution,
216            has_bias,
217        }
218    }
219
220    /// Get the expected size of the input covariance matrix
221    pub fn input_cov_size(&self) -> usize {
222        match self.layer_type {
223            LayerType::Dense => {
224                if self.has_bias {
225                    self.input_dim + 1
226                } else {
227                    self.input_dim
228                }
229            }
230            LayerType::Convolution => {
231                // For convolutional layers, this depends on the specific implementation
232                // This is a simplified calculation
233                self.input_dim
234            }
235            LayerType::GroupedConvolution { .. } => self.input_dim,
236            LayerType::Embedding => self.input_dim,
237            LayerType::BatchNorm => self.input_dim,
238        }
239    }
240
241    /// Get the expected size of the output gradient covariance matrix
242    pub fn output_cov_size(&self) -> usize {
243        self.output_dim
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn test_config_default() {
253        let config = KFACConfig::<f32>::default();
254        assert!(config.learning_rate > 0.0);
255        assert!(config.damping >= 0.0);
256        assert!(config.cov_update_freq > 0);
257        assert!(config.inv_update_freq > 0);
258        assert!(config.validate().is_ok());
259    }
260
261    #[test]
262    fn test_config_presets() {
263        let large_config = KFACConfig::<f64>::for_large_model();
264        assert!(large_config.validate().is_ok());
265        assert_eq!(large_config.cov_update_freq, 20);
266
267        let small_config = KFACConfig::<f64>::for_small_model();
268        assert!(small_config.validate().is_ok());
269        assert_eq!(small_config.cov_update_freq, 5);
270
271        let stable_config = KFACConfig::<f64>::for_stability();
272        assert!(stable_config.validate().is_ok());
273        assert!(!stable_config.auto_damping);
274    }
275
276    #[test]
277    fn test_layer_info_creation() {
278        let dense_layer = LayerInfo::dense("layer1".to_string(), 128, 64, true);
279        assert_eq!(dense_layer.layer_type, LayerType::Dense);
280        assert_eq!(dense_layer.input_cov_size(), 129); // +1 for bias
281
282        let conv_layer = LayerInfo::convolution("conv1".to_string(), 64, 32, false);
283        assert_eq!(conv_layer.layer_type, LayerType::Convolution);
284        assert_eq!(conv_layer.input_cov_size(), 64);
285    }
286
287    #[test]
288    fn test_config_validation() {
289        let mut config = KFACConfig::<f32> {
290            learning_rate: -0.1,
291            ..Default::default()
292        };
293
294        assert!(config.validate().is_err());
295
296        config.learning_rate = 0.001;
297        config.damping = -0.1;
298        assert!(config.validate().is_err());
299
300        config.damping = 0.001;
301        config.stat_decay = 1.5;
302        assert!(config.validate().is_err());
303
304        config.stat_decay = 0.95;
305        assert!(config.validate().is_ok());
306    }
307}
optirs_core/second_order/kfac/config.rs

optirs_core/second_order/kfac/
config.rs
