Struct MultiHeadAttention

pub struct MultiHeadAttention<F: Float + Debug> { /* private fields */ }

Multi-head attention layer as used in transformer architectures

This layer performs the attention operation described in “Attention Is All You Need” by Vaswani et al. It projects the queries, keys, and values into multiple heads, computes scaled dot-product attention for each head, concatenates the results, and projects the result back to the original dimension.

Examples

use scirs2_neural::layers::{MultiHeadAttention, Layer};
use scirs2_neural::layers::AttentionConfig;
use ndarray::Array3;
use rand::rngs::SmallRng;
use rand::SeedableRng;

// Create multi-head attention with 2 heads and 64-dim embeddings
let mut rng = SmallRng::seed_from_u64(42);
let config = AttentionConfig {
    num_heads: 2,
    head_dim: 32,
    dropout_prob: 0.0,
    causal: false,
    scale: None,
};

let mha = MultiHeadAttention::new(64, config, &mut rng).unwrap();

// Forward pass with a batch of 2 samples, sequence length 3
let batch_size = 2;
let seq_len = 3;
let d_model = 64;
let input = Array3::<f64>::from_elem((batch_size, seq_len, d_model), 0.1).into_dyn();
let output = mha.forward(&input).unwrap();

// Output shape should match input shape
assert_eq!(output.shape(), input.shape());

Implementations

impl<F: Float + Debug + ScalarOperand + 'static> MultiHeadAttention<F>

pub fn new<R: Rng>( d_model: usize, config: AttentionConfig, rng: &mut R, ) -> Result<Self>

Create a new multi-head attention layer

Arguments

• d_model - Embedding dimension
• config - Attention configuration
• rng - Random number generator for weight initialization

Returns

• A new multi-head attention layer

Trait Implementations

impl<F: Float + Debug + ScalarOperand + 'static> Clone for MultiHeadAttention<F>

fn clone(&self) -> Self

Returns a duplicate of the value.

const fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<F: Debug + Float + Debug> Debug for MultiHeadAttention<F>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<F: Float + Debug + ScalarOperand + 'static> Layer<F> for MultiHeadAttention<F>

fn as_any(&self) -> &dyn Any

Get the layer as a dyn Any for downcasting

fn as_any_mut(&mut self) -> &mut dyn Any

Get the layer as a mutable dyn Any for downcasting

fn forward(&self, input: &Array<F, IxDyn>) -> Result<Array<F, IxDyn>>

Forward pass of the layer

fn backward( &self, input: &Array<F, IxDyn>, _grad_output: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Backward pass of the layer to compute gradients

fn update(&mut self, learning_rate: F) -> Result<()>

Update the layer parameters with the given gradients

fn params(&self) -> Vec<Array<F, IxDyn>>

Get the parameters of the layer

fn gradients(&self) -> Vec<Array<F, IxDyn>>

Get the gradients of the layer parameters

fn set_gradients(&mut self, _gradients: &[Array<F, IxDyn>]) -> Result<()>

Set the gradients of the layer parameters

fn set_params(&mut self, _params: &[Array<F, IxDyn>]) -> Result<()>

Set the parameters of the layer

fn set_training(&mut self, _training: bool)

Set the layer to training mode (true) or evaluation mode (false)

fn is_training(&self) -> bool

Get the current training mode

fn layer_type(&self) -> &str

Get the type of the layer (e.g., “Dense”, “Conv2D”)

fn parameter_count(&self) -> usize

Get the number of trainable parameters in this layer

fn layer_description(&self) -> String

Get a detailed description of this layer

impl<F: Float + Debug + ScalarOperand + 'static> ParamLayer<F> for MultiHeadAttention<F>

fn get_parameters(&self) -> Vec<&Array<F, IxDyn>>

Get the parameters of the layer as a vector of arrays

fn get_gradients(&self) -> Vec<&Array<F, IxDyn>>

Get the gradients of the parameters

fn set_parameters(&mut self, params: Vec<Array<F, IxDyn>>) -> Result<()>

Set the parameters of the layer