optirs_core/regularizers/
shakedrop.rs

1use scirs2_core::ndarray::{Array, ArrayBase, Data, Dimension, ScalarOperand};
2use scirs2_core::numeric::{Float, FromPrimitive};
3use std::cell::RefCell;
4use std::fmt::Debug;
5
6use crate::error::{OptimError, Result};
7use crate::regularizers::Regularizer;
8
9/// Result type for ShakeDrop forward pass: transformed activations and gate parameters (b, alpha, beta).
10type ShakeDropResult<A, D> = (Array<A, D>, (A, A, A));
11
12/// ShakeDrop regularization
13///
14/// ShakeDrop is a regularization method that extends Stochastic Depth and
15/// is often used in very deep neural networks. It randomly scales activations
16/// during training.
17///
18/// # Parameters
19///
20/// * `p` - The probability of activating ShakeDrop (probability of activating the forward
21///   pass transformation), value between 0 and 1.
22/// * `alpha_range` - The range for the alpha parameter used in forward pass (default: [-1.0, 1.0]).
23/// * `beta_range` - The range for the beta parameter used in backward pass (default: [0.0, 1.0]).
24///
25/// # References
26///
27/// * Yamada, Y., Iwamura, M., & Kise, K. (2018). ShakeDrop regularization.
28///   arXiv preprint arXiv:1802.02375.
29///
30#[derive(Debug)]
31pub struct ShakeDrop<A: Float + FromPrimitive + Debug> {
32    /// Probability of applying ShakeDrop
33    pub p: A,
34    /// Range for the alpha parameter
35    pub alpha_range: (A, A),
36    /// Range for the beta parameter
37    pub beta_range: (A, A),
38    /// Random number generator (wrapped in RefCell for interior mutability)
39    rng: RefCell<scirs2_core::random::Random<scirs2_core::random::rngs::StdRng>>,
40}
41
42impl<A: Float + FromPrimitive + Debug + Send + Sync> ShakeDrop<A> {
43    /// Create a new ShakeDrop regularizer
44    ///
45    /// # Arguments
46    ///
47    /// * `p` - Probability of applying ShakeDrop, between 0 and 1
48    /// * `alpha_range` - Range for the alpha parameter (default: [-1.0, 1.0])
49    /// * `beta_range` - Range for the beta parameter (default: [0.0, 1.0])
50    ///
51    /// # Returns
52    ///
53    /// A ShakeDrop regularizer
54    pub fn new(p: A) -> Self {
55        let zero = A::zero();
56        let one = A::one();
57        let neg_one = zero - one;
58
59        Self {
60            p,
61            alpha_range: (neg_one, one),
62            beta_range: (zero, one),
63            rng: RefCell::new(scirs2_core::random::Random::seed(42)),
64        }
65    }
66
67    /// Create a new ShakeDrop regularizer with custom ranges
68    ///
69    /// # Arguments
70    ///
71    /// * `p` - Probability of applying ShakeDrop, between 0 and 1
72    /// * `alpha_range` - Range for the alpha parameter
73    /// * `beta_range` - Range for the beta parameter
74    ///
75    /// # Returns
76    ///
77    /// A ShakeDrop regularizer
78    pub fn new_with_ranges(p: A, alpharange: (A, A), beta_range: (A, A)) -> Self {
79        Self {
80            p,
81            alpha_range: alpharange,
82            beta_range,
83            rng: RefCell::new(scirs2_core::random::Random::seed(42)),
84        }
85    }
86
87    /// Get a random value between the given range
88    fn random_in_range(&self, range: (A, A)) -> Result<A> {
89        let (min, max) = range;
90        let min_f = min
91            .to_f64()
92            .ok_or_else(|| OptimError::InvalidConfig("Failed to convert min to f64".to_string()))?;
93        let max_f = max
94            .to_f64()
95            .ok_or_else(|| OptimError::InvalidConfig("Failed to convert max to f64".to_string()))?;
96
97        // Handle equal min and max to avoid "empty range" error
98        if (max_f - min_f).abs() < 1e-10 {
99            return Ok(min);
100        }
101
102        let random_val = self.rng.borrow_mut().gen_range(min_f..max_f);
103        A::from_f64(random_val).ok_or_else(|| {
104            OptimError::InvalidConfig("Failed to convert random value from f64".to_string())
105        })
106    }
107
108    /// Get a forward pass gate for the ShakeDrop
109    ///
110    /// # Returns
111    ///
112    /// A tuple (b, alpha, beta):
113    /// - b: Binary gate (1 or 0) based on the probability p
114    /// - alpha: Random value within alpha_range if b is 1, otherwise 0
115    /// - beta: Random value within beta_range
116    fn get_gate(&self) -> Result<(A, A, A)> {
117        let zero = A::zero();
118        let one = A::one();
119
120        // Determine if the gate is active
121        let u: f64 = self.rng.borrow_mut().gen_range(0.0..1.0);
122        let p_f64 = self
123            .p
124            .to_f64()
125            .ok_or_else(|| OptimError::InvalidConfig("Failed to convert p to f64".to_string()))?;
126        let b = if u < p_f64 { one } else { zero };
127
128        // Get random alpha if gate is active..otherwise 0
129        let alpha = if b > zero {
130            self.random_in_range(self.alpha_range)?
131        } else {
132            zero
133        };
134
135        // Get random beta regardless of gate
136        let beta = self.random_in_range(self.beta_range)?;
137
138        Ok((b, alpha, beta))
139    }
140
141    /// Apply ShakeDrop to input activations
142    ///
143    /// # Arguments
144    ///
145    /// * `x` - Input activation tensor
146    ///
147    /// # Returns
148    ///
149    /// The transformed activations and gate parameters for use in backward pass
150    pub fn forward<S, D>(&self, x: &ArrayBase<S, D>) -> Result<ShakeDropResult<A, D>>
151    where
152        S: Data<Elem = A>,
153        D: Dimension,
154    {
155        // Get the gate values
156        let (b, alpha, beta) = self.get_gate()?;
157
158        // Apply ShakeDrop transformation
159        // During forward pass: x' = x * (b + alpha - b*alpha)
160        let factor = b + alpha - b * alpha;
161        let result = x.mapv(|v| v * factor);
162
163        Ok((result, (b, alpha, beta)))
164    }
165
166    /// Backward pass for ShakeDrop
167    ///
168    /// # Arguments
169    ///
170    /// * `grad_output` - Gradient from the next layer
171    /// * `gate_params` - The gate parameters (b, alpha, beta) from the forward pass
172    ///
173    /// # Returns
174    ///
175    /// The modified gradients
176    pub fn backward<S, D>(
177        &self,
178        grad_output: &ArrayBase<S, D>,
179        gate_params: (A, A, A),
180    ) -> Array<A, D>
181    where
182        S: Data<Elem = A>,
183        D: Dimension,
184    {
185        let (b, _alpha, beta) = gate_params;
186
187        // During backward pass: grad_x = grad_output * (b + beta - b*beta)
188        // (alpha, drawn for the forward pass, is intentionally not reused here:
189        // ShakeDrop decorrelates the forward and backward scaling factors.)
190        let factor = b + beta - b * beta;
191        grad_output.mapv(|g| g * factor)
192    }
193}
194
195impl<A: Float + FromPrimitive + Debug + ScalarOperand, D: Dimension + Send + Sync> Regularizer<A, D>
196    for ShakeDrop<A>
197{
198    fn apply(&self, _params: &Array<A, D>, _gradients: &mut Array<A, D>) -> Result<A> {
199        // ShakeDrop is typically applied to activations, not parameters
200        // In this implementation, apply() isn't the primary usage pattern
201        // Instead, users would call forward() during the forward pass
202        // and backward() during the backward pass
203        Err(OptimError::InvalidConfig(
204            "ShakeDrop should be applied to activations during forward/backward passes, \
205             not through the Regularizer trait's apply method"
206                .to_string(),
207        ))
208    }
209
210    fn penalty(&self, _params: &Array<A, D>) -> Result<A> {
211        // ShakeDrop doesn't add a penalty term to the loss function
212        Ok(A::zero())
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use approx::assert_abs_diff_eq;
220    use scirs2_core::ndarray::{Array1, Array2};
221
222    #[test]
223    fn test_shakedrop_new() {
224        let sd = ShakeDrop::new(0.5f64);
225        assert_eq!(sd.p, 0.5);
226        assert_eq!(sd.alpha_range, (-1.0, 1.0));
227        assert_eq!(sd.beta_range, (0.0, 1.0));
228    }
229
230    #[test]
231    fn test_shakedrop_new_with_ranges() {
232        let sd = ShakeDrop::new_with_ranges(0.7f64, (-0.5, 0.5), (0.2, 0.8));
233        assert_eq!(sd.p, 0.7);
234        assert_eq!(sd.alpha_range, (-0.5, 0.5));
235        assert_eq!(sd.beta_range, (0.2, 0.8));
236    }
237
238    #[test]
239    fn test_shakedrop_forward_backward() {
240        // Create a simple 2D array
241        let x = Array2::from_elem((2, 3), 1.0f64);
242
243        // Initialize ShakeDrop with p=1.0 to ensure gate is always active
244        // Use slightly different values for min and max to avoid empty range error
245        let sd = ShakeDrop::new_with_ranges(1.0f64, (0.5, 0.500001), (0.5, 0.500001));
246
247        // Forward pass
248        let (output, gate_params) = sd.forward(&x).expect("forward failed");
249
250        // Verify the gate parameters
251        assert_eq!(gate_params.0, 1.0); // b should be 1 since p=1.0
252        assert_abs_diff_eq!(gate_params.1, 0.5, epsilon = 1e-5); // alpha should be approximately 0.5
253        assert_abs_diff_eq!(gate_params.2, 0.5, epsilon = 1e-5); // beta should be approximately 0.5
254
255        // The expected output is x * (b + alpha - b*alpha) = x * (1 + 0.5 - 1*0.5) = x * 1
256        for &val in output.iter() {
257            assert_abs_diff_eq!(val, 1.0, epsilon = 1e-5);
258        }
259
260        // Backward pass
261        let grad_output = Array2::from_elem((2, 3), 2.0f64);
262        let grad_input = sd.backward(&grad_output, gate_params);
263
264        // The expected gradient is grad_output * (b + beta - b*beta) = grad_output * (1 + 0.5 - 1*0.5) = grad_output * 1
265        for &val in grad_input.iter() {
266            assert_abs_diff_eq!(val, 2.0, epsilon = 1e-5);
267        }
268    }
269
270    #[test]
271    fn test_shakedrop_forward_inactive() {
272        // Create a simple 1D array
273        let x = Array1::from_vec(vec![1.0f64, 2.0, 3.0]);
274
275        // Initialize ShakeDrop with p=0.0 to ensure gate is always inactive
276        // Use slightly different values for min and max to avoid empty range error
277        let sd = ShakeDrop::new_with_ranges(0.0f64, (-0.5, -0.499999), (0.5, 0.500001));
278
279        // Forward pass - gate should be inactive
280        let (output, gate_params) = sd.forward(&x).expect("forward failed");
281
282        // Verify the gate parameters
283        assert_eq!(gate_params.0, 0.0); // b should be 0 since p=0.0
284        assert_eq!(gate_params.1, 0.0); // alpha should be 0 when gate is inactive
285        assert_abs_diff_eq!(gate_params.2, 0.5, epsilon = 1e-5); // beta should be approximately 0.5
286
287        // The expected output is x * (b + alpha - b*alpha) = x * (0 + 0 - 0*0) = x * 0
288        for &val in output.iter() {
289            assert_abs_diff_eq!(val, 0.0, epsilon = 1e-10);
290        }
291    }
292
293    #[test]
294    fn test_shakedrop_gen_range() {
295        let sd = ShakeDrop::new(0.5f64);
296
297        // Test random value generation within range
298        for _ in 0..100 {
299            let value = sd
300                .random_in_range((-0.5, 0.5))
301                .expect("random_in_range failed");
302            assert!((-0.5..=0.5).contains(&value));
303        }
304
305        // Test with very small range (should not panic)
306        let value = sd
307            .random_in_range((0.5, 0.5))
308            .expect("random_in_range failed");
309        assert_eq!(value, 0.5);
310    }
311
312    #[test]
313    fn test_shakedrop_get_gate() {
314        // Test with p=1.0 - gate should always be active
315        let sd = ShakeDrop::new(1.0f64);
316        for _ in 0..10 {
317            let (b, alpha, beta) = sd.get_gate().expect("get_gate failed");
318            assert_eq!(b, 1.0);
319            assert!((-1.0..=1.0).contains(&alpha));
320            assert!((0.0..=1.0).contains(&beta));
321        }
322
323        // Test with p=0.0 - gate should always be inactive
324        let sd = ShakeDrop::new(0.0f64);
325        for _ in 0..10 {
326            let (b, alpha, beta) = sd.get_gate().expect("get_gate failed");
327            assert_eq!(b, 0.0);
328            assert_eq!(alpha, 0.0);
329            assert!((0.0..=1.0).contains(&beta));
330        }
331    }
332
333    #[test]
334    fn test_regularizer_trait() {
335        let sd = ShakeDrop::new(0.5f64);
336        let params = Array2::from_elem((2, 3), 1.0f64);
337        let mut grads = Array2::from_elem((2, 3), 1.0f64);
338
339        // apply() should return an error for ShakeDrop
340        assert!(sd.apply(&params, &mut grads).is_err());
341
342        // penalty() should return zero
343        let penalty = sd.penalty(&params).expect("penalty failed");
344        assert_eq!(penalty, 0.0);
345    }
346}
optirs_core/regularizers/shakedrop.rs

optirs_core/regularizers/
shakedrop.rs
