rust_lstm/layers/
bilstm_network.rs

1use crate::layers::lstm_cell::{masked, LSTMCell, LSTMCellCache, LSTMCellGradients};
2use crate::optimizers::Optimizer;
3use ndarray::Array2;
4
5/// Cache for bidirectional LSTM forward pass
6///
7/// All vectors are indexed `[layer][time]` in the original time order.
8#[derive(Clone)]
9pub struct BiLSTMNetworkCache {
10    pub forward_caches: Vec<Vec<LSTMCellCache>>,
11    pub backward_caches: Vec<Vec<LSTMCellCache>>,
12    pub forward_output_masks: Vec<Vec<Option<Array2<f64>>>>,
13    pub backward_output_masks: Vec<Vec<Option<Array2<f64>>>>,
14}
15
16/// Configuration for combining forward and backward outputs
17#[derive(Clone, Debug)]
18pub enum CombineMode {
19    Concat,
20    Sum,
21    Average,
22}
23
24/// Bidirectional LSTM network for sequence modeling
25#[derive(Clone)]
26pub struct BiLSTMNetwork {
27    forward_cells: Vec<LSTMCell>,
28    backward_cells: Vec<LSTMCell>,
29    pub input_size: usize,
30    pub hidden_size: usize,
31    pub num_layers: usize,
32    pub combine_mode: CombineMode,
33    pub is_training: bool,
34}
35
36impl BiLSTMNetwork {
37    /// Creates a new bidirectional LSTM network
38    ///
39    /// # Arguments
40    /// * `input_size` - Size of input features
41    /// * `hidden_size` - Size of hidden state for each direction
42    /// * `num_layers` - Number of bidirectional layers
43    /// * `combine_mode` - How to combine forward and backward outputs
44    pub fn new(
45        input_size: usize,
46        hidden_size: usize,
47        num_layers: usize,
48        combine_mode: CombineMode,
49    ) -> Self {
50        let mut forward_cells = Vec::new();
51        let mut backward_cells = Vec::new();
52
53        for i in 0..num_layers {
54            let layer_input_size = if i == 0 {
55                input_size
56            } else {
57                match combine_mode {
58                    CombineMode::Concat => 2 * hidden_size,
59                    CombineMode::Sum | CombineMode::Average => hidden_size,
60                }
61            };
62
63            forward_cells.push(LSTMCell::new(layer_input_size, hidden_size));
64            backward_cells.push(LSTMCell::new(layer_input_size, hidden_size));
65        }
66
67        BiLSTMNetwork {
68            forward_cells,
69            backward_cells,
70            input_size,
71            hidden_size,
72            num_layers,
73            combine_mode,
74            is_training: true,
75        }
76    }
77
78    /// Create BiLSTM with concatenated outputs (most common)
79    pub fn new_concat(input_size: usize, hidden_size: usize, num_layers: usize) -> Self {
80        Self::new(input_size, hidden_size, num_layers, CombineMode::Concat)
81    }
82
83    /// Create BiLSTM with summed outputs
84    pub fn new_sum(input_size: usize, hidden_size: usize, num_layers: usize) -> Self {
85        Self::new(input_size, hidden_size, num_layers, CombineMode::Sum)
86    }
87
88    /// Create BiLSTM with averaged outputs
89    pub fn new_average(input_size: usize, hidden_size: usize, num_layers: usize) -> Self {
90        Self::new(input_size, hidden_size, num_layers, CombineMode::Average)
91    }
92
93    /// Get the output size based on combine mode
94    pub fn output_size(&self) -> usize {
95        match self.combine_mode {
96            CombineMode::Concat => 2 * self.hidden_size,
97            CombineMode::Sum | CombineMode::Average => self.hidden_size,
98        }
99    }
100
101    /// Apply dropout configuration to all cells
102    pub fn with_input_dropout(mut self, dropout_rate: f64, variational: bool) -> Self {
103        for cell in &mut self.forward_cells {
104            *cell = cell.clone().with_input_dropout(dropout_rate, variational);
105        }
106        for cell in &mut self.backward_cells {
107            *cell = cell.clone().with_input_dropout(dropout_rate, variational);
108        }
109        self
110    }
111
112    pub fn with_recurrent_dropout(mut self, dropout_rate: f64, variational: bool) -> Self {
113        for cell in &mut self.forward_cells {
114            *cell = cell
115                .clone()
116                .with_recurrent_dropout(dropout_rate, variational);
117        }
118        for cell in &mut self.backward_cells {
119            *cell = cell
120                .clone()
121                .with_recurrent_dropout(dropout_rate, variational);
122        }
123        self
124    }
125
126    pub fn with_output_dropout(mut self, dropout_rate: f64) -> Self {
127        // Apply output dropout to all layers except the last
128        for (i, cell) in self.forward_cells.iter_mut().enumerate() {
129            if i < self.num_layers - 1 {
130                *cell = cell.clone().with_output_dropout(dropout_rate);
131            }
132        }
133        for (i, cell) in self.backward_cells.iter_mut().enumerate() {
134            if i < self.num_layers - 1 {
135                *cell = cell.clone().with_output_dropout(dropout_rate);
136            }
137        }
138        self
139    }
140
141    pub fn with_cell_update_dropout(mut self, dropout_rate: f64, variational: bool) -> Self {
142        for cell in self
143            .forward_cells
144            .iter_mut()
145            .chain(self.backward_cells.iter_mut())
146        {
147            *cell = cell
148                .clone()
149                .with_cell_update_dropout(dropout_rate, variational);
150        }
151        self
152    }
153
154    pub fn with_zoneout(mut self, cell_zoneout_rate: f64, hidden_zoneout_rate: f64) -> Self {
155        for cell in &mut self.forward_cells {
156            *cell = cell
157                .clone()
158                .with_zoneout(cell_zoneout_rate, hidden_zoneout_rate);
159        }
160        for cell in &mut self.backward_cells {
161            *cell = cell
162                .clone()
163                .with_zoneout(cell_zoneout_rate, hidden_zoneout_rate);
164        }
165        self
166    }
167
168    pub fn train(&mut self) {
169        self.is_training = true;
170        for cell in &mut self.forward_cells {
171            cell.train();
172        }
173        for cell in &mut self.backward_cells {
174            cell.train();
175        }
176    }
177
178    pub fn eval(&mut self) {
179        self.is_training = false;
180        for cell in &mut self.forward_cells {
181            cell.eval();
182        }
183        for cell in &mut self.backward_cells {
184            cell.eval();
185        }
186    }
187
188    /// Combine forward and backward outputs according to the combine mode
189    fn combine_outputs(&self, forward: &Array2<f64>, backward: &Array2<f64>) -> Array2<f64> {
190        match self.combine_mode {
191            CombineMode::Concat => {
192                // Stack forward and backward outputs vertically
193                let mut combined =
194                    Array2::zeros((forward.nrows() + backward.nrows(), forward.ncols()));
195                combined
196                    .slice_mut(ndarray::s![..forward.nrows(), ..])
197                    .assign(forward);
198                combined
199                    .slice_mut(ndarray::s![forward.nrows().., ..])
200                    .assign(backward);
201                combined
202            }
203            CombineMode::Sum => forward + backward,
204            CombineMode::Average => (forward + backward) * 0.5,
205        }
206    }
207
208    /// Forward pass for a complete sequence
209    ///
210    /// Each layer runs one LSTM from start to end and another from end to start;
211    /// their outputs are combined and fed to the next layer.
212    pub fn forward_sequence(&mut self, sequence: &[Array2<f64>]) -> Vec<Array2<f64>> {
213        self.forward_sequence_with_cache(sequence).0
214    }
215
216    /// Forward pass with caching for training
217    pub fn forward_sequence_with_cache(
218        &mut self,
219        sequence: &[Array2<f64>],
220    ) -> (Vec<Array2<f64>>, BiLSTMNetworkCache) {
221        let seq_len = sequence.len();
222        let batch_size = sequence.first().map_or(1, |x| x.ncols());
223        let mut cache = BiLSTMNetworkCache {
224            forward_caches: Vec::with_capacity(self.num_layers),
225            backward_caches: Vec::with_capacity(self.num_layers),
226            forward_output_masks: Vec::with_capacity(self.num_layers),
227            backward_output_masks: Vec::with_capacity(self.num_layers),
228        };
229        if seq_len == 0 {
230            return (Vec::new(), cache);
231        }
232
233        let last = self.num_layers - 1;
234        let mut layer_input = sequence.to_vec();
235        for l in 0..self.num_layers {
236            let order: Vec<usize> = (0..seq_len).collect();
237            let reversed: Vec<usize> = (0..seq_len).rev().collect();
238            let (f_out, f_cache, f_masks) = run_direction(
239                &mut self.forward_cells[l],
240                &layer_input,
241                &order,
242                batch_size,
243                l < last,
244            );
245            let (b_out, b_cache, b_masks) = run_direction(
246                &mut self.backward_cells[l],
247                &layer_input,
248                &reversed,
249                batch_size,
250                l < last,
251            );
252            layer_input = f_out
253                .iter()
254                .zip(&b_out)
255                .map(|(f, b)| self.combine_outputs(f, b))
256                .collect();
257            cache.forward_caches.push(f_cache);
258            cache.backward_caches.push(b_cache);
259            cache.forward_output_masks.push(f_masks);
260            cache.backward_output_masks.push(b_masks);
261        }
262        (layer_input, cache)
263    }
264
265    /// Backpropagation through time for both directions.
266    ///
267    /// `d_outputs[t]` is the gradient of the loss with respect to the combined output
268    /// at step `t`. Returns the gradients of the forward cells, of the backward cells
269    /// and of every input.
270    pub fn backward_sequence(
271        &self,
272        d_outputs: &[Array2<f64>],
273        cache: &BiLSTMNetworkCache,
274    ) -> (
275        Vec<LSTMCellGradients>,
276        Vec<LSTMCellGradients>,
277        Vec<Array2<f64>>,
278    ) {
279        let (mut f_grads, mut b_grads) = self.zero_gradients();
280        let seq_len = d_outputs.len();
281        let mut d_layer = d_outputs.to_vec();
282        let h = self.hidden_size;
283
284        for l in (0..self.num_layers).rev() {
285            let (d_f, d_b): (Vec<_>, Vec<_>) = d_layer
286                .iter()
287                .enumerate()
288                .map(|(t, d)| {
289                    let (df, db) = match self.combine_mode {
290                        CombineMode::Concat => (
291                            d.slice(ndarray::s![..h, ..]).to_owned(),
292                            d.slice(ndarray::s![h.., ..]).to_owned(),
293                        ),
294                        CombineMode::Sum => (d.clone(), d.clone()),
295                        CombineMode::Average => (d * 0.5, d * 0.5),
296                    };
297                    (
298                        masked(&df, &cache.forward_output_masks[l][t]),
299                        masked(&db, &cache.backward_output_masks[l][t]),
300                    )
301                })
302                .unzip();
303
304            let forward_bptt: Vec<usize> = (0..seq_len).rev().collect();
305            let backward_bptt: Vec<usize> = (0..seq_len).collect();
306            let dx_f = backprop_direction(
307                &self.forward_cells[l],
308                &d_f,
309                &cache.forward_caches[l],
310                &forward_bptt,
311                &mut f_grads[l],
312            );
313            let dx_b = backprop_direction(
314                &self.backward_cells[l],
315                &d_b,
316                &cache.backward_caches[l],
317                &backward_bptt,
318                &mut b_grads[l],
319            );
320            d_layer = dx_f.iter().zip(&dx_b).map(|(a, b)| a + b).collect();
321        }
322        (f_grads, b_grads, d_layer)
323    }
324
325    /// Get references to forward and backward cells for serialization
326    pub fn get_forward_cells(&self) -> &[LSTMCell] {
327        &self.forward_cells
328    }
329
330    pub fn get_backward_cells(&self) -> &[LSTMCell] {
331        &self.backward_cells
332    }
333
334    /// Get mutable references for training mode changes
335    pub fn get_forward_cells_mut(&mut self) -> &mut [LSTMCell] {
336        &mut self.forward_cells
337    }
338
339    pub fn get_backward_cells_mut(&mut self) -> &mut [LSTMCell] {
340        &mut self.backward_cells
341    }
342
343    /// Update parameters for both directions
344    pub fn update_parameters<O: Optimizer>(
345        &mut self,
346        forward_gradients: &[LSTMCellGradients],
347        backward_gradients: &[LSTMCellGradients],
348        optimizer: &mut O,
349    ) {
350        // Update forward cells
351        for (i, (cell, gradients)) in self
352            .forward_cells
353            .iter_mut()
354            .zip(forward_gradients.iter())
355            .enumerate()
356        {
357            cell.update_parameters(gradients, optimizer, &format!("forward_layer_{}", i));
358        }
359
360        // Update backward cells
361        for (i, (cell, gradients)) in self
362            .backward_cells
363            .iter_mut()
364            .zip(backward_gradients.iter())
365            .enumerate()
366        {
367            cell.update_parameters(gradients, optimizer, &format!("backward_layer_{}", i));
368        }
369    }
370
371    /// Zero gradients for all cells
372    pub fn zero_gradients(&self) -> (Vec<LSTMCellGradients>, Vec<LSTMCellGradients>) {
373        let forward_gradients: Vec<_> = self
374            .forward_cells
375            .iter()
376            .map(|cell| cell.zero_gradients())
377            .collect();
378
379        let backward_gradients: Vec<_> = self
380            .backward_cells
381            .iter()
382            .map(|cell| cell.zero_gradients())
383            .collect();
384
385        (forward_gradients, backward_gradients)
386    }
387}
388
389/// Runs one direction of a layer over `inputs` in the given time order and returns
390/// outputs, caches and inter-layer dropout masks, all in the original time order.
391fn run_direction(
392    cell: &mut LSTMCell,
393    inputs: &[Array2<f64>],
394    order: &[usize],
395    batch_size: usize,
396    drop_output: bool,
397) -> (
398    Vec<Array2<f64>>,
399    Vec<LSTMCellCache>,
400    Vec<Option<Array2<f64>>>,
401) {
402    cell.reset_dropout_masks();
403    let n = inputs.len();
404    let mut outputs = vec![Array2::zeros((0, 0)); n];
405    let mut caches: Vec<Option<LSTMCellCache>> = vec![None; n];
406    let mut masks = vec![None; n];
407    let mut hx = Array2::zeros((cell.hidden_size, batch_size));
408    let mut cx = Array2::zeros((cell.hidden_size, batch_size));
409    for &t in order {
410        let (hy, cy, cache) = cell.forward_with_cache(&inputs[t], &hx, &cx);
411        let mask = if drop_output {
412            cell.output_dropout_mask(hy.raw_dim())
413        } else {
414            None
415        };
416        outputs[t] = masked(&hy, &mask);
417        masks[t] = mask;
418        caches[t] = Some(cache);
419        hx = hy;
420        cx = cy;
421    }
422    (outputs, caches.into_iter().flatten().collect(), masks)
423}
424
425/// BPTT for one direction; `order` visits the steps in reverse processing order.
426fn backprop_direction(
427    cell: &LSTMCell,
428    d_outputs: &[Array2<f64>],
429    caches: &[LSTMCellCache],
430    order: &[usize],
431    gradients: &mut LSTMCellGradients,
432) -> Vec<Array2<f64>> {
433    let batch_size = d_outputs.first().map_or(1, |d| d.ncols());
434    let mut dh_next = Array2::zeros((cell.hidden_size, batch_size));
435    let mut dc_next = Array2::zeros((cell.hidden_size, batch_size));
436    let mut d_inputs = vec![Array2::zeros((0, 0)); d_outputs.len()];
437    for &t in order {
438        let dh = &d_outputs[t] + &dh_next;
439        let (g, dx, dhx, dcx) = cell.backward(&dh, &dc_next, &caches[t]);
440        gradients.accumulate(&g);
441        dh_next = dhx;
442        dc_next = dcx;
443        d_inputs[t] = dx;
444    }
445    d_inputs
446}
447
448#[cfg(test)]
449mod tests {
450    use super::*;
451    use ndarray::arr2;
452
453    #[test]
454    fn test_bilstm_creation() {
455        let network = BiLSTMNetwork::new_concat(3, 5, 2);
456        assert_eq!(network.input_size, 3);
457        assert_eq!(network.hidden_size, 5);
458        assert_eq!(network.num_layers, 2);
459        assert_eq!(network.output_size(), 10); // 2 * hidden_size for concat mode
460    }
461
462    #[test]
463    fn test_bilstm_combine_modes() {
464        let forward = arr2(&[[1.0], [2.0]]);
465        let backward = arr2(&[[3.0], [4.0]]);
466
467        let concat_network = BiLSTMNetwork::new_concat(2, 2, 1);
468        let concat_result = concat_network.combine_outputs(&forward, &backward);
469        assert_eq!(concat_result.shape(), &[4, 1]);
470        assert_eq!(concat_result[[0, 0]], 1.0);
471        assert_eq!(concat_result[[1, 0]], 2.0);
472        assert_eq!(concat_result[[2, 0]], 3.0);
473        assert_eq!(concat_result[[3, 0]], 4.0);
474
475        let sum_network = BiLSTMNetwork::new_sum(2, 2, 1);
476        let sum_result = sum_network.combine_outputs(&forward, &backward);
477        assert_eq!(sum_result.shape(), &[2, 1]);
478        assert_eq!(sum_result[[0, 0]], 4.0);
479        assert_eq!(sum_result[[1, 0]], 6.0);
480
481        let avg_network = BiLSTMNetwork::new_average(2, 2, 1);
482        let avg_result = avg_network.combine_outputs(&forward, &backward);
483        assert_eq!(avg_result.shape(), &[2, 1]);
484        assert_eq!(avg_result[[0, 0]], 2.0);
485        assert_eq!(avg_result[[1, 0]], 3.0);
486    }
487
488    #[test]
489    fn test_bilstm_forward_sequence() {
490        let mut network = BiLSTMNetwork::new_concat(2, 3, 1);
491
492        let sequence = vec![
493            arr2(&[[1.0], [0.5]]),
494            arr2(&[[0.8], [0.2]]),
495            arr2(&[[0.3], [0.9]]),
496        ];
497
498        let outputs = network.forward_sequence(&sequence);
499
500        assert_eq!(outputs.len(), 3);
501        for output in &outputs {
502            assert_eq!(output.shape(), &[6, 1]); // 2 * hidden_size for concat
503        }
504    }
505
506    #[test]
507    fn test_bilstm_training_mode() {
508        let mut network = BiLSTMNetwork::new_concat(2, 3, 1)
509            .with_input_dropout(0.1, false)
510            .with_recurrent_dropout(0.2, true);
511
512        // Test mode switching
513        network.train();
514        assert!(network.is_training);
515
516        network.eval();
517        assert!(!network.is_training);
518    }
519}
rust_lstm/layers/bilstm_network.rs

rust_lstm/layers/
bilstm_network.rs
