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_zoneout(mut self, cell_zoneout_rate: f64, hidden_zoneout_rate: f64) -> Self {
142        for cell in &mut self.forward_cells {
143            *cell = cell
144                .clone()
145                .with_zoneout(cell_zoneout_rate, hidden_zoneout_rate);
146        }
147        for cell in &mut self.backward_cells {
148            *cell = cell
149                .clone()
150                .with_zoneout(cell_zoneout_rate, hidden_zoneout_rate);
151        }
152        self
153    }
154
155    pub fn train(&mut self) {
156        self.is_training = true;
157        for cell in &mut self.forward_cells {
158            cell.train();
159        }
160        for cell in &mut self.backward_cells {
161            cell.train();
162        }
163    }
164
165    pub fn eval(&mut self) {
166        self.is_training = false;
167        for cell in &mut self.forward_cells {
168            cell.eval();
169        }
170        for cell in &mut self.backward_cells {
171            cell.eval();
172        }
173    }
174
175    /// Combine forward and backward outputs according to the combine mode
176    fn combine_outputs(&self, forward: &Array2<f64>, backward: &Array2<f64>) -> Array2<f64> {
177        match self.combine_mode {
178            CombineMode::Concat => {
179                // Stack forward and backward outputs vertically
180                let mut combined =
181                    Array2::zeros((forward.nrows() + backward.nrows(), forward.ncols()));
182                combined
183                    .slice_mut(ndarray::s![..forward.nrows(), ..])
184                    .assign(forward);
185                combined
186                    .slice_mut(ndarray::s![forward.nrows().., ..])
187                    .assign(backward);
188                combined
189            }
190            CombineMode::Sum => forward + backward,
191            CombineMode::Average => (forward + backward) * 0.5,
192        }
193    }
194
195    /// Forward pass for a complete sequence
196    ///
197    /// Each layer runs one LSTM from start to end and another from end to start;
198    /// their outputs are combined and fed to the next layer.
199    pub fn forward_sequence(&mut self, sequence: &[Array2<f64>]) -> Vec<Array2<f64>> {
200        self.forward_sequence_with_cache(sequence).0
201    }
202
203    /// Forward pass with caching for training
204    pub fn forward_sequence_with_cache(
205        &mut self,
206        sequence: &[Array2<f64>],
207    ) -> (Vec<Array2<f64>>, BiLSTMNetworkCache) {
208        let seq_len = sequence.len();
209        let batch_size = sequence.first().map_or(1, |x| x.ncols());
210        let mut cache = BiLSTMNetworkCache {
211            forward_caches: Vec::with_capacity(self.num_layers),
212            backward_caches: Vec::with_capacity(self.num_layers),
213            forward_output_masks: Vec::with_capacity(self.num_layers),
214            backward_output_masks: Vec::with_capacity(self.num_layers),
215        };
216        if seq_len == 0 {
217            return (Vec::new(), cache);
218        }
219
220        let last = self.num_layers - 1;
221        let mut layer_input = sequence.to_vec();
222        for l in 0..self.num_layers {
223            let order: Vec<usize> = (0..seq_len).collect();
224            let reversed: Vec<usize> = (0..seq_len).rev().collect();
225            let (f_out, f_cache, f_masks) = run_direction(
226                &mut self.forward_cells[l],
227                &layer_input,
228                &order,
229                batch_size,
230                l < last,
231            );
232            let (b_out, b_cache, b_masks) = run_direction(
233                &mut self.backward_cells[l],
234                &layer_input,
235                &reversed,
236                batch_size,
237                l < last,
238            );
239            layer_input = f_out
240                .iter()
241                .zip(&b_out)
242                .map(|(f, b)| self.combine_outputs(f, b))
243                .collect();
244            cache.forward_caches.push(f_cache);
245            cache.backward_caches.push(b_cache);
246            cache.forward_output_masks.push(f_masks);
247            cache.backward_output_masks.push(b_masks);
248        }
249        (layer_input, cache)
250    }
251
252    /// Backpropagation through time for both directions.
253    ///
254    /// `d_outputs[t]` is the gradient of the loss with respect to the combined output
255    /// at step `t`. Returns the gradients of the forward cells, of the backward cells
256    /// and of every input.
257    pub fn backward_sequence(
258        &self,
259        d_outputs: &[Array2<f64>],
260        cache: &BiLSTMNetworkCache,
261    ) -> (
262        Vec<LSTMCellGradients>,
263        Vec<LSTMCellGradients>,
264        Vec<Array2<f64>>,
265    ) {
266        let (mut f_grads, mut b_grads) = self.zero_gradients();
267        let seq_len = d_outputs.len();
268        let mut d_layer = d_outputs.to_vec();
269        let h = self.hidden_size;
270
271        for l in (0..self.num_layers).rev() {
272            let (d_f, d_b): (Vec<_>, Vec<_>) = d_layer
273                .iter()
274                .enumerate()
275                .map(|(t, d)| {
276                    let (df, db) = match self.combine_mode {
277                        CombineMode::Concat => (
278                            d.slice(ndarray::s![..h, ..]).to_owned(),
279                            d.slice(ndarray::s![h.., ..]).to_owned(),
280                        ),
281                        CombineMode::Sum => (d.clone(), d.clone()),
282                        CombineMode::Average => (d * 0.5, d * 0.5),
283                    };
284                    (
285                        masked(&df, &cache.forward_output_masks[l][t]),
286                        masked(&db, &cache.backward_output_masks[l][t]),
287                    )
288                })
289                .unzip();
290
291            let forward_bptt: Vec<usize> = (0..seq_len).rev().collect();
292            let backward_bptt: Vec<usize> = (0..seq_len).collect();
293            let dx_f = backprop_direction(
294                &self.forward_cells[l],
295                &d_f,
296                &cache.forward_caches[l],
297                &forward_bptt,
298                &mut f_grads[l],
299            );
300            let dx_b = backprop_direction(
301                &self.backward_cells[l],
302                &d_b,
303                &cache.backward_caches[l],
304                &backward_bptt,
305                &mut b_grads[l],
306            );
307            d_layer = dx_f.iter().zip(&dx_b).map(|(a, b)| a + b).collect();
308        }
309        (f_grads, b_grads, d_layer)
310    }
311
312    /// Get references to forward and backward cells for serialization
313    pub fn get_forward_cells(&self) -> &[LSTMCell] {
314        &self.forward_cells
315    }
316
317    pub fn get_backward_cells(&self) -> &[LSTMCell] {
318        &self.backward_cells
319    }
320
321    /// Get mutable references for training mode changes
322    pub fn get_forward_cells_mut(&mut self) -> &mut [LSTMCell] {
323        &mut self.forward_cells
324    }
325
326    pub fn get_backward_cells_mut(&mut self) -> &mut [LSTMCell] {
327        &mut self.backward_cells
328    }
329
330    /// Update parameters for both directions
331    pub fn update_parameters<O: Optimizer>(
332        &mut self,
333        forward_gradients: &[LSTMCellGradients],
334        backward_gradients: &[LSTMCellGradients],
335        optimizer: &mut O,
336    ) {
337        // Update forward cells
338        for (i, (cell, gradients)) in self
339            .forward_cells
340            .iter_mut()
341            .zip(forward_gradients.iter())
342            .enumerate()
343        {
344            cell.update_parameters(gradients, optimizer, &format!("forward_layer_{}", i));
345        }
346
347        // Update backward cells
348        for (i, (cell, gradients)) in self
349            .backward_cells
350            .iter_mut()
351            .zip(backward_gradients.iter())
352            .enumerate()
353        {
354            cell.update_parameters(gradients, optimizer, &format!("backward_layer_{}", i));
355        }
356    }
357
358    /// Zero gradients for all cells
359    pub fn zero_gradients(&self) -> (Vec<LSTMCellGradients>, Vec<LSTMCellGradients>) {
360        let forward_gradients: Vec<_> = self
361            .forward_cells
362            .iter()
363            .map(|cell| cell.zero_gradients())
364            .collect();
365
366        let backward_gradients: Vec<_> = self
367            .backward_cells
368            .iter()
369            .map(|cell| cell.zero_gradients())
370            .collect();
371
372        (forward_gradients, backward_gradients)
373    }
374}
375
376/// Runs one direction of a layer over `inputs` in the given time order and returns
377/// outputs, caches and inter-layer dropout masks, all in the original time order.
378fn run_direction(
379    cell: &mut LSTMCell,
380    inputs: &[Array2<f64>],
381    order: &[usize],
382    batch_size: usize,
383    drop_output: bool,
384) -> (
385    Vec<Array2<f64>>,
386    Vec<LSTMCellCache>,
387    Vec<Option<Array2<f64>>>,
388) {
389    cell.reset_dropout_masks();
390    let n = inputs.len();
391    let mut outputs = vec![Array2::zeros((0, 0)); n];
392    let mut caches: Vec<Option<LSTMCellCache>> = vec![None; n];
393    let mut masks = vec![None; n];
394    let mut hx = Array2::zeros((cell.hidden_size, batch_size));
395    let mut cx = Array2::zeros((cell.hidden_size, batch_size));
396    for &t in order {
397        let (hy, cy, cache) = cell.forward_with_cache(&inputs[t], &hx, &cx);
398        let mask = if drop_output {
399            cell.output_dropout_mask(hy.raw_dim())
400        } else {
401            None
402        };
403        outputs[t] = masked(&hy, &mask);
404        masks[t] = mask;
405        caches[t] = Some(cache);
406        hx = hy;
407        cx = cy;
408    }
409    (outputs, caches.into_iter().flatten().collect(), masks)
410}
411
412/// BPTT for one direction; `order` visits the steps in reverse processing order.
413fn backprop_direction(
414    cell: &LSTMCell,
415    d_outputs: &[Array2<f64>],
416    caches: &[LSTMCellCache],
417    order: &[usize],
418    gradients: &mut LSTMCellGradients,
419) -> Vec<Array2<f64>> {
420    let batch_size = d_outputs.first().map_or(1, |d| d.ncols());
421    let mut dh_next = Array2::zeros((cell.hidden_size, batch_size));
422    let mut dc_next = Array2::zeros((cell.hidden_size, batch_size));
423    let mut d_inputs = vec![Array2::zeros((0, 0)); d_outputs.len()];
424    for &t in order {
425        let dh = &d_outputs[t] + &dh_next;
426        let (g, dx, dhx, dcx) = cell.backward(&dh, &dc_next, &caches[t]);
427        gradients.accumulate(&g);
428        dh_next = dhx;
429        dc_next = dcx;
430        d_inputs[t] = dx;
431    }
432    d_inputs
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438    use ndarray::arr2;
439
440    #[test]
441    fn test_bilstm_creation() {
442        let network = BiLSTMNetwork::new_concat(3, 5, 2);
443        assert_eq!(network.input_size, 3);
444        assert_eq!(network.hidden_size, 5);
445        assert_eq!(network.num_layers, 2);
446        assert_eq!(network.output_size(), 10); // 2 * hidden_size for concat mode
447    }
448
449    #[test]
450    fn test_bilstm_combine_modes() {
451        let forward = arr2(&[[1.0], [2.0]]);
452        let backward = arr2(&[[3.0], [4.0]]);
453
454        let concat_network = BiLSTMNetwork::new_concat(2, 2, 1);
455        let concat_result = concat_network.combine_outputs(&forward, &backward);
456        assert_eq!(concat_result.shape(), &[4, 1]);
457        assert_eq!(concat_result[[0, 0]], 1.0);
458        assert_eq!(concat_result[[1, 0]], 2.0);
459        assert_eq!(concat_result[[2, 0]], 3.0);
460        assert_eq!(concat_result[[3, 0]], 4.0);
461
462        let sum_network = BiLSTMNetwork::new_sum(2, 2, 1);
463        let sum_result = sum_network.combine_outputs(&forward, &backward);
464        assert_eq!(sum_result.shape(), &[2, 1]);
465        assert_eq!(sum_result[[0, 0]], 4.0);
466        assert_eq!(sum_result[[1, 0]], 6.0);
467
468        let avg_network = BiLSTMNetwork::new_average(2, 2, 1);
469        let avg_result = avg_network.combine_outputs(&forward, &backward);
470        assert_eq!(avg_result.shape(), &[2, 1]);
471        assert_eq!(avg_result[[0, 0]], 2.0);
472        assert_eq!(avg_result[[1, 0]], 3.0);
473    }
474
475    #[test]
476    fn test_bilstm_forward_sequence() {
477        let mut network = BiLSTMNetwork::new_concat(2, 3, 1);
478
479        let sequence = vec![
480            arr2(&[[1.0], [0.5]]),
481            arr2(&[[0.8], [0.2]]),
482            arr2(&[[0.3], [0.9]]),
483        ];
484
485        let outputs = network.forward_sequence(&sequence);
486
487        assert_eq!(outputs.len(), 3);
488        for output in &outputs {
489            assert_eq!(output.shape(), &[6, 1]); // 2 * hidden_size for concat
490        }
491    }
492
493    #[test]
494    fn test_bilstm_training_mode() {
495        let mut network = BiLSTMNetwork::new_concat(2, 3, 1)
496            .with_input_dropout(0.1, false)
497            .with_recurrent_dropout(0.2, true);
498
499        // Test mode switching
500        network.train();
501        assert!(network.is_training);
502
503        network.eval();
504        assert!(!network.is_training);
505    }
506}
rust_lstm/layers/bilstm_network.rs

rust_lstm/layers/
bilstm_network.rs
