Struct ClassificationHead 

pub struct ClassificationHead {
    pub weight: Tensor,
    pub bias: Tensor,
    /* private fields */
}

Classification head: mean pool + linear projection.

Maps transformer hidden states to class logits. Weight shape: [hidden_size * num_classes] (flattened row-major). Bias shape: [num_classes].

Fields

weight: Tensor

Linear weight [hidden_size, num_classes] flattened

bias: Tensor

Bias [num_classes]

Implementations

impl ClassificationHead

pub fn new(hidden_size: usize, num_classes: usize) -> Self

Create a new classification head with Xavier-initialized weights.

Arguments

• hidden_size - Transformer hidden dimension (e.g., 896 for Qwen2-0.5B)
• num_classes - Number of output classes (e.g., 5 for shell safety)

Contract (F-CLASS-004)

Validates hidden_size > 0 and num_classes >= 2.

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

Forward pass: hidden_states → mean pool → linear → logits.

Arguments

• hidden_states - Transformer output [seq_len * hidden_size] flattened
• seq_len - Sequence length

Returns

Logits tensor [num_classes]

Contract (F-CLASS-001)

Output always has exactly num_classes elements.

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

Mean pool hidden states across sequence length.

hidden_states: [seq_len * hidden_size] → [hidden_size]

pub fn cls_pool(&self, hidden_states: &Tensor) -> Tensor

CLS pooling: extract the first token’s hidden state.

In BERT-family models, position 0 is the [CLS] token which attends bidirectionally to all other tokens, making it suitable for classification.

Contract (ENC-007)

Output has exactly hidden_size elements (first position of input).

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

Last-token pooling: extract the last token’s hidden state.

In decoder models, the last token has seen all prior context through causal attention, making it suitable for classification.

pub fn pool( &self, hidden_states: &Tensor, seq_len: usize, strategy: PoolingStrategy, ) -> Tensor

Pool hidden states using the specified strategy.

pub fn forward_with_pooling( &self, hidden_states: &Tensor, seq_len: usize, strategy: PoolingStrategy, ) -> Tensor

Forward pass with configurable pooling strategy.

pub fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Get trainable parameters (weight + bias).

pub fn parameters(&self) -> Vec<&Tensor>

Get parameters (immutable).

pub fn num_classes(&self) -> usize

Number of classes.

pub fn hidden_size(&self) -> usize

Hidden size.

pub fn num_parameters(&self) -> usize

Total trainable parameter count.