burn_nn/modules/transformer/
encoder.rs

1use burn_core as burn;
2
3use alloc::vec::Vec;
4
5use super::{PositionWiseFeedForward, PositionWiseFeedForwardConfig};
6use crate::{
7    Dropout, DropoutConfig, LayerNorm, LayerNormConfig,
8    attention::{MhaCache, MhaInput, MultiHeadAttention, MultiHeadAttentionConfig},
9    cache::TensorCache,
10};
11use burn::config::Config;
12use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay};
13use burn::tensor::{Bool, Tensor, backend::Backend};
14
15/// Configuration to create a [Transformer Encoder](TransformerEncoder) layer using the [init function](TransformerEncoderConfig::init).
16#[derive(Config, Debug)]
17pub struct TransformerEncoderConfig {
18    /// The size of the model.
19    pub d_model: usize,
20    /// The size of the position-wise feed-forward network.
21    pub d_ff: usize,
22    /// The number of attention heads.
23    pub n_heads: usize,
24    /// The number of layers.
25    pub n_layers: usize,
26    /// The dropout rate. Default: 0.1
27    #[config(default = 0.1)]
28    pub dropout: f64,
29    /// Layer norm will be applied first instead of after the other modules.
30    #[config(default = false)]
31    pub norm_first: bool,
32    /// Use "quiet softmax" instead of regular softmax.
33    ///
34    /// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head).
35    /// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression.
36    ///
37    /// Reference: <https://www.evanmiller.org/attention-is-off-by-one.html>
38    #[config(default = false)]
39    pub quiet_softmax: bool,
40    /// The type of function used to initialize neural network parameters
41    #[config(
42        default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}"
43    )]
44    pub initializer: Initializer,
45}
46
47/// The transformer encoder module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
48///
49/// # Params
50///
51/// - layers: transformer encoder layers with `d_model` input and output features.
52///
53/// Should be created using [TransformerEncoderConfig]
54#[derive(Module, Debug)]
55#[module(custom_display)]
56pub struct TransformerEncoder<B: Backend> {
57    /// The transformer encoder layers.
58    pub layers: Vec<TransformerEncoderLayer<B>>,
59
60    /// The size of the model.
61    pub d_model: usize,
62
63    /// The size of the position-wise feed-forward network.
64    pub d_ff: usize,
65
66    /// The number of attention heads.
67    pub n_heads: usize,
68
69    /// The number of layers.
70    pub n_layers: usize,
71
72    /// The dropout rate. Default: 0.1
73    pub dropout: f64,
74
75    /// Layer norm will be applied first instead of after the other modules.
76    pub norm_first: bool,
77
78    /// Use "quiet softmax" instead of regular softmax.
79    pub quiet_softmax: bool,
80}
81
82impl<B: Backend> ModuleDisplay for TransformerEncoder<B> {
83    fn custom_settings(&self) -> Option<DisplaySettings> {
84        DisplaySettings::new()
85            .with_new_line_after_attribute(false)
86            .optional()
87    }
88
89    fn custom_content(&self, content: Content) -> Option<Content> {
90        content
91            .add("d_model", &self.d_model)
92            .add("d_ff", &self.d_ff)
93            .add("n_heads", &self.n_heads)
94            .add("n_layers", &self.n_layers)
95            .add("dropout", &self.dropout)
96            .add("norm_first", &self.norm_first)
97            .add("quiet_softmax", &self.quiet_softmax)
98            .optional()
99    }
100}
101
102/// [Transformer Encoder](TransformerEncoder) forward pass input argument.
103#[derive(Debug)]
104pub struct TransformerEncoderInput<B: Backend> {
105    tensor: Tensor<B, 3>,
106    mask_pad: Option<Tensor<B, 2, Bool>>,
107    mask_attn: Option<Tensor<B, 3, Bool>>,
108}
109
110impl<B: Backend> TransformerEncoderInput<B> {
111    /// Create a [transformer encoder](TransformerEncoder) input argument.
112    pub fn new(tensor: Tensor<B, 3>) -> Self {
113        Self {
114            tensor,
115            mask_pad: None,
116            mask_attn: None,
117        }
118    }
119
120    /// Register the padding mask.
121    pub fn mask_pad(mut self, mask_pad: Tensor<B, 2, Bool>) -> Self {
122        self.mask_pad = Some(mask_pad);
123        self
124    }
125
126    /// Register the attention mask.
127    pub fn mask_attn(mut self, mask_attn: Tensor<B, 3, Bool>) -> Self {
128        self.mask_attn = Some(mask_attn);
129        self
130    }
131}
132impl TransformerEncoderConfig {
133    /// Initialize a new [transformer encoder](TransformerEncoder) module.
134    pub fn init<B: Backend>(&self, device: &B::Device) -> TransformerEncoder<B> {
135        let layers = (0..self.n_layers)
136            .map(|_| TransformerEncoderLayer::new(self, device))
137            .collect::<Vec<_>>();
138
139        TransformerEncoder {
140            layers,
141            d_model: self.d_model,
142            d_ff: self.d_ff,
143            n_heads: self.n_heads,
144            n_layers: self.n_layers,
145            dropout: self.dropout,
146            norm_first: self.norm_first,
147            quiet_softmax: self.quiet_softmax,
148        }
149    }
150}
151
152impl<B: Backend> TransformerEncoder<B> {
153    /// Applies the forward pass on the input tensor.
154    ///
155    /// # Shapes
156    ///
157    /// - tensor: `[batch_size, seq_length, d_model]`
158    /// - output: `[batch_size, seq_length, d_model]`
159    pub fn forward(&self, input: TransformerEncoderInput<B>) -> Tensor<B, 3> {
160        let mut x = input.tensor;
161
162        for layer in self.layers.iter() {
163            x = layer.forward(x, input.mask_pad.clone(), input.mask_attn.clone());
164        }
165
166        x
167    }
168    /// Applies the forward pass on the input tensor using autoregressive cache.
169    ///
170    /// # Shapes
171    ///
172    /// - tensor: `[batch_size, seq_length, d_model]`
173    /// - output: `[batch_size, seq_length, d_model]`
174    pub fn forward_autoregressive_inference(
175        &self,
176        input: TransformerEncoderInput<B>,
177        cache: &mut TransformerEncoderAutoregressiveCache<B>,
178    ) -> Tensor<B, 3> {
179        let mut x = input.tensor;
180
181        for i in 0..self.layers.len() {
182            let layer = self.layers.get(i).unwrap();
183            let cache = cache.layers.get_mut(i).unwrap();
184
185            x = layer.forward_autoregressive_inference(
186                x,
187                input.mask_pad.clone(),
188                input.mask_attn.clone(),
189                cache,
190            );
191        }
192
193        x
194    }
195
196    /// Create an empty autoregressive cache.
197    pub fn new_autoregressive_cache(&self) -> TransformerEncoderAutoregressiveCache<B> {
198        TransformerEncoderAutoregressiveCache::empty(self.layers.len())
199    }
200}
201
202/// Transformer encoder layer module.
203#[derive(Module, Debug)]
204pub struct TransformerEncoderLayer<B: Backend> {
205    mha: MultiHeadAttention<B>,
206    pwff: PositionWiseFeedForward<B>,
207    norm_1: LayerNorm<B>,
208    norm_2: LayerNorm<B>,
209    dropout: Dropout,
210    norm_first: bool,
211}
212
213impl<B: Backend> TransformerEncoderLayer<B> {
214    fn new(config: &TransformerEncoderConfig, device: &B::Device) -> Self {
215        let mha = MultiHeadAttentionConfig::new(config.d_model, config.n_heads)
216            .with_initializer(config.initializer.clone())
217            .with_dropout(config.dropout)
218            .with_quiet_softmax(config.quiet_softmax)
219            .init(device);
220        let norm_1 = LayerNormConfig::new(config.d_model).init(device);
221        let norm_2 = LayerNormConfig::new(config.d_model).init(device);
222        let dropout = DropoutConfig::new(config.dropout).init();
223        let pwff = PositionWiseFeedForwardConfig::new(config.d_model, config.d_ff)
224            .with_initializer(config.initializer.clone())
225            .with_dropout(config.dropout)
226            .init(device);
227
228        Self {
229            mha,
230            norm_1,
231            norm_2,
232            pwff,
233            dropout,
234            norm_first: config.norm_first,
235        }
236    }
237
238    fn forward(
239        &self,
240        input: Tensor<B, 3>,
241        mask_pad: Option<Tensor<B, 2, Bool>>,
242        mask_attn: Option<Tensor<B, 3, Bool>>,
243    ) -> Tensor<B, 3> {
244        // Multi-head attention residual path.
245        let x = input;
246        let mut residual_path = x.clone();
247
248        // Normalize.
249        if self.norm_first {
250            residual_path = self.norm_2.forward(residual_path)
251        }
252
253        // Multi-head attention.
254        let mut input_mhs = MhaInput::self_attn(residual_path);
255        if let Some(mask_pad) = mask_pad {
256            input_mhs = input_mhs.mask_pad(mask_pad);
257        }
258        if let Some(mask_attn) = mask_attn {
259            input_mhs = input_mhs.mask_attn(mask_attn);
260        }
261        let residual_path = self.mha.forward(input_mhs).context;
262
263        let residual_path = self.dropout.forward(residual_path);
264        let mut x = x + residual_path;
265
266        // Feed forward residual path.
267        // Normalize.
268        let residual_path = if self.norm_first {
269            self.norm_1.forward(x.clone())
270        } else {
271            x = self.norm_1.forward(x);
272            x.clone()
273        };
274
275        // Feed forward.
276        let residual_path = self.pwff.forward(residual_path);
277        let residual_path = self.dropout.forward(residual_path);
278        let mut x = x + residual_path;
279
280        // Main path.
281        // Normalize.
282        if !self.norm_first {
283            x = self.norm_2.forward(x)
284        }
285
286        x
287    }
288
289    fn forward_autoregressive_inference(
290        &self,
291        input: Tensor<B, 3>,
292        mask_pad: Option<Tensor<B, 2, Bool>>,
293        mask_attn: Option<Tensor<B, 3, Bool>>,
294        cache: &mut TransformerEncoderLayerAutoregressiveCache<B>,
295    ) -> Tensor<B, 3> {
296        // Multi-head attention residual path.
297        let x = input;
298        let mut residual_path = x.clone();
299
300        // Normalize.
301        if self.norm_first {
302            residual_path = cache
303                .norm_2
304                .forward_autoregressive(residual_path, 1, |x| self.norm_2.forward(x))
305        }
306
307        // Multi-head attention.
308        let mut input_mhs = MhaInput::self_attn(residual_path);
309        if let Some(mask_pad) = mask_pad {
310            input_mhs = input_mhs.mask_pad(mask_pad);
311        }
312        if let Some(mask_attn) = mask_attn {
313            input_mhs = input_mhs.mask_attn(mask_attn);
314        }
315        let residual_path = self.mha.forward_cache(input_mhs, &mut cache.mha).context;
316
317        let residual_path = self.dropout.forward(residual_path);
318        let mut x = x + residual_path;
319
320        // Feed forward residual path.
321        // Normalize.
322        let residual_path = if self.norm_first {
323            cache
324                .norm_1
325                .forward_autoregressive(x.clone(), 1, |x| self.norm_1.forward(x))
326        } else {
327            x = cache
328                .norm_1
329                .forward_autoregressive(x, 1, |x| self.norm_1.forward(x));
330            x.clone()
331        };
332
333        // Feed forward.
334        let residual_path = cache
335            .pwff
336            .forward_autoregressive(residual_path, 1, |x| self.pwff.forward(x));
337        let residual_path = self.dropout.forward(residual_path);
338        let mut x = x + residual_path;
339
340        // Main path.
341        // Normalize.
342        if !self.norm_first {
343            x = cache
344                .norm_2
345                .forward_autoregressive(x, 1, |x| self.norm_2.forward(x))
346        }
347
348        x
349    }
350}
351
352struct TransformerEncoderLayerAutoregressiveCache<B: Backend> {
353    mha: MhaCache<B>,
354    pwff: TensorCache<B, 3>,
355    norm_1: TensorCache<B, 3>,
356    norm_2: TensorCache<B, 3>,
357}
358
359impl<B: Backend> TransformerEncoderLayerAutoregressiveCache<B> {
360    fn empty() -> Self {
361        Self {
362            mha: MhaCache::autoregressive(),
363            pwff: TensorCache::empty(),
364            norm_1: TensorCache::empty(),
365            norm_2: TensorCache::empty(),
366        }
367    }
368}
369
370/// Autoregressive cache for the [Transformer Encoder](TransformerEncoder) layer.
371///
372/// To be used during inference when decoding tokens.
373pub struct TransformerEncoderAutoregressiveCache<B: Backend> {
374    layers: Vec<TransformerEncoderLayerAutoregressiveCache<B>>,
375}
376
377impl<B: Backend> TransformerEncoderAutoregressiveCache<B> {
378    fn empty(num_layers: usize) -> Self {
379        Self {
380            layers: (0..num_layers)
381                .map(|_| TransformerEncoderLayerAutoregressiveCache::empty())
382                .collect(),
383        }
384    }
385}
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390    use crate::{TestBackend, attention::generate_autoregressive_mask};
391    use burn::tensor::Distribution;
392    use burn::tensor::{Tolerance, ops::FloatElem};
393    type FT = FloatElem<TestBackend>;
394
395    #[test]
396    fn test_autoregressive_norm_last() {
397        let [d_model, d_ff, n_heads, num_layers] = [12, 24, 2, 3];
398        test_autoregressive(
399            TransformerEncoderConfig::new(d_model, d_ff, n_heads, num_layers)
400                .with_norm_first(false),
401        )
402    }
403
404    #[test]
405    fn test_autoregressive_norm_first() {
406        let [d_model, d_ff, n_heads, num_layers] = [12, 24, 2, 3];
407        test_autoregressive(
408            TransformerEncoderConfig::new(d_model, d_ff, n_heads, num_layers).with_norm_first(true),
409        )
410    }
411
412    fn test_autoregressive(config: TransformerEncoderConfig) {
413        let [batch_size, seq_length, d_model] = [3, 4, config.d_model];
414        let device = Default::default();
415        let transformer = config.init(&device);
416
417        let tensor = Tensor::<TestBackend, 3>::random(
418            [batch_size, seq_length, d_model],
419            Distribution::Default,
420            &device,
421        );
422        let mask_attn = generate_autoregressive_mask(batch_size, seq_length, &tensor.device());
423        let input = TransformerEncoderInput::new(tensor.clone()).mask_attn(mask_attn);
424
425        let output_1 = transformer.forward(input);
426        let mut output_2 = Vec::new();
427        let mut cache = transformer.new_autoregressive_cache();
428
429        for i in 1..seq_length + 1 {
430            let tensor = tensor.clone().slice([0..batch_size, 0..i, 0..d_model]);
431            let input = TransformerEncoderInput::new(tensor.clone());
432            let next_tok = transformer
433                .forward_autoregressive_inference(input, &mut cache)
434                .slice([0..batch_size, i - 1..i, 0..d_model]);
435            output_2.push(next_tok);
436        }
437
438        let output_2 = Tensor::cat(output_2, 1);
439
440        output_1
441            .into_data()
442            .assert_approx_eq::<FT>(&output_2.into_data(), Tolerance::permissive());
443    }
444
445    #[test]
446    fn display() {
447        let config = TransformerEncoderConfig::new(2, 4, 2, 3);
448        let transformer = config.init::<TestBackend>(&Default::default());
449
450        assert_eq!(
451            alloc::format!("{transformer}"),
452            "TransformerEncoder {d_model: 2, d_ff: 4, n_heads: 2, \
453            n_layers: 3, dropout: 0.1, norm_first: false, quiet_softmax: false, params: 162}"
454        );
455    }
456}
burn_nn/modules/transformer/encoder.rs

burn_nn/modules/transformer/
encoder.rs
