Struct MultiHeadAttention 

pub struct MultiHeadAttention {
    pub d_model: usize,
    pub n_heads: usize,
    pub d_head: usize,
    pub w_q: Vec<f64>,
    pub w_k: Vec<f64>,
    pub w_v: Vec<f64>,
    pub w_o: Vec<f64>,
    pub b_o: Vec<f64>,
}

Multi-head attention module.

Projects Q, K, V with learned linear projections, runs n_heads parallel attention heads, then concatenates and projects the output.

All weight matrices are stored flat row-major.

Fields

d_model: usize

Model dimensionality.

n_heads: usize

Number of attention heads.

d_head: usize

Dimensionality per head: d_model / n_heads.

w_q: Vec<f64>

W_Q projection [d_model × d_model].

w_k: Vec<f64>

W_K projection [d_model × d_model].

w_v: Vec<f64>

W_V projection [d_model × d_model].

w_o: Vec<f64>

W_O output projection [d_model × d_model].

b_o: Vec<f64>

Output bias [d_model].

Implementations

impl MultiHeadAttention

pub fn new(d_model: usize, n_heads: usize) -> Self

Create a new MHA module with zero-initialised projections.

pub fn init_identity(&mut self)

Initialise W_Q, W_K, W_V, W_O with identity-like weights for testing.

pub fn forward(&self, x: &[f64], seq_len: usize) -> Vec<f64>

Forward pass.

x has shape [seq_len × d_model] (flat row-major). Returns output of shape [seq_len × d_model].

pub fn num_params(&self) -> usize

Total number of trainable parameters.

Trait Implementations

impl Clone for MultiHeadAttention

fn clone(&self) -> MultiHeadAttention

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for MultiHeadAttention

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

Formats the value using the given formatter.