Struct MultiHeadAttention 

pub struct MultiHeadAttention { /* private fields */ }

Multi-Head Attention (Vaswani et al., 2017).

Allows the model to jointly attend to information from different representation subspaces at different positions.

Example

let mha = MultiHeadAttention::new(512, 8);  // d_model=512, num_heads=8
let q = Tensor::randn(&[32, 10, 512]);
let k = Tensor::randn(&[32, 20, 512]);
let v = Tensor::randn(&[32, 20, 512]);
let (output, attn_weights) = mha.forward_qkv(&q, &k, &v, None);

Implementations

impl MultiHeadAttention

pub fn new(embed_dim: usize, num_heads: usize) -> Self

Create a new Multi-Head Attention layer.

Arguments

• embed_dim - Total dimension of the model (must be divisible by num_heads)
• num_heads - Number of attention heads

Panics

Panics if embed_dim is not divisible by num_heads.

pub fn with_dropout(self, dropout_p: f32) -> Self

Set dropout probability.

pub fn forward_qkv( &self, query: &Tensor, key: &Tensor, value: &Tensor, attn_mask: Option<&Tensor>, ) -> (Tensor, Tensor)

Forward pass with separate query, key, value inputs.

Arguments

• query - Query tensor [batch, target_len, embed_dim]
• key - Key tensor [batch, source_len, embed_dim]
• value - Value tensor [batch, source_len, embed_dim]
• attn_mask - Optional attention mask [batch, target_len, source_len]

Returns

Tuple of (output, attention_weights)

pub fn forward_self( &self, x: &Tensor, attn_mask: Option<&Tensor>, ) -> (Tensor, Tensor)

Self-attention: query, key, value are the same.

pub fn embed_dim(&self) -> usize

Get embed_dim.

pub fn num_heads(&self) -> usize

Get num_heads.

Trait Implementations

impl Debug for MultiHeadAttention

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

Formats the value using the given formatter.

impl Module for MultiHeadAttention

fn forward(&self, input: &Tensor) -> Tensor

Perform forward computation.

fn parameters(&self) -> Vec<&Tensor>

Get references to all learnable parameters.

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Get mutable references to all learnable parameters.

fn train(&mut self)

Set the module to training mode.

fn eval(&mut self)

Set the module to evaluation mode.

fn training(&self) -> bool

Check if the module is in training mode.

fn refresh_caches(&mut self)

Refresh any cached computations after parameters have been modified.

fn zero_grad(&mut self)

Zero out gradients for all parameters.

fn num_parameters(&self) -> usize

Get the number of learnable parameters.