Struct GraphAttentionClassifier 

pub struct GraphAttentionClassifier {
    pub embedding: Vec<f32>,
    pub w_q: Vec<f32>,
    pub w_k: Vec<f32>,
    pub w_v: Vec<f32>,
    pub edge_weights_raw: Vec<f32>,
    pub mlp: MlpClassifier,
    /* private fields */
}

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.

Fields

embedding: Vec<f32>

Token embedding table, row-major [vocab_size, embed_dim].

w_q: Vec<f32>

Query projection, row-major [embed_dim, embed_dim].

w_k: Vec<f32>

Key projection, row-major [embed_dim, embed_dim].

w_v: Vec<f32>

Value projection, row-major [embed_dim, embed_dim].

edge_weights_raw: 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.

mlp: MlpClassifier

MLP head: embed_dim -> hidden_dim -> output_dim.

Implementations

impl GraphAttentionClassifier

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

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 set_geometric_attention_only(&mut self, enabled: bool)

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 forward(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> Vec<f32>

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_hybrid( &self, token_ids: &[u32], neighbors: &[&[usize]], ) -> Vec<f32>

Forward pass using the full hybrid attention graph.

pub fn forward_geometric_only( &self, token_ids: &[u32], neighbors: &[&[usize]], ) -> Vec<f32>

Forward pass using only geometric (neighbor) attention.

pub fn forward_hidden( &self, token_ids: &[u32], neighbors: &[&[usize]], ) -> (Vec<f32>, Vec<f32>)

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_only( &self, token_ids: &[u32], neighbors: &[&[usize]], ) -> Vec<f32>

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 loss( &self, token_ids: &[u32], neighbors: &[&[usize]], target: usize, ) -> f32

Cross-entropy loss for a single example.

pub fn predict(&self, token_ids: &[u32], neighbors: &[&[usize]]) -> usize

Predicted class for a single example.

pub fn attend_mixed( &self, token_ids: &[u32], neighbors: &[&[usize]], per_position_geometric: &[bool], ) -> Vec<Vec<f32>>

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 forward_mixed( &self, token_ids: &[u32], neighbors: &[&[usize]], per_position_geometric: &[bool], ) -> Vec<f32>

Forward pass with a per-position mixed attention mode.

pub fn forward_tiebreak( &self, token_ids: &[u32], neighbors: &[&[usize]], config: &TiebreakConfig, ) -> Vec<f32>

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 tiebreak_mask( &self, token_ids: &[u32], neighbors: &[&[usize]], config: &TiebreakConfig, ) -> Vec<bool>

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 forward_batch( &self, token_ids: &[u32], neighbors: &[Vec<Vec<usize>>], batch_size: usize, ) -> Vec<f32>

Forward pass for a batch.

pub fn backward_batch( &self, token_ids: &[u32], neighbors: &[Vec<Vec<usize>>], targets: &[usize], batch_size: usize, ) -> (GraphAttentionGradients, f32)

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 apply_sgd(&mut self, grad: &GraphAttentionGradients, lr: f32)

Apply a single SGD step using the supplied gradients.

pub fn flatten_params(&self) -> Vec<f32>

Flatten all learnable parameters into one vector, ordered [embedding, w_q, w_k, w_v, edge_weights_raw, mlp].

pub fn load_flat_params(&mut self, params: &[f32])

Restore parameters from a flat vector produced by Self::flatten_params.

Trait Implementations

impl Clone for GraphAttentionClassifier

fn clone(&self) -> GraphAttentionClassifier

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for GraphAttentionClassifier

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

Formats the value using the given formatter.