Struct ALiBi 

pub struct ALiBi { /* private fields */ }

ALiBi (Attention with Linear Biases) (Press et al., 2022).

Adds linear position biases directly to attention scores instead of positional embeddings. Enables length extrapolation.

Method

attention = softmax(Q @ K^T / sqrt(d) - m * |i - j|)

where m is a head-specific slope.

Reference

• Press, O., et al. (2022). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR.

Implementations

impl ALiBi

pub fn new(num_heads: usize) -> Self

Create ALiBi with specified number of attention heads.

Slopes follow geometric sequence: 2^(-8/n), 2^(-16/n), …

pub fn compute_bias(&self, seq_len: usize) -> Tensor

Compute ALiBi bias matrix for given sequence length.

Arguments

• seq_len - Current sequence length

Returns

Bias tensor [num_heads, seq_len, seq_len] to add to attention scores.

pub fn apply(&self, scores: &Tensor) -> Tensor

Apply ALiBi to attention scores.

Arguments

• scores - Attention scores [batch, num_heads, seq_len, seq_len]

Returns

Scores with ALiBi bias applied.

pub fn slopes(&self) -> &[f32]

Get slopes for each head.

pub fn num_heads(&self) -> usize

Get number of heads.

Trait Implementations

impl Clone for ALiBi

fn clone(&self) -> ALiBi

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for ALiBi

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

Formats the value using the given formatter.