Skip to main content

gaussian_attention

Function gaussian_attention 

Source
pub fn gaussian_attention(
    queries: &[f64],
    keys: &[f64],
    values: &[f64],
    nq: usize,
    nk: usize,
    d: usize,
    sigma: f64,
) -> Vec<f64>
Expand description

Gaussian attention. Li et al. (2025), scKGBERT.

α_ij = exp(-||q_i - k_j||² / (2σ²)) / Σ_m exp(-||q_i - k_m||² / (2σ²))

Q [nq × d], K [nk × d], V [nk × d]. Output: [nq × d].