geographdb_core/algorithms/

attention.rs

use crate::acceleration::cpu_kernels;
use crate::algorithms::mlp::{cross_entropy_from_logits, MlpClassifier, MlpGradients};

/// Single-head graph-attention classifier.
///
/// Each context token attends to itself plus its pre-computed geometric
/// neighbors.  The attended representation of the *last* context token is
/// passed through a two-layer MLP head for classification.
#[derive(Debug, Clone)]
pub struct GraphAttentionClassifier {
    vocab_size: usize,
    embed_dim: usize,
    num_neighbors: usize,
    /// Optional RoPE theta. When `Some(theta)`, context-token embeddings are
    /// rotated by their sequence distance before computing Q/K/V.
    rope_theta: Option<f32>,
    /// Token embedding table, row-major `[vocab_size, embed_dim]`.
    pub embedding: Vec<f32>,
    /// Query projection, row-major `[embed_dim, embed_dim]`.
    pub w_q: Vec<f32>,
    /// Key projection, row-major `[embed_dim, embed_dim]`.
    pub w_k: Vec<f32>,
    /// Value projection, row-major `[embed_dim, embed_dim]`.
    pub w_v: Vec<f32>,
    /// Learnable raw edge weights, row-major `[vocab_size, num_neighbors]`.
    /// The actual positive weight is `exp(raw)` so gradients are well behaved
    /// and weights stay non-negative.  When plasticity is disabled this vector
    /// is empty.
    pub edge_weights_raw: Vec<f32>,
    /// MLP head: embed_dim -> hidden_dim -> output_dim.
    pub mlp: MlpClassifier,
    /// When true, attention is restricted to the current token plus its
    /// geometric neighbors.  This drops the full self-attention term and
    /// reduces complexity from O(L^2) to O(L * k).  Defaults to false to
    /// preserve the original hybrid attention behaviour.
    geometric_attention_only: bool,
}

/// Gradients for [`GraphAttentionClassifier`].
#[derive(Debug, Clone)]
pub struct GraphAttentionGradients {
    pub dembedding: Vec<f32>,
    pub dw_q: Vec<f32>,
    pub dw_k: Vec<f32>,
    pub dw_v: Vec<f32>,
    pub dedge_weights_raw: Vec<f32>,
    pub mlp: MlpGradients,
}

/// Configuration for the tiebreak residual-correction prototype.
#[derive(Debug, Clone, Copy)]
pub struct TiebreakConfig {
    /// Positions whose maximum geometric attention weight is below this
    /// threshold are recomputed with full hybrid attention.
    pub min_max_weight: f32,
    /// Minimum gap between the top two geometric attention weights.  If the
    /// gap is smaller than this, the position is marked uncertain.
    pub min_margin: f32,
    /// Maximum box-counting dimension of the local geometric-attention signal
    /// for a position to be kept as geometric-only.  A sharply peaked or smooth
    /// (low-dimensional) pattern is considered reliable; a spread out or
    /// rapidly varying (high-dimensional) pattern triggers a full hybrid
    /// recompute.  `f32::INFINITY` disables the fractal criterion.
    pub max_fractal_dimension: f32,
    /// Size of the sliding window used to estimate the fractal dimension of the
    /// per-position maximum geometric attention weight sequence.
    pub fractal_window_size: usize,
}

impl Default for TiebreakConfig {
    fn default() -> Self {
        Self {
            min_max_weight: 0.5,
            min_margin: 0.1,
            max_fractal_dimension: f32::INFINITY,
            fractal_window_size: 8,
        }
    }
}

/// Multiply a row vector `[cols]` by a matrix `[cols, out]` to get `[out]`.
fn vec_matmul(v: &[f32], cols: usize, m: &[f32], out: usize) -> Vec<f32> {
    assert_eq!(v.len(), cols);
    assert_eq!(m.len(), cols * out);
    let mut res = vec![0.0f32; out];
    for k in 0..cols {
        for j in 0..out {
            res[j] += v[k] * m[k * out + j];
        }
    }
    res
}

/// Multiply a transposed matrix `[rows, cols]^T` by a column vector `[rows]` to
/// get `[cols]`.
fn mat_t_vec(m: &[f32], rows: usize, cols: usize, v: &[f32]) -> Vec<f32> {
    assert_eq!(m.len(), rows * cols);
    assert_eq!(v.len(), rows);
    let mut res = vec![0.0f32; cols];
    for i in 0..rows {
        for j in 0..cols {
            res[j] += m[i * cols + j] * v[i];
        }
    }
    res
}

fn softmax(scores: &mut [f32]) -> Vec<f32> {
    cpu_kernels::softmax_in_place(scores);
    scores.to_vec()
}

fn top_two(weights: &[f32]) -> (f32, f32) {
    let mut max = f32::NEG_INFINITY;
    let mut second = f32::NEG_INFINITY;
    for &w in weights {
        if w > max {
            second = max;
            max = w;
        } else if w > second {
            second = w;
        }
    }
    (max, second)
}

fn xavier_init(rows: usize, cols: usize, rng: &mut impl FnMut() -> f32) -> Vec<f32> {
    let scale = (2.0 / (rows + cols) as f32).sqrt();
    (0..rows * cols).map(|_| rng() * scale).collect()
}

/// Apply standard RoPE rotation to a single embedding vector in place.
///
/// `position` is the 1-based sequence distance from the target (e.g. 1 for the
/// previous token).  `embed_dim` must be even.
fn apply_rope_in_place(vec: &mut [f32], embed_dim: usize, position: usize, theta: f32) {
    cpu_kernels::apply_rope_in_place(vec, embed_dim, position, theta);
}

/// Apply the inverse RoPE rotation to a single embedding vector in place.
fn apply_rope_inv_in_place(vec: &mut [f32], embed_dim: usize, position: usize, theta: f32) {
    cpu_kernels::apply_rope_inv_in_place(vec, embed_dim, position, theta);
}

impl GraphAttentionClassifier {
    /// Create a fresh classifier with Xavier-like random weights.
    ///
    /// When `plasticity` is `true`, a learnable positive weight is attached to
    /// each geometric neighbor edge and updated during training.
    pub fn new(
        vocab_size: usize,
        embed_dim: usize,
        hidden_dim: usize,
        output_dim: usize,
        num_neighbors: usize,
        seed: u32,
        rope_theta: Option<f32>,
        plasticity: bool,
    ) -> Self {
        if let Some(theta) = rope_theta {
            assert!(
                theta > 0.0,
                "GraphAttentionClassifier: rope_theta must be positive"
            );
            assert!(
                embed_dim % 2 == 0,
                "GraphAttentionClassifier: RoPE requires even embed_dim"
            );
        }

        let mut rng_state = seed;
        let mut next_rand = || {
            rng_state ^= rng_state << 13;
            rng_state ^= rng_state >> 17;
            rng_state ^= rng_state << 5;
            (rng_state as f32 / u32::MAX as f32) - 0.5
        };

        let embedding = xavier_init(vocab_size, embed_dim, &mut next_rand);
        let w_q = xavier_init(embed_dim, embed_dim, &mut next_rand);
        let w_k = xavier_init(embed_dim, embed_dim, &mut next_rand);
        let w_v = xavier_init(embed_dim, embed_dim, &mut next_rand);
        let mlp = MlpClassifier::new(embed_dim, hidden_dim, output_dim, seed.wrapping_add(1));
        let edge_weights_raw = if plasticity {
            // Initialize exp(raw) == 1.0 so training starts from the original
            // unweighted graph and learns deviations.
            vec![0.0f32; vocab_size * num_neighbors.max(1)]
        } else {
            Vec::new()
        };

        Self {
            vocab_size,
            embed_dim,
            num_neighbors,
            rope_theta,
            embedding,
            w_q,
            w_k,
            w_v,
            edge_weights_raw,
            mlp,
            geometric_attention_only: false,
        }
    }

    /// Enable or disable geometric-only attention.
    ///
    /// In geometric-only mode each context token attends only to itself and
    /// its pre-computed geometric neighbors, removing the O(L^2) full
    /// self-attention term.  This is intended as a fast CPU-native variant
    /// for long contexts.
    pub fn set_geometric_attention_only(&mut self, enabled: bool) {
        self.geometric_attention_only = enabled;
    }

    fn plasticity_enabled(&self) -> bool {
        !self.edge_weights_raw.is_empty()
    }

    fn edge_weight(&self, token_idx: usize, nbr_pos: usize) -> f32 {
        if !self.plasticity_enabled() {
            return 1.0f32;
        }
        let idx = token_idx * self.num_neighbors + nbr_pos;
        self.edge_weights_raw[idx.min(self.edge_weights_raw.len() - 1)].exp()
    }

    /// Build the attended-index list for position `j` and the number of
    /// entries in that list that count as "context" positions (as opposed to
    /// geometric neighbours).  Edge weights are only applied to neighbours.
    fn build_attended_indices(
        &self,
        j: usize,
        n_context: usize,
        nbrs: &[usize],
        geometric_only: bool,
    ) -> (Vec<usize>, usize) {
        let mut attended = Vec::with_capacity(1 + n_context + self.num_neighbors);
        attended.push(j);
        if !geometric_only {
            for k in 0..n_context {
                if k != j {
                    attended.push(k);
                }
            }
        }
        for &nbr in nbrs.iter().take(self.num_neighbors) {
            attended.push(nbr);
        }
        let n_context_attended = if geometric_only { 1 } else { n_context };
        (attended, n_context_attended)
    }

    fn lookup(&self, token_id: u32) -> &[f32] {
        let idx = (token_id as usize).min(self.vocab_size - 1);
        &self.embedding[idx * self.embed_dim..(idx + 1) * self.embed_dim]
    }

    /// Forward pass for a single example.
    ///
    /// `token_ids` has length `n_context`.  `neighbors[j]` lists the dense
    /// neighbor indices for `token_ids[j]` (may be empty).  The token itself is
    /// implicitly included as the first attended element.
    pub fn forward(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> Vec<f32> {
        let h = self.attend(token_ids, neighbors, self.geometric_attention_only);
        let last = h.last().expect("empty context");
        self.mlp.forward(last).logits
    }

    /// Forward pass using the full hybrid attention graph.
    pub fn forward_hybrid(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> Vec<f32> {
        let h = self.attend(token_ids, neighbors, false);
        let last = h.last().expect("empty context");
        self.mlp.forward(last).logits
    }

    /// Forward pass using only geometric (neighbor) attention.
    pub fn forward_geometric_only(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> Vec<f32> {
        let h = self.attend(token_ids, neighbors, true);
        let last = h.last().expect("empty context");
        self.mlp.forward(last).logits
    }

    /// Forward pass that also returns the MLP hidden state for the last
    /// context position.  This is useful for diagnosing the FFN head, e.g.
    /// low-rank approximation experiments that want real hidden-state
    /// activations rather than random probes.
    pub fn forward_hidden(
        &self,
        token_ids: &[u32],
        neighbors: &[&[usize]],
    ) -> (Vec<f32>, Vec<f32>) {
        let h = self.attend(token_ids, neighbors, self.geometric_attention_only);
        let last = h.last().expect("empty context");
        let fwd = self.mlp.forward(last);
        (fwd.logits, fwd.hidden)
    }

    /// Hidden state for the last context position without the output projection.
    /// Used by confidence-gated FFN routing to decide whether to compute the
    /// cheap low-rank logits or the full dense logits.
    pub fn forward_hidden_only(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> Vec<f32> {
        let h = self.attend(token_ids, neighbors, self.geometric_attention_only);
        let last = h.last().expect("empty context");
        self.mlp.forward_hidden(last)
    }

    /// Cross-entropy loss for a single example.
    pub fn loss(&self, token_ids: &[u32], neighbors: &[&[usize]], target: usize) -> f32 {
        let logits = self.forward(token_ids, neighbors);
        cross_entropy_from_logits(&logits, target)
    }

    /// Predicted class for a single example.
    pub fn predict(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> usize {
        let logits = self.forward(token_ids, neighbors);
        logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap_or(0)
    }

    /// Compute the attended hidden states for all context positions.
    pub(crate) fn attend(
        &self,
        token_ids: &[u32],
        neighbors: &[&[usize]],
        geometric_only: bool,
    ) -> Vec<Vec<f32>> {
        let n_context = token_ids.len();
        assert_eq!(neighbors.len(), n_context);

        let d = self.embed_dim;
        let scale = (d as f32).sqrt();

        // Pre-compute embeddings and RoPE-rotated versions for context tokens.
        let embeddings: Vec<Vec<f32>> =
            token_ids.iter().map(|&t| self.lookup(t).to_vec()).collect();
        let mut rotated_embeddings = embeddings.clone();
        if let Some(theta) = self.rope_theta {
            for (j, emb) in rotated_embeddings.iter_mut().enumerate() {
                let position = n_context - j;
                apply_rope_in_place(emb, d, position, theta);
            }
        }

        let mut hiddens = Vec::with_capacity(n_context);

        for j in 0..n_context {
            let q = vec_matmul(&rotated_embeddings[j], d, &self.w_q, d);

            // Build attended set: self + (optionally all other context positions)
            // + geometric neighbors.  In geometric-only mode this becomes a pure
            // graph-attention step with O(k) complexity per position.
            let (attended_indices, n_context_attended) =
                self.build_attended_indices(j, n_context, neighbors[j], geometric_only);

            let mut scores = Vec::with_capacity(attended_indices.len());
            let mut values = Vec::with_capacity(attended_indices.len());

            for (pos, &idx) in attended_indices.iter().enumerate() {
                let e_owned: Vec<f32>;
                let e: &[f32] = if idx < n_context {
                    // Context tokens use RoPE-rotated embeddings so sequence
                    // position matters.
                    &rotated_embeddings[idx]
                } else {
                    // Geometric neighbors are not sequence positions; use the
                    // unrotated embedding.
                    let dense_idx = idx.min(self.vocab_size - 1);
                    e_owned = self.lookup(dense_idx as u32).to_vec();
                    &e_owned
                };
                let k = vec_matmul(e, d, &self.w_k, d);
                let v = vec_matmul(e, d, &self.w_v, d);

                let mut score: f32 =
                    q.iter().zip(k.iter()).map(|(a, b)| a * b).sum::<f32>() / scale;
                // Edge weights apply only to the geometric neighbours, which
                // appear after all context positions in `attended_indices`.
                if pos >= n_context_attended {
                    let nbr_pos = pos - n_context_attended;
                    let edge_w = self.edge_weight(token_ids[j] as usize, nbr_pos);
                    score *= edge_w;
                }
                scores.push(score);
                values.push(v);
            }

            let weights = softmax(&mut scores);
            let mut context = vec![0.0f32; d];
            for (w, v) in weights.iter().zip(values.iter()) {
                for i in 0..d {
                    context[i] += w * v[i];
                }
            }

            // Residual update uses the unrotated embedding.
            let h: Vec<f32> = embeddings[j]
                .iter()
                .zip(context.iter())
                .map(|(e, c)| e + c)
                .collect();
            hiddens.push(h);
        }

        hiddens
    }

    /// Compute attended hidden states using a per-position attention mode.
    ///
    /// `per_position_geometric[j]` selects geometric-only attention for
    /// position `j`; otherwise the full context is attended.  This is the
    /// primitive used to measure approximation divergence and to prototype
    /// residual correction tie-breaks.
    pub fn attend_mixed(
        &self,
        token_ids: &[u32],
        neighbors: &[&[usize]],
        per_position_geometric: &[bool],
    ) -> Vec<Vec<f32>> {
        let n_context = token_ids.len();
        assert_eq!(neighbors.len(), n_context);
        assert_eq!(per_position_geometric.len(), n_context);

        let d = self.embed_dim;
        let scale = (d as f32).sqrt();

        let embeddings: Vec<Vec<f32>> =
            token_ids.iter().map(|&t| self.lookup(t).to_vec()).collect();
        let mut rotated_embeddings = embeddings.clone();
        if let Some(theta) = self.rope_theta {
            for (j, emb) in rotated_embeddings.iter_mut().enumerate() {
                let position = n_context - j;
                apply_rope_in_place(emb, d, position, theta);
            }
        }

        let mut hiddens = Vec::with_capacity(n_context);

        for j in 0..n_context {
            let q = vec_matmul(&rotated_embeddings[j], d, &self.w_q, d);
            let geometric_only = per_position_geometric[j];
            let (attended_indices, n_context_attended) =
                self.build_attended_indices(j, n_context, neighbors[j], geometric_only);

            let mut scores = Vec::with_capacity(attended_indices.len());
            let mut values = Vec::with_capacity(attended_indices.len());

            for (pos, &idx) in attended_indices.iter().enumerate() {
                let e_owned: Vec<f32>;
                let e: &[f32] = if idx < n_context {
                    &rotated_embeddings[idx]
                } else {
                    let dense_idx = idx.min(self.vocab_size - 1);
                    e_owned = self.lookup(dense_idx as u32).to_vec();
                    &e_owned
                };
                let k = vec_matmul(e, d, &self.w_k, d);
                let v = vec_matmul(e, d, &self.w_v, d);

                let mut score: f32 =
                    q.iter().zip(k.iter()).map(|(a, b)| a * b).sum::<f32>() / scale;
                if pos >= n_context_attended {
                    let nbr_pos = pos - n_context_attended;
                    let edge_w = self.edge_weight(token_ids[j] as usize, nbr_pos);
                    score *= edge_w;
                }
                scores.push(score);
                values.push(v);
            }

            let weights = softmax(&mut scores);
            let mut context = vec![0.0f32; d];
            for (w, v) in weights.iter().zip(values.iter()) {
                for i in 0..d {
                    context[i] += w * v[i];
                }
            }

            let h: Vec<f32> = embeddings[j]
                .iter()
                .zip(context.iter())
                .map(|(e, c)| e + c)
                .collect();
            hiddens.push(h);
        }

        hiddens
    }

    /// Forward pass with a per-position mixed attention mode.
    pub fn forward_mixed(
        &self,
        token_ids: &[u32],
        neighbors: &[&[usize]],
        per_position_geometric: &[bool],
    ) -> Vec<f32> {
        let h = self.attend_mixed(token_ids, neighbors, per_position_geometric);
        let last = h.last().expect("empty context");
        self.mlp.forward(last).logits
    }

    /// Tiebreak-corrected forward pass.
    ///
    /// Runs cheap geometric-only attention, identifies uncertain positions by
    /// inspecting the geometric attention distribution, then recomputes those
    /// positions with full hybrid attention via [`Self::forward_mixed`].  This
    /// is the prototype residual-correction stage for the approximate
    /// transformer math experiment.
    pub fn forward_tiebreak(
        &self,
        token_ids: &[u32],
        neighbors: &[&[usize]],
        config: &TiebreakConfig,
    ) -> Vec<f32> {
        let mask = self.tiebreak_mask(token_ids, neighbors, config);
        self.forward_mixed(token_ids, neighbors, &mask)
    }

    /// Build the per-position mixed-attention mask used by
    /// [`Self::forward_tiebreak`] and [`Self::forward_mixed`].
    ///
    /// A returned value of `true` means "keep the cheap geometric-only result
    /// for this position"; `false` means "recompute this position with full
    /// hybrid attention".  This matches the semantics expected by
    /// [`Self::forward_mixed`].
    pub fn tiebreak_mask(
        &self,
        token_ids: &[u32],
        neighbors: &[&[usize]],
        config: &TiebreakConfig,
    ) -> Vec<bool> {
        let n_context = token_ids.len();
        assert_eq!(neighbors.len(), n_context);

        let d = self.embed_dim;
        let scale = (d as f32).sqrt();

        let embeddings: Vec<Vec<f32>> =
            token_ids.iter().map(|&t| self.lookup(t).to_vec()).collect();
        let mut rotated_embeddings = embeddings.clone();
        if let Some(theta) = self.rope_theta {
            for (j, emb) in rotated_embeddings.iter_mut().enumerate() {
                let position = n_context - j;
                apply_rope_in_place(emb, d, position, theta);
            }
        }

        // First pass: cheap geometric-only attention for every position.
        // Record classic attention-threshold uncertainty and the maximum
        // geometric attention weight per position.
        let mut max_weights = Vec::with_capacity(n_context);
        let mut attention_uncertain = Vec::with_capacity(n_context);
        for j in 0..n_context {
            let q = vec_matmul(&rotated_embeddings[j], d, &self.w_q, d);
            let (attended, n_context_attended) =
                self.build_attended_indices(j, n_context, neighbors[j], true);

            let mut scores = Vec::with_capacity(attended.len());
            for (pos, &idx) in attended.iter().enumerate() {
                let e_owned: Vec<f32>;
                let e: &[f32] = if idx < n_context {
                    &rotated_embeddings[idx]
                } else {
                    let dense_idx = idx.min(self.vocab_size - 1);
                    e_owned = self.lookup(dense_idx as u32).to_vec();
                    &e_owned
                };
                let k = vec_matmul(e, d, &self.w_k, d);
                let mut score: f32 =
                    q.iter().zip(k.iter()).map(|(a, b)| a * b).sum::<f32>() / scale;
                if pos >= n_context_attended {
                    let nbr_pos = pos - n_context_attended;
                    let edge_w = self.edge_weight(token_ids[j] as usize, nbr_pos);
                    score *= edge_w;
                }
                scores.push(score);
            }

            let weights = softmax(&mut scores);
            let (max, second) = top_two(&weights);
            max_weights.push(max);
            attention_uncertain
                .push(max < config.min_max_weight || (max - second) < config.min_margin);
        }

        // Second pass: apply the fractal-dimension criterion over a sliding
        // window of the maximum-weight signal.  A rapidly changing attention
        // landscape (high local dimension) is treated as unreliable.
        let mut mask = Vec::with_capacity(n_context);
        for j in 0..n_context {
            let start = j.saturating_sub(config.fractal_window_size.saturating_sub(1));
            let window = &max_weights[start..=j];
            let fractal_dim = if window.len() >= 4 {
                super::fractal::box_counting_dimension_1d(window)
            } else {
                // Not enough samples: assume line-like complexity so the
                // threshold must be explicitly set above one to force hybrid.
                1.0
            };
            let uncertain = attention_uncertain[j] || fractal_dim > config.max_fractal_dimension;
            mask.push(!uncertain);
        }

        mask
    }

    /// Forward pass for a batch.
    pub fn forward_batch(
        &self,
        token_ids: &[u32],
        neighbors: &[Vec<Vec<usize>>],
        batch_size: usize,
    ) -> Vec<f32> {
        let n_context = token_ids.len() / batch_size;
        let mut h_matrix = Vec::with_capacity(batch_size * self.embed_dim);
        for b in 0..batch_size {
            let start = b * n_context;
            let ids = &token_ids[start..start + n_context];
            let nbr_refs: Vec<&[usize]> = neighbors[b].iter().map(|v| v.as_slice()).collect();
            let h = self.attend(ids, &nbr_refs, self.geometric_attention_only);
            h_matrix.extend_from_slice(h.last().unwrap());
        }
        self.mlp
            .forward_batch(&h_matrix, batch_size)
            .logits
            .iter()
            .copied()
            .collect()
    }

    /// Back-propagation for a batch.
    ///
    /// `neighbors[b][j]` is the neighbor list for the j-th context token of
    /// example b.  Returns gradients averaged over the batch and the average
    /// cross-entropy loss.
    pub fn backward_batch(
        &self,
        token_ids: &[u32],
        neighbors: &[Vec<Vec<usize>>],
        targets: &[usize],
        batch_size: usize,
    ) -> (GraphAttentionGradients, f32) {
        let n_context = token_ids.len() / batch_size;
        let d = self.embed_dim;
        let scale = (d as f32).sqrt();

        // Forward cache.
        let mut embeddings_cache: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch_size);
        let mut rotated_embeddings_cache: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch_size);
        let mut queries_cache: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch_size);
        let mut values_cache: Vec<Vec<Vec<Vec<f32>>>> = Vec::with_capacity(batch_size);
        let mut scores_cache: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch_size);
        let mut weights_cache: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch_size);
        let mut hiddens_cache: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch_size);
        let mut h_matrix: Vec<f32> = Vec::with_capacity(batch_size * d);

        for b in 0..batch_size {
            let start = b * n_context;
            let ids = &token_ids[start..start + n_context];
            let nbr_refs: Vec<&[usize]> = neighbors[b].iter().map(|v| v.as_slice()).collect();

            let embeddings: Vec<Vec<f32>> = ids.iter().map(|&t| self.lookup(t).to_vec()).collect();
            let mut rotated_embeddings = embeddings.clone();
            if let Some(theta) = self.rope_theta {
                for (j, emb) in rotated_embeddings.iter_mut().enumerate() {
                    let position = n_context - j;
                    apply_rope_in_place(emb, d, position, theta);
                }
            }

            let mut queries = Vec::with_capacity(n_context);
            let mut batch_values = Vec::with_capacity(n_context);
            let mut batch_scores = Vec::with_capacity(n_context);
            let mut batch_weights = Vec::with_capacity(n_context);
            let mut hiddens = Vec::with_capacity(n_context);

            for j in 0..n_context {
                let q = vec_matmul(&rotated_embeddings[j], d, &self.w_q, d);

                let (attended, n_context_attended) = self.build_attended_indices(
                    j,
                    n_context,
                    nbr_refs[j],
                    self.geometric_attention_only,
                );

                let mut scores = Vec::with_capacity(attended.len());
                let mut values = Vec::with_capacity(attended.len());

                for (pos, &idx) in attended.iter().enumerate() {
                    let e_owned: Vec<f32>;
                    let e: &[f32] = if idx < n_context {
                        &rotated_embeddings[idx]
                    } else {
                        let dense_idx = idx.min(self.vocab_size - 1);
                        e_owned = self.lookup(dense_idx as u32).to_vec();
                        &e_owned
                    };
                    let k = vec_matmul(e, d, &self.w_k, d);
                    let v = vec_matmul(e, d, &self.w_v, d);
                    let mut score = q.iter().zip(k.iter()).map(|(a, b)| a * b).sum::<f32>() / scale;
                    if pos >= n_context_attended {
                        let nbr_pos = pos - n_context_attended;
                        let edge_w = self.edge_weight(ids[j] as usize, nbr_pos);
                        score *= edge_w;
                    }
                    scores.push(score);
                    values.push(v);
                }

                let weights = softmax(&mut scores);
                let mut context = vec![0.0f32; d];
                for (w, v) in weights.iter().zip(values.iter()) {
                    for i in 0..d {
                        context[i] += w * v[i];
                    }
                }

                let h: Vec<f32> = embeddings[j]
                    .iter()
                    .zip(context.iter())
                    .map(|(e, c)| e + c)
                    .collect();

                queries.push(q);
                batch_values.push(values);
                batch_scores.push(scores);
                batch_weights.push(weights);
                hiddens.push(h.clone());
                if j == n_context - 1 {
                    h_matrix.extend_from_slice(&h);
                }
            }

            embeddings_cache.push(embeddings);
            rotated_embeddings_cache.push(rotated_embeddings);
            queries_cache.push(queries);
            values_cache.push(batch_values);
            scores_cache.push(batch_scores);
            weights_cache.push(batch_weights);
            hiddens_cache.push(hiddens);
        }

        // MLP head backward.
        let (mlp_grad, dh_matrix, loss) = self.mlp.backward_batch(&h_matrix, targets, batch_size);

        // Backprop through attention.
        let mut dembedding = vec![0.0f32; self.vocab_size * d];
        let mut dw_q = vec![0.0f32; d * d];
        let mut dw_k = vec![0.0f32; d * d];
        let mut dw_v = vec![0.0f32; d * d];
        let mut dedge_weights_raw = vec![0.0f32; self.edge_weights_raw.len()];

        for b in 0..batch_size {
            let start = b * n_context;
            let ids = &token_ids[start..start + n_context];

            for j in 0..n_context {
                // Only the last context position contributes to the MLP head.
                let dh = if j == n_context - 1 {
                    &dh_matrix[b * d..(b + 1) * d]
                } else {
                    &[] as &[f32]
                };

                if dh.is_empty() {
                    continue;
                }
                let dh: Vec<f32> = dh.to_vec();

                let q = &queries_cache[b][j];
                let values = &values_cache[b][j];
                let weights = &weights_cache[b][j];

                let (attended, n_context_attended) = self.build_attended_indices(
                    j,
                    n_context,
                    &neighbors[b][j],
                    self.geometric_attention_only,
                );

                // dcontext = dh (residual: h = e + context)
                let dcontext = dh.clone();

                // dvalue_i = weight_i * dcontext
                // dweight_i = value_i · dcontext
                let mut dvalues: Vec<Vec<f32>> = Vec::with_capacity(attended.len());
                let mut dscores = Vec::with_capacity(attended.len());
                for (w, v) in weights.iter().zip(values.iter()) {
                    let dv: Vec<f32> = dcontext.iter().map(|&dc| w * dc).collect();
                    let dw: f32 = v
                        .iter()
                        .zip(dcontext.iter())
                        .map(|(vi, dci)| vi * dci)
                        .sum();
                    dvalues.push(dv);
                    dscores.push(dw);
                }

                // Backprop through softmax.
                // dscore_i = weight_i * (dscore_i - Σ_j weight_j * dscore_j)
                let weighted_sum: f32 = weights
                    .iter()
                    .zip(dscores.iter())
                    .map(|(w, ds)| w * ds)
                    .sum();
                let dscores: Vec<f32> = weights
                    .iter()
                    .zip(dscores.iter())
                    .map(|(w, ds)| w * (ds - weighted_sum))
                    .collect();

                // Backprop through the learnable edge weights.
                // score_i = (q·k_i / scale) * exp(raw_i), so
                // draw_i = dscore_i * score_i.
                if self.plasticity_enabled() {
                    let scores = &scores_cache[b][j];
                    let src_idx = ids[j] as usize;
                    for (pos, (&ds, &score)) in dscores.iter().zip(scores.iter()).enumerate() {
                        if pos < n_context_attended {
                            continue;
                        }
                        let nbr_pos = pos - n_context_attended;
                        let edge_idx = src_idx * self.num_neighbors + nbr_pos;
                        dedge_weights_raw[edge_idx] += ds * score;
                    }
                }

                // Backprop through score = q·k / scale.
                // dq = Σ_i dscore_i * k_i / scale
                // dk_i = dscore_i * q / scale
                let mut dq = vec![0.0f32; d];
                let mut dk_list: Vec<Vec<f32>> = Vec::with_capacity(attended.len());
                for (ds, _v) in dscores.iter().zip(values.iter()) {
                    // k_i is not cached; recompute from embedding and w_k below.
                    let idx = attended[dk_list.len()];
                    let e_owned: Vec<f32>;
                    let e: &[f32] = if idx < n_context {
                        &rotated_embeddings_cache[b][idx]
                    } else {
                        let dense_idx = idx.min(self.vocab_size - 1);
                        e_owned = self.lookup(dense_idx as u32).to_vec();
                        &e_owned
                    };
                    let k = vec_matmul(e, d, &self.w_k, d);
                    for i in 0..d {
                        dq[i] += ds * k[i] / scale;
                    }
                    let dk: Vec<f32> = q.iter().map(|&qi| ds * qi / scale).collect();
                    dk_list.push(dk);
                }

                // Backprop through projections.
                // de_j += dh (residual) — uses the unrotated embedding.
                for i in 0..d {
                    dembedding[ids[j] as usize * d + i] += dh[i];
                }

                // de_j (rotated) += dq · w_q^T
                let de_q_rotated = mat_t_vec(&self.w_q, d, d, &dq);
                // Apply inverse RoPE to get gradient w.r.t. unrotated embedding.
                let mut de_q = de_q_rotated;
                if let Some(theta) = self.rope_theta {
                    let position = n_context - j;
                    apply_rope_inv_in_place(&mut de_q, d, position, theta);
                }
                for i in 0..d {
                    dembedding[ids[j] as usize * d + i] += de_q[i];
                }

                // dw_q += rotated_e_j^T · dq
                for k in 0..d {
                    for l in 0..d {
                        dw_q[k * d + l] += rotated_embeddings_cache[b][j][k] * dq[l];
                    }
                }

                // For each attended token i:
                // de_i (rotated or not) += dvalue_i · w_v^T
                // de_i (rotated or not) += dk_i · w_k^T
                // dw_k += e_i^T · dk_i
                // dw_v += e_i^T · dvalue_i
                for (idx, (dv, dk)) in attended.iter().zip(dvalues.iter().zip(dk_list.iter())) {
                    let idx = *idx;
                    let is_context = idx < n_context;

                    let e_unrot_owned: Vec<f32>;
                    let e_for_proj: &[f32] = if is_context {
                        &rotated_embeddings_cache[b][idx]
                    } else {
                        let dense_idx = idx.min(self.vocab_size - 1);
                        e_unrot_owned = self.lookup(dense_idx as u32).to_vec();
                        &e_unrot_owned
                    };

                    let de_v_rotated = mat_t_vec(&self.w_v, d, d, dv);
                    let de_k_rotated = mat_t_vec(&self.w_k, d, d, dk);

                    let token_idx = if is_context {
                        ids[idx] as usize
                    } else {
                        idx.min(self.vocab_size - 1)
                    };

                    // If the attended token is a context token, its K/V used a
                    // RoPE-rotated embedding, so apply the inverse rotation.
                    let mut de_v = de_v_rotated;
                    let mut de_k = de_k_rotated;
                    if is_context {
                        if let Some(theta) = self.rope_theta {
                            let position = n_context - idx;
                            apply_rope_inv_in_place(&mut de_v, d, position, theta);
                            apply_rope_inv_in_place(&mut de_k, d, position, theta);
                        }
                    }

                    for i in 0..d {
                        dembedding[token_idx * d + i] += de_v[i] + de_k[i];
                    }

                    for k in 0..d {
                        for l in 0..d {
                            dw_k[k * d + l] += e_for_proj[k] * dk[l];
                            dw_v[k * d + l] += e_for_proj[k] * dv[l];
                        }
                    }
                }
            }
        }

        // Average gradients over batch.
        let inv_b = 1.0 / batch_size as f32;
        for v in dembedding.iter_mut() {
            *v *= inv_b;
        }
        for v in dw_q.iter_mut() {
            *v *= inv_b;
        }
        for v in dw_k.iter_mut() {
            *v *= inv_b;
        }
        for v in dw_v.iter_mut() {
            *v *= inv_b;
        }
        for v in dedge_weights_raw.iter_mut() {
            *v *= inv_b;
        }

        (
            GraphAttentionGradients {
                dembedding,
                dw_q,
                dw_k,
                dw_v,
                dedge_weights_raw,
                mlp: mlp_grad,
            },
            loss,
        )
    }

    /// Apply a single SGD step using the supplied gradients.
    pub fn apply_sgd(&mut self, grad: &GraphAttentionGradients, lr: f32) {
        for i in 0..self.embedding.len() {
            self.embedding[i] -= lr * grad.dembedding[i];
        }
        for i in 0..self.w_q.len() {
            self.w_q[i] -= lr * grad.dw_q[i];
        }
        for i in 0..self.w_k.len() {
            self.w_k[i] -= lr * grad.dw_k[i];
        }
        for i in 0..self.w_v.len() {
            self.w_v[i] -= lr * grad.dw_v[i];
        }
        for i in 0..self.edge_weights_raw.len() {
            self.edge_weights_raw[i] -= lr * grad.dedge_weights_raw[i];
        }
        self.mlp.apply_sgd(&grad.mlp, lr);
    }

    /// Flatten all learnable parameters into one vector, ordered
    /// `[embedding, w_q, w_k, w_v, edge_weights_raw, mlp]`.
    pub fn flatten_params(&self) -> Vec<f32> {
        let mut params = Vec::with_capacity(
            self.embedding.len()
                + self.w_q.len()
                + self.w_k.len()
                + self.w_v.len()
                + self.edge_weights_raw.len()
                + self.mlp.flatten_params().len(),
        );
        params.extend_from_slice(&self.embedding);
        params.extend_from_slice(&self.w_q);
        params.extend_from_slice(&self.w_k);
        params.extend_from_slice(&self.w_v);
        params.extend_from_slice(&self.edge_weights_raw);
        params.extend_from_slice(&self.mlp.flatten_params());
        params
    }

    /// Restore parameters from a flat vector produced by [`Self::flatten_params`].
    pub fn load_flat_params(&mut self, params: &[f32]) {
        let embed_len = self.embedding.len();
        let wq_len = self.w_q.len();
        let wk_len = self.w_k.len();
        let wv_len = self.w_v.len();
        let edge_len = self.edge_weights_raw.len();
        let mlp_len = self.mlp.flatten_params().len();
        let expected = embed_len + wq_len + wk_len + wv_len + edge_len + mlp_len;
        assert_eq!(params.len(), expected, "load_flat_params: size mismatch");

        let mut off = 0;
        self.embedding
            .copy_from_slice(&params[off..off + embed_len]);
        off += embed_len;
        self.w_q.copy_from_slice(&params[off..off + wq_len]);
        off += wq_len;
        self.w_k.copy_from_slice(&params[off..off + wk_len]);
        off += wk_len;
        self.w_v.copy_from_slice(&params[off..off + wv_len]);
        off += wv_len;
        self.edge_weights_raw
            .copy_from_slice(&params[off..off + edge_len]);
        off += edge_len;
        self.mlp.load_flat_params(&params[off..off + mlp_len]);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn forward_produces_logits() {
        let model = GraphAttentionClassifier::new(10, 8, 16, 3, 2, 1, None, false);
        let token_ids = vec![0u32, 1, 2];
        let neighbors_owned: Vec<Vec<usize>> = vec![vec![0], vec![1], vec![2]];
        let neighbors: Vec<&[usize]> = neighbors_owned.iter().map(|v| v.as_slice()).collect();
        let logits = model.forward(&token_ids, &neighbors);
        assert_eq!(logits.len(), 3);
    }

    #[test]
    fn learns_simple_task() {
        // Task: if the first and second tokens are the same, predict class 0,
        // otherwise class 1.
        let mut model = GraphAttentionClassifier::new(4, 8, 16, 2, 1, 7, None, false);
        let lr = 0.5;

        for _ in 0..2000 {
            let examples: Vec<(Vec<u32>, Vec<Vec<usize>>, usize)> = vec![
                (vec![0, 0, 1], vec![vec![0], vec![1], vec![2]], 0),
                (vec![0, 1, 2], vec![vec![0], vec![1], vec![2]], 1),
                (vec![1, 1, 0], vec![vec![0], vec![1], vec![2]], 0),
                (vec![1, 2, 0], vec![vec![0], vec![1], vec![2]], 1),
            ];
            for (ids, nbrs, target) in examples {
                let (grad, _) = model.backward_batch(&ids, &[nbrs], &[target], 1);
                model.apply_sgd(&grad, lr);
            }
        }

        let neighbors_owned: Vec<Vec<usize>> = vec![vec![0], vec![1], vec![2]];
        let neighbors: Vec<&[usize]> = neighbors_owned.iter().map(|v| v.as_slice()).collect();
        assert_eq!(model.predict(&[0, 0, 1], &neighbors), 0);
        assert_eq!(model.predict(&[0, 1, 2], &neighbors), 1);
    }

    #[test]
    fn learns_with_rope() {
        // Same task but with RoPE enabled; even embed_dim is required.
        let mut model = GraphAttentionClassifier::new(4, 8, 16, 2, 1, 7, Some(10000.0), false);
        let lr = 0.5;

        for _ in 0..3000 {
            let examples: Vec<(Vec<u32>, Vec<Vec<usize>>, usize)> = vec![
                (vec![0, 0, 1], vec![vec![0], vec![1], vec![2]], 0),
                (vec![0, 1, 2], vec![vec![0], vec![1], vec![2]], 1),
                (vec![1, 1, 0], vec![vec![0], vec![1], vec![2]], 0),
                (vec![1, 2, 0], vec![vec![0], vec![1], vec![2]], 1),
            ];
            for (ids, nbrs, target) in examples {
                let (grad, _) = model.backward_batch(&ids, &[nbrs], &[target], 1);
                model.apply_sgd(&grad, lr);
            }
        }

        let neighbors_owned: Vec<Vec<usize>> = vec![vec![0], vec![1], vec![2]];
        let neighbors: Vec<&[usize]> = neighbors_owned.iter().map(|v| v.as_slice()).collect();
        assert_eq!(model.predict(&[0, 0, 1], &neighbors), 0);
        assert_eq!(model.predict(&[0, 1, 2], &neighbors), 1);
    }

    #[test]
    fn learns_simple_task_geometric_only() {
        // Geometric-only attention removes full self-attention but the task
        // is small enough that attending to self + one neighbour per position
        // still carries enough signal.
        let mut model = GraphAttentionClassifier::new(4, 8, 16, 2, 1, 7, None, false);
        model.set_geometric_attention_only(true);
        let lr = 0.5;

        for _ in 0..2000 {
            let examples: Vec<(Vec<u32>, Vec<Vec<usize>>, usize)> = vec![
                (vec![0, 0, 1], vec![vec![0], vec![1], vec![2]], 0),
                (vec![0, 1, 2], vec![vec![0], vec![1], vec![2]], 1),
                (vec![1, 1, 0], vec![vec![0], vec![1], vec![2]], 0),
                (vec![1, 2, 0], vec![vec![0], vec![1], vec![2]], 1),
            ];
            for (ids, nbrs, target) in examples {
                let (grad, _) = model.backward_batch(&ids, &[nbrs], &[target], 1);
                model.apply_sgd(&grad, lr);
            }
        }

        let neighbors_owned: Vec<Vec<usize>> = vec![vec![0], vec![1], vec![2]];
        let neighbors: Vec<&[usize]> = neighbors_owned.iter().map(|v| v.as_slice()).collect();
        assert_eq!(model.predict(&[0, 0, 1], &neighbors), 0);
        assert_eq!(model.predict(&[0, 1, 2], &neighbors), 1);
    }
}