optirs_learned/transformer_based_optimizer/
feedforward.rs

1// Feed-forward network implementations for transformer layers
2
3use super::layers::ActivationLayer;
4use crate::error::{OptimError, Result};
5use scirs2_core::ndarray::{Array1, Array2, Axis};
6use scirs2_core::numeric::Float;
7use std::fmt::Debug;
8
9/// Activations cached by [`FeedForwardNetwork::forward_with_cache`] so the
10/// backward pass can reuse them instead of recomputing the forward pass.
11#[derive(Debug, Clone)]
12pub struct FeedForwardCache<
13    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
14> {
15    /// Input to the first linear layer
16    pub input: Array2<T>,
17
18    /// Pre-activation hidden state
19    pub hidden: Array2<T>,
20
21    /// Post-activation hidden state (input to the second linear layer)
22    pub activated: Array2<T>,
23}
24
25/// Feed-forward network implementation
26pub struct FeedForwardNetwork<
27    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
28> {
29    /// First linear layer (expansion)
30    linear1: LinearLayer<T>,
31
32    /// Second linear layer (projection)
33    linear2: LinearLayer<T>,
34
35    /// Activation function
36    activation: ActivationFunction,
37
38    /// Dropout layer for regularization
39    dropout: super::layers::DropoutLayer,
40}
41
42impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
43    FeedForwardNetwork<T>
44{
45    /// Create new feed-forward network
46    pub fn new(
47        input_dimension: usize,
48        hidden_dimension: usize,
49        activation: ActivationFunction,
50    ) -> Result<Self> {
51        let linear1 = LinearLayer::new(input_dimension, hidden_dimension)?;
52        let linear2 = LinearLayer::new(hidden_dimension, input_dimension)?;
53        let dropout = super::layers::DropoutLayer::new(0.1);
54
55        Ok(Self {
56            linear1,
57            linear2,
58            activation,
59            dropout,
60        })
61    }
62
63    /// Create with custom dropout rate
64    pub fn new_with_dropout(
65        input_dimension: usize,
66        hidden_dimension: usize,
67        activation: ActivationFunction,
68        dropout_rate: f64,
69    ) -> Result<Self> {
70        let linear1 = LinearLayer::new(input_dimension, hidden_dimension)?;
71        let linear2 = LinearLayer::new(hidden_dimension, input_dimension)?;
72        let dropout = super::layers::DropoutLayer::new(dropout_rate);
73
74        Ok(Self {
75            linear1,
76            linear2,
77            activation,
78            dropout,
79        })
80    }
81
82    /// Forward pass through feed-forward network
83    pub fn forward(&mut self, input: &Array2<T>) -> Result<Array2<T>> {
84        // First linear transformation
85        let hidden = self.linear1.forward(input)?;
86
87        // Apply activation function
88        let activated = ActivationLayer::apply(&hidden, self.activation);
89
90        // Apply dropout
91        let dropout_output = self.dropout.forward(&activated);
92
93        // Second linear transformation
94        let output = self.linear2.forward(&dropout_output)?;
95
96        Ok(output)
97    }
98
99    /// Forward pass that also returns the activations needed for the backward
100    /// pass. Dropout is bypassed here so that the cached activations match the
101    /// ones the gradients are computed against.
102    pub fn forward_with_cache(
103        &mut self,
104        input: &Array2<T>,
105    ) -> Result<(Array2<T>, FeedForwardCache<T>)> {
106        let hidden = self.linear1.forward(input)?;
107        let activated = ActivationLayer::apply(&hidden, self.activation);
108        let output = self.linear2.forward(&activated)?;
109
110        Ok((
111            output,
112            FeedForwardCache {
113                input: input.clone(),
114                hidden,
115                activated,
116            },
117        ))
118    }
119
120    /// Backward pass. Returns `dL/d(input)` and applies SGD steps to both
121    /// linear layers.
122    pub fn backward(
123        &mut self,
124        cache: &FeedForwardCache<T>,
125        grad_output: &Array2<T>,
126        learning_rate: T,
127    ) -> Result<Array2<T>> {
128        let grad_activated = self
129            .linear2
130            .backward(&cache.activated, grad_output, learning_rate)?;
131        let derivative = ActivationLayer::derivative(&cache.hidden, self.activation);
132        if derivative.dim() != grad_activated.dim() {
133            return Err(OptimError::InvalidConfig(
134                "Feed-forward backward shape mismatch at the activation".to_string(),
135            ));
136        }
137        let grad_hidden = &grad_activated * &derivative;
138        self.linear1
139            .backward(&cache.input, &grad_hidden, learning_rate)
140    }
141
142    /// Get parameter count
143    pub fn parameter_count(&self) -> usize {
144        self.linear1.parameter_count() + self.linear2.parameter_count()
145    }
146
147    /// Reset all parameters
148    pub fn reset(&mut self) -> Result<()> {
149        self.linear1.reset()?;
150        self.linear2.reset()?;
151        Ok(())
152    }
153
154    /// Set training mode
155    pub fn set_training(&mut self, training: bool) {
156        self.dropout.set_training(training);
157    }
158
159    /// Get activation function
160    pub fn get_activation(&self) -> ActivationFunction {
161        self.activation
162    }
163
164    /// Set activation function
165    pub fn set_activation(&mut self, activation: ActivationFunction) {
166        self.activation = activation;
167    }
168}
169
170/// Linear layer implementation
171pub struct LinearLayer<
172    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
173> {
174    /// Weight matrix
175    weight: Array2<T>,
176
177    /// Bias vector
178    bias: Array1<T>,
179
180    /// Input dimension
181    input_dim: usize,
182
183    /// Output dimension
184    output_dim: usize,
185}
186
187impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
188    LinearLayer<T>
189{
190    /// Create new linear layer with Xavier/Glorot uniform weights,
191    /// limit `sqrt(6 / (fan_in + fan_out))`.
192    pub fn new(input_dim: usize, output_dim: usize) -> Result<Self> {
193        if input_dim == 0 || output_dim == 0 {
194            return Err(OptimError::InvalidConfig(
195                "Linear layer dimensions must be positive".to_string(),
196            ));
197        }
198
199        let mut layer = Self {
200            weight: Array2::zeros((input_dim, output_dim)),
201            bias: Array1::zeros(output_dim),
202            input_dim,
203            output_dim,
204        };
205        layer.randomize_xavier();
206        Ok(layer)
207    }
208
209    /// Fill the weights with a fresh Xavier/Glorot uniform draw.
210    ///
211    /// The generator handle is acquired **once** for the whole matrix. Calling
212    /// the free `random()` function per element re-entered the thread-local
213    /// generator ~25M times when building a default-configuration optimizer,
214    /// which dominated construction time.
215    fn randomize_xavier(&mut self) {
216        let bound = (6.0 / (self.input_dim + self.output_dim).max(1) as f64).sqrt();
217        let mut rng = scirs2_core::random::thread_rng();
218        for elem in self.weight.iter_mut() {
219            let sample = rng.random::<f64>() * 2.0 - 1.0;
220            *elem =
221                scirs2_core::numeric::NumCast::from(sample * bound).unwrap_or_else(|| T::zero());
222        }
223    }
224
225    /// Create with He initialization (better for ReLU)
226    pub fn new_he_init(input_dim: usize, output_dim: usize) -> Result<Self> {
227        if input_dim == 0 || output_dim == 0 {
228            return Err(OptimError::InvalidConfig(
229                "Linear layer dimensions must be positive".to_string(),
230            ));
231        }
232
233        let scale = (2.0 / input_dim as f64).sqrt();
234        let mut weight = Array2::zeros((input_dim, output_dim));
235        let mut rng = scirs2_core::random::thread_rng();
236        for elem in weight.iter_mut() {
237            let sample = rng.random::<f64>() * 2.0 - 1.0;
238            *elem =
239                scirs2_core::numeric::NumCast::from(sample * scale).unwrap_or_else(|| T::zero());
240        }
241
242        Ok(Self {
243            weight,
244            bias: Array1::zeros(output_dim),
245            input_dim,
246            output_dim,
247        })
248    }
249
250    /// Forward pass through linear layer: `input @ weight + bias`.
251    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
252        let input_features = input.ncols();
253
254        if input_features != self.input_dim {
255            return Err(OptimError::InvalidConfig(format!(
256                "Input dimension mismatch: expected {}, got {}",
257                self.input_dim, input_features
258            )));
259        }
260
261        Ok(input.dot(&self.weight) + &self.bias)
262    }
263
264    /// Backward pass. Returns `dL/d(input)` and applies an SGD step.
265    pub fn backward(
266        &mut self,
267        input: &Array2<T>,
268        grad_output: &Array2<T>,
269        learning_rate: T,
270    ) -> Result<Array2<T>> {
271        if input.ncols() != self.input_dim || grad_output.ncols() != self.output_dim {
272            return Err(OptimError::InvalidConfig(format!(
273                "Linear backward shape mismatch: input {:?}, grad {:?}, layer ({}, {})",
274                input.dim(),
275                grad_output.dim(),
276                self.input_dim,
277                self.output_dim
278            )));
279        }
280        if input.nrows() != grad_output.nrows() {
281            return Err(OptimError::InvalidConfig(
282                "Linear backward received mismatched row counts".to_string(),
283            ));
284        }
285
286        let grad_weight = input.t().dot(grad_output);
287        let mut grad_bias = Array1::zeros(self.output_dim);
288        for j in 0..self.output_dim {
289            grad_bias[j] = grad_output.column(j).iter().fold(T::zero(), |a, &b| a + b);
290        }
291        let grad_input = grad_output.dot(&self.weight.t());
292
293        self.update_weights(&(grad_weight * learning_rate), &(grad_bias * learning_rate))?;
294
295        Ok(grad_input)
296    }
297
298    /// Get parameter count
299    pub fn parameter_count(&self) -> usize {
300        self.input_dim * self.output_dim + self.output_dim
301    }
302
303    /// Reset parameters, re-drawing the Xavier initialization.
304    pub fn reset(&mut self) -> Result<()> {
305        self.randomize_xavier();
306        self.bias.fill(T::zero());
307        Ok(())
308    }
309
310    /// Get weight matrix reference
311    pub fn get_weights(&self) -> &Array2<T> {
312        &self.weight
313    }
314
315    /// Get bias vector reference
316    pub fn get_bias(&self) -> &Array1<T> {
317        &self.bias
318    }
319
320    /// Update weights (for training)
321    pub fn update_weights(
322        &mut self,
323        weight_delta: &Array2<T>,
324        bias_delta: &Array1<T>,
325    ) -> Result<()> {
326        if weight_delta.shape() != self.weight.shape() {
327            return Err(OptimError::InvalidConfig(format!(
328                "Weight delta shape {:?} does not match {:?}",
329                weight_delta.dim(),
330                self.weight.dim()
331            )));
332        }
333
334        if bias_delta.len() != self.bias.len() {
335            return Err(OptimError::InvalidConfig(format!(
336                "Bias delta length {} does not match {}",
337                bias_delta.len(),
338                self.bias.len()
339            )));
340        }
341
342        self.weight = &self.weight - weight_delta;
343        self.bias = &self.bias - bias_delta;
344
345        Ok(())
346    }
347}
348
349/// Gated Linear Unit (GLU) implementation
350pub struct GatedLinearUnit<
351    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
352> {
353    /// Linear layer for gate
354    gate_linear: LinearLayer<T>,
355
356    /// Linear layer for values
357    value_linear: LinearLayer<T>,
358}
359
360impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
361    GatedLinearUnit<T>
362{
363    /// Create new GLU
364    pub fn new(input_dimension: usize, hidden_dimension: usize) -> Result<Self> {
365        let gate_linear = LinearLayer::new(input_dimension, hidden_dimension)?;
366        let value_linear = LinearLayer::new(input_dimension, hidden_dimension)?;
367
368        Ok(Self {
369            gate_linear,
370            value_linear,
371        })
372    }
373
374    /// Forward pass through GLU
375    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
376        let gate = self.gate_linear.forward(input)?;
377        let value = self.value_linear.forward(input)?;
378
379        // Apply sigmoid to gate and element-wise multiply
380        let sigmoid_gate = ActivationLayer::apply(&gate, ActivationFunction::Sigmoid);
381        let output = &sigmoid_gate * &value;
382
383        Ok(output)
384    }
385
386    /// Get parameter count
387    pub fn parameter_count(&self) -> usize {
388        self.gate_linear.parameter_count() + self.value_linear.parameter_count()
389    }
390
391    /// Reset parameters
392    pub fn reset(&mut self) -> Result<()> {
393        self.gate_linear.reset()?;
394        self.value_linear.reset()?;
395        Ok(())
396    }
397}
398
399/// Swish GLU variant
400pub struct SwiGLU<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static> {
401    /// Linear layer for gate
402    gate_linear: LinearLayer<T>,
403
404    /// Linear layer for values
405    value_linear: LinearLayer<T>,
406}
407
408impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static> SwiGLU<T> {
409    /// Create new SwiGLU
410    pub fn new(input_dimension: usize, hidden_dimension: usize) -> Result<Self> {
411        let gate_linear = LinearLayer::new(input_dimension, hidden_dimension)?;
412        let value_linear = LinearLayer::new(input_dimension, hidden_dimension)?;
413
414        Ok(Self {
415            gate_linear,
416            value_linear,
417        })
418    }
419
420    /// Forward pass through SwiGLU
421    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
422        let gate = self.gate_linear.forward(input)?;
423        let value = self.value_linear.forward(input)?;
424
425        // Apply Swish to gate and element-wise multiply
426        let swish_gate = ActivationLayer::apply(&gate, ActivationFunction::Swish);
427        let output = &swish_gate * &value;
428
429        Ok(output)
430    }
431
432    /// Get parameter count
433    pub fn parameter_count(&self) -> usize {
434        self.gate_linear.parameter_count() + self.value_linear.parameter_count()
435    }
436
437    /// Reset parameters
438    pub fn reset(&mut self) -> Result<()> {
439        self.gate_linear.reset()?;
440        self.value_linear.reset()?;
441        Ok(())
442    }
443}
444
445/// Expert mixture for sparse feed-forward networks
446pub struct MixtureOfExperts<
447    T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
448> {
449    /// Individual expert networks
450    experts: Vec<FeedForwardNetwork<T>>,
451
452    /// Gating network
453    gate: LinearLayer<T>,
454
455    /// Number of experts to activate (top-k)
456    top_k: usize,
457
458    /// Input dimension
459    input_dimension: usize,
460}
461
462impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
463    MixtureOfExperts<T>
464{
465    /// Create new mixture of experts
466    pub fn new(
467        input_dimension: usize,
468        hidden_dimension: usize,
469        num_experts: usize,
470        top_k: usize,
471        activation: ActivationFunction,
472    ) -> Result<Self> {
473        let mut experts = Vec::new();
474        for _ in 0..num_experts {
475            experts.push(FeedForwardNetwork::new(
476                input_dimension,
477                hidden_dimension,
478                activation,
479            )?);
480        }
481
482        let gate = LinearLayer::new(input_dimension, num_experts)?;
483
484        Ok(Self {
485            experts,
486            gate,
487            top_k: top_k.min(num_experts),
488            input_dimension,
489        })
490    }
491
492    /// Forward pass through mixture of experts
493    pub fn forward(&mut self, input: &Array2<T>) -> Result<Array2<T>> {
494        let batch_size = input.shape()[0];
495
496        // Compute gating scores
497        let gate_scores = self.gate.forward(input)?;
498        let gate_probs = self.softmax(&gate_scores);
499
500        // Find top-k experts for each sample
501        let mut output = Array2::zeros((batch_size, self.input_dimension));
502
503        for i in 0..batch_size {
504            let sample_input = input.row(i).insert_axis(Axis(0)).to_owned();
505            let sample_probs = gate_probs.row(i);
506
507            // Get top-k indices
508            let mut prob_indices: Vec<(usize, T)> = sample_probs
509                .iter()
510                .enumerate()
511                .map(|(idx, &prob)| (idx, prob))
512                .collect();
513
514            // `total_cmp` is a total order over floats, so a NaN gate score
515            // sorts deterministically instead of panicking inside `sort_by`.
516            prob_indices.sort_by(|a, b| {
517                b.1.to_f64()
518                    .unwrap_or(f64::NAN)
519                    .total_cmp(&a.1.to_f64().unwrap_or(f64::NAN))
520            });
521
522            let top_k_indices: Vec<usize> = prob_indices
523                .iter()
524                .take(self.top_k)
525                .map(|(idx, _)| *idx)
526                .collect();
527
528            // Compute weighted sum from top-k experts
529            let mut sample_output = Array1::zeros(self.input_dimension);
530            let mut total_weight = T::zero();
531
532            for &expert_idx in &top_k_indices {
533                let expert_output = self.experts[expert_idx].forward(&sample_input)?;
534                let weight = sample_probs[expert_idx];
535
536                total_weight = total_weight + weight;
537
538                for j in 0..self.input_dimension {
539                    sample_output[j] = sample_output[j] + weight * expert_output[[0, j]];
540                }
541            }
542
543            // Normalize by total weight
544            if total_weight > T::zero() {
545                for j in 0..self.input_dimension {
546                    sample_output[j] = sample_output[j] / total_weight;
547                    output[[i, j]] = sample_output[j];
548                }
549            }
550        }
551
552        Ok(output)
553    }
554
555    /// Softmax activation for gating
556    fn softmax(&self, input: &Array2<T>) -> Array2<T> {
557        let mut output = Array2::zeros(input.raw_dim());
558        let batch_size = input.shape()[0];
559
560        for i in 0..batch_size {
561            let row = input.row(i);
562            let max_val = row.iter().fold(T::neg_infinity(), |a, &b| a.max(b));
563
564            let mut exp_sum = T::zero();
565            let mut exp_row = Array1::zeros(row.len());
566
567            for (j, &val) in row.iter().enumerate() {
568                exp_row[j] = (val - max_val).exp();
569                exp_sum = exp_sum + exp_row[j];
570            }
571
572            for (j, &exp_val) in exp_row.iter().enumerate() {
573                output[[i, j]] = exp_val / exp_sum;
574            }
575        }
576
577        output
578    }
579
580    /// Get parameter count
581    pub fn parameter_count(&self) -> usize {
582        let expert_params: usize = self
583            .experts
584            .iter()
585            .map(|expert| expert.parameter_count())
586            .sum();
587
588        expert_params + self.gate.parameter_count()
589    }
590
591    /// Reset all parameters
592    pub fn reset(&mut self) -> Result<()> {
593        for expert in &mut self.experts {
594            expert.reset()?;
595        }
596        self.gate.reset()?;
597        Ok(())
598    }
599
600    /// Set training mode for all experts
601    pub fn set_training(&mut self, training: bool) {
602        for expert in &mut self.experts {
603            expert.set_training(training);
604        }
605    }
606}
607
608// Re-export for backward compatibility
609pub use super::config::ActivationFunction;
610
611#[cfg(test)]
612mod tests {
613    use super::*;
614
615    #[test]
616    fn test_feedforward_network() {
617        let ffn = FeedForwardNetwork::<f32>::new(
618            128,
619            512,
620            crate::transformer_based_optimizer::config::ActivationFunction::ReLU,
621        );
622        assert!(ffn.is_ok());
623
624        let mut network = ffn.expect("FeedForwardNetwork::new should succeed");
625        let input = Array2::<f32>::ones((4, 128));
626        let result = network.forward(&input);
627        assert!(result.is_ok());
628
629        let output = result.expect("forward should succeed");
630        assert_eq!(output.shape(), &[4, 128]);
631    }
632
633    #[test]
634    fn test_linear_layer() {
635        let linear = LinearLayer::<f32>::new(64, 128);
636        assert!(linear.is_ok());
637
638        let layer = linear.expect("LinearLayer::new should succeed");
639        let input = Array2::<f32>::zeros((2, 64));
640        let result = layer.forward(&input);
641        assert!(result.is_ok());
642
643        let output = result.expect("forward should succeed");
644        assert_eq!(output.shape(), &[2, 128]);
645        assert_eq!(layer.parameter_count(), 64 * 128 + 128);
646    }
647
648    #[test]
649    fn test_gated_linear_unit() {
650        let glu = GatedLinearUnit::<f32>::new(128, 256);
651        assert!(glu.is_ok());
652
653        let unit = glu.expect("GatedLinearUnit::new should succeed");
654        let input = Array2::<f32>::ones((2, 128));
655        let result = unit.forward(&input);
656        assert!(result.is_ok());
657
658        let output = result.expect("forward should succeed");
659        assert_eq!(output.shape(), &[2, 256]);
660    }
661
662    #[test]
663    fn test_swiglu() {
664        let swiglu = SwiGLU::<f32>::new(128, 256);
665        assert!(swiglu.is_ok());
666
667        let unit = swiglu.expect("SwiGLU::new should succeed");
668        let input = Array2::<f32>::ones((2, 128));
669        let result = unit.forward(&input);
670        assert!(result.is_ok());
671
672        let output = result.expect("forward should succeed");
673        assert_eq!(output.shape(), &[2, 256]);
674    }
675
676    #[test]
677    fn test_mixture_of_experts() {
678        let moe = MixtureOfExperts::<f32>::new(128, 256, 4, 2, ActivationFunction::ReLU);
679        assert!(moe.is_ok());
680
681        let mut mixture = moe.expect("MixtureOfExperts::new should succeed");
682        let input = Array2::<f32>::ones((3, 128));
683        let result = mixture.forward(&input);
684        assert!(result.is_ok());
685
686        let output = result.expect("forward should succeed");
687        assert_eq!(output.shape(), &[3, 128]);
688    }
689
690    #[test]
691    fn test_linear_layer_initialization() {
692        let xavier_layer =
693            LinearLayer::<f32>::new(64, 128).expect("LinearLayer::new should succeed");
694        let he_layer = LinearLayer::<f32>::new_he_init(64, 128)
695            .expect("LinearLayer::new_he_init should succeed");
696
697        assert_eq!(xavier_layer.parameter_count(), he_layer.parameter_count());
698        assert_eq!(
699            xavier_layer.get_weights().shape(),
700            he_layer.get_weights().shape()
701        );
702    }
703}
optirs_learned/transformer_based_optimizer/feedforward.rs

optirs_learned/transformer_based_optimizer/
feedforward.rs
