optirs_core/regularizers/
entropy.rs

1use scirs2_core::ndarray::{Array, ArrayBase, Data, Dimension, ScalarOperand};
2use scirs2_core::numeric::{Float, FromPrimitive};
3use std::fmt::Debug;
4
5use crate::error::Result;
6use crate::regularizers::Regularizer;
7
8/// Entropy regularization
9///
10/// Entropy regularization encourages a model to produce more confident outputs (lower entropy),
11/// or more uncertain outputs (higher entropy), depending on the settings. This is often used
12/// in reinforcement learning, semi-supervised learning, and some classifier applications.
13///
14/// # Types
15///
16/// * `MaximizeEntropy`: Encourages high entropy (uniform, uncertain predictions)
17/// * `MinimizeEntropy`: Encourages low entropy (confident, peaked predictions)
18///
19/// # Parameters
20///
21/// * `lambda`: Regularization strength coefficient, controls the amount of regularization applied
22/// * `epsilon`: Small value for numerical stability (prevents log(0))
23///
24#[derive(Debug, Clone, Copy)]
25pub enum EntropyRegularizerType {
26    /// Maximize entropy (encourages uniform distributions)
27    MaximizeEntropy,
28    /// Minimize entropy (encourages confident predictions)
29    MinimizeEntropy,
30}
31
32/// Entropy regularization for probability distributions
33///
34/// This regularizer can either encourage high entropy (more uniform distributions) or
35/// low entropy (more peaked distributions) depending on the selected regularizer type.
36///
37/// It's commonly used in reinforcement learning algorithms, semi-supervised learning,
38/// and some classification tasks where controlling the certainty of outputs is desired.
39#[derive(Debug, Clone, Copy)]
40pub struct EntropyRegularization<A: Float + FromPrimitive + Debug> {
41    /// Regularization strength
42    pub lambda: A,
43    /// Small value for numerical stability
44    pub epsilon: A,
45    /// Type of entropy regularization
46    pub reg_type: EntropyRegularizerType,
47}
48
49impl<A: Float + FromPrimitive + Debug + Send + Sync> EntropyRegularization<A> {
50    /// Create a new entropy regularization
51    ///
52    /// # Arguments
53    ///
54    /// * `lambda` - Regularization strength coefficient
55    /// * `reg_type` - Type of entropy regularization (maximize or minimize)
56    ///
57    /// # Returns
58    ///
59    /// An entropy regularization with default epsilon
60    pub fn new(lambda: A, regtype: EntropyRegularizerType) -> Self {
61        let epsilon =
62            A::from_f64(1e-8).expect("EntropyRegularization: default epsilon (1e-8) must fit in A");
63        Self {
64            lambda,
65            epsilon,
66            reg_type: regtype,
67        }
68    }
69
70    /// Create a new entropy regularization with custom epsilon
71    ///
72    /// # Arguments
73    ///
74    /// * `lambda` - Regularization strength coefficient
75    /// * `epsilon` - Small value for numerical stability
76    /// * `reg_type` - Type of entropy regularization (maximize or minimize)
77    ///
78    /// # Returns
79    ///
80    /// An entropy regularization with custom epsilon
81    pub fn new_with_epsilon(lambda: A, epsilon: A, regtype: EntropyRegularizerType) -> Self {
82        Self {
83            lambda,
84            epsilon,
85            reg_type: regtype,
86        }
87    }
88
89    /// Calculate the entropy of a probability distribution
90    ///
91    /// # Arguments
92    ///
93    /// * `probs` - Probability distribution (should sum to 1 along the appropriate axis)
94    ///
95    /// # Returns
96    ///
97    /// The entropy value
98    pub fn calculate_entropy<S, D>(&self, probs: &ArrayBase<S, D>) -> A
99    where
100        S: Data<Elem = A>,
101        D: Dimension,
102    {
103        // Clip probabilities to avoid log(0)
104        let safe_probs = probs.mapv(|p| {
105            if p < self.epsilon {
106                self.epsilon
107            } else if p > (A::one() - self.epsilon) {
108                A::one() - self.epsilon
109            } else {
110                p
111            }
112        });
113
114        // Calculate entropy: -sum(p * log(p))
115        let neg_entropy = safe_probs.mapv(|p| p * p.ln()).sum();
116        -neg_entropy
117    }
118
119    /// Calculate gradient of entropy with respect to input probabilities
120    ///
121    /// # Arguments
122    ///
123    /// * `probs` - Probability distribution
124    ///
125    /// # Returns
126    ///
127    /// The gradient of entropy with respect to probabilities
128    fn entropy_gradient<S, D>(&self, probs: &ArrayBase<S, D>) -> Array<A, D>
129    where
130        S: Data<Elem = A>,
131        D: Dimension,
132    {
133        // Clip probabilities to avoid log(0)
134        let safe_probs = probs.mapv(|p| {
135            if p < self.epsilon {
136                self.epsilon
137            } else if p > (A::one() - self.epsilon) {
138                A::one() - self.epsilon
139            } else {
140                p
141            }
142        });
143
144        // Gradient of entropy: -(1 + log(p))
145        let gradient = safe_probs.mapv(|p| -(A::one() + p.ln()));
146
147        // For minimizing entropy, we negate the gradient
148        match self.reg_type {
149            EntropyRegularizerType::MaximizeEntropy => gradient,
150            EntropyRegularizerType::MinimizeEntropy => gradient.mapv(|g| -g),
151        }
152    }
153}
154
155impl<A, D> Regularizer<A, D> for EntropyRegularization<A>
156where
157    A: Float + ScalarOperand + Debug + FromPrimitive + Send + Sync,
158    D: Dimension,
159{
160    fn apply(&self, params: &Array<A, D>, gradients: &mut Array<A, D>) -> Result<A> {
161        // Calculate entropy penalty
162        let entropy = self.calculate_entropy(params);
163
164        // Calculate entropy gradients
165        let entropy_grads = self.entropy_gradient(params);
166
167        // Scale gradients by lambda and add to input gradients
168        gradients.zip_mut_with(&entropy_grads, |g, &e| *g = *g + self.lambda * e);
169
170        // Return the regularization term to be added to the loss:
171        // For maximizing entropy, we return -lambda * entropy (to minimize -entropy)
172        // For minimizing entropy, we return lambda * entropy (to minimize entropy)
173        let penalty = match self.reg_type {
174            EntropyRegularizerType::MaximizeEntropy => -self.lambda * entropy,
175            EntropyRegularizerType::MinimizeEntropy => self.lambda * entropy,
176        };
177
178        Ok(penalty)
179    }
180
181    fn penalty(&self, params: &Array<A, D>) -> Result<A> {
182        // Calculate entropy penalty
183        let entropy = self.calculate_entropy(params);
184
185        // For maximizing entropy, we return -lambda * entropy (to minimize -entropy)
186        // For minimizing entropy, we return lambda * entropy (to minimize entropy)
187        let penalty = match self.reg_type {
188            EntropyRegularizerType::MaximizeEntropy => -self.lambda * entropy,
189            EntropyRegularizerType::MinimizeEntropy => self.lambda * entropy,
190        };
191
192        Ok(penalty)
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use approx::assert_abs_diff_eq;
200    use scirs2_core::ndarray::Array1;
201
202    #[test]
203    fn test_entropy_regularization_creation() {
204        let er = EntropyRegularization::new(0.1f64, EntropyRegularizerType::MaximizeEntropy);
205        assert_eq!(er.lambda, 0.1);
206        assert_eq!(er.epsilon, 1e-8);
207        match er.reg_type {
208            EntropyRegularizerType::MaximizeEntropy => (),
209            _ => panic!("Wrong regularizer type"),
210        }
211
212        let er = EntropyRegularization::new_with_epsilon(
213            0.2f64,
214            1e-10,
215            EntropyRegularizerType::MinimizeEntropy,
216        );
217        assert_eq!(er.lambda, 0.2);
218        assert_eq!(er.epsilon, 1e-10);
219        match er.reg_type {
220            EntropyRegularizerType::MinimizeEntropy => (),
221            _ => panic!("Wrong regularizer type"),
222        }
223    }
224
225    #[test]
226    fn test_calculate_entropy() {
227        // Uniform distribution (maximum entropy)
228        let uniform = Array1::from_vec(vec![0.25f64, 0.25, 0.25, 0.25]);
229        let er = EntropyRegularization::new(1.0f64, EntropyRegularizerType::MaximizeEntropy);
230        let entropy = er.calculate_entropy(&uniform);
231
232        // Entropy of uniform distribution should be ln(n)
233        let expected = (4.0f64).ln();
234        assert_abs_diff_eq!(entropy, expected, epsilon = 1e-6);
235
236        // Peaked distribution (low entropy)
237        let peaked = Array1::from_vec(vec![0.01f64, 0.01, 0.97, 0.01]);
238        let entropy = er.calculate_entropy(&peaked);
239        assert!(entropy < expected); // Should be less than uniform entropy
240    }
241
242    #[test]
243    fn test_entropy_gradient() {
244        let er = EntropyRegularization::new(1.0f64, EntropyRegularizerType::MaximizeEntropy);
245
246        // For uniform distribution, gradients should be approximately equal
247        let uniform = Array1::from_vec(vec![0.25f64, 0.25, 0.25, 0.25]);
248        let grads = er.entropy_gradient(&uniform);
249
250        // Expected gradient: -(1 + ln(0.25))
251        let expected = -(1.0 + 0.25f64.ln());
252        for &g in grads.iter() {
253            assert_abs_diff_eq!(g, expected, epsilon = 1e-6);
254        }
255
256        // For peaked distribution, gradients should be different for different probabilities
257        let peaked = Array1::from_vec(vec![0.1f64, 0.1, 0.7, 0.1]);
258        let grads = er.entropy_gradient(&peaked);
259
260        // The gradient for larger probability should have a smaller absolute value
261        // because ln(0.7) is greater (less negative) than ln(0.1)
262        // So -(1 + ln(0.7)) has smaller magnitude than -(1 + ln(0.1))
263        assert!(grads[2].abs() < grads[0].abs());
264    }
265
266    #[test]
267    fn test_maximize_entropy_penalty() {
268        // For maximizing entropy, we want to minimize -entropy
269        let er = EntropyRegularization::new(1.0f64, EntropyRegularizerType::MaximizeEntropy);
270
271        // Uniform distribution (high entropy)
272        let uniform = Array1::from_vec(vec![0.25f64, 0.25, 0.25, 0.25]);
273        let penalty = er
274            .penalty(&uniform)
275            .expect("er.penalty succeeds in test_maximize_entropy_penalty");
276
277        // Peaked distribution (low entropy)
278        let peaked = Array1::from_vec(vec![0.01f64, 0.01, 0.97, 0.01]);
279        let peaked_penalty = er
280            .penalty(&peaked)
281            .expect("er.penalty succeeds in test_maximize_entropy_penalty");
282
283        // The penalty for peaked should be greater than for uniform
284        // because we're trying to maximize entropy
285        assert!(peaked_penalty > penalty);
286    }
287
288    #[test]
289    fn test_minimize_entropy_penalty() {
290        // For minimizing entropy, we want to minimize entropy
291        let er = EntropyRegularization::new(1.0f64, EntropyRegularizerType::MinimizeEntropy);
292
293        // Uniform distribution (high entropy)
294        let uniform = Array1::from_vec(vec![0.25f64, 0.25, 0.25, 0.25]);
295        let penalty = er
296            .penalty(&uniform)
297            .expect("er.penalty succeeds in test_minimize_entropy_penalty");
298
299        // Peaked distribution (low entropy)
300        let peaked = Array1::from_vec(vec![0.01f64, 0.01, 0.97, 0.01]);
301        let peaked_penalty = er
302            .penalty(&peaked)
303            .expect("er.penalty succeeds in test_minimize_entropy_penalty");
304
305        // The penalty for uniform should be greater than for peaked
306        // because we're trying to minimize entropy
307        assert!(penalty > peaked_penalty);
308    }
309
310    #[test]
311    fn test_apply_gradients() {
312        let lambda = 0.5f64;
313        let er = EntropyRegularization::new(lambda, EntropyRegularizerType::MaximizeEntropy);
314
315        let probs = Array1::from_vec(vec![0.25f64, 0.25, 0.25, 0.25]);
316        let mut gradients = Array1::zeros(4);
317
318        let penalty = er
319            .apply(&probs, &mut gradients)
320            .expect("er.apply succeeds in test_apply_gradients");
321
322        // Check that gradients have been modified
323        assert!(gradients.iter().all(|&g| g != 0.0));
324
325        // For uniform distribution, all gradients should be equal
326        let first = gradients[0];
327        assert!(gradients.iter().all(|&g| (g - first).abs() < 1e-6));
328
329        // Expected gradient: -lambda * (1 + ln(0.25))
330        let expected_grad = -lambda * (1.0 + 0.25f64.ln());
331        assert_abs_diff_eq!(gradients[0], expected_grad, epsilon = 1e-6);
332
333        // Check penalty matches expected value
334        let entropy = (4.0f64).ln(); // Entropy of uniform distribution
335        let expected_penalty = -lambda * entropy; // For maximizing entropy
336        assert_abs_diff_eq!(penalty, expected_penalty, epsilon = 1e-6);
337    }
338
339    #[test]
340    fn test_regularizer_trait() {
341        // Test that EntropyRegularization implements Regularizer trait correctly
342        let er = EntropyRegularization::new(0.1f64, EntropyRegularizerType::MaximizeEntropy);
343
344        let probs = Array1::from_vec(vec![0.25f64, 0.25, 0.25, 0.25]);
345        let mut gradients = Array1::zeros(4);
346
347        // Both methods should return the same penalty for the same input
348        let penalty1 = er
349            .apply(&probs, &mut gradients)
350            .expect("er.apply succeeds in test_regularizer_trait");
351        let penalty2 = er
352            .penalty(&probs)
353            .expect("er.penalty succeeds in test_regularizer_trait");
354
355        assert_abs_diff_eq!(penalty1, penalty2, epsilon = 1e-10);
356    }
357}
optirs_core/regularizers/entropy.rs

optirs_core/regularizers/
entropy.rs
