Struct MlpProbe 

pub struct MlpProbe {
    pub w1: Vec<f32>,
    pub b1: Vec<f32>,
    pub w2: Vec<f32>,
    pub b2: Vec<f32>,
    pub hidden_size: usize,
    pub mlp_hidden: usize,
    pub num_classes: usize,
}

MLP probe: frozen embeddings + trainable 2-layer MLP head (Level 0.5).

Adds a hidden layer with ReLU between embeddings and classification head. This allows non-linear decision boundaries, which can capture patterns that a linear probe cannot (e.g., shell safety from CodeBERT embeddings).

Architecture: embedding → Linear(hidden_size, mlp_hidden) → ReLU → Linear(mlp_hidden, num_classes)

Fields

w1: Vec<f32>

First layer weights [hidden_size × mlp_hidden] flattened row-major

b1: Vec<f32>

First layer bias [mlp_hidden]

w2: Vec<f32>

Second layer weights [mlp_hidden × num_classes] flattened row-major

b2: Vec<f32>

Second layer bias [num_classes]

hidden_size: usize

Input dimension

mlp_hidden: usize

Hidden layer dimension

num_classes: usize

Number of output classes

Implementations

impl MlpProbe

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

Create with Xavier initialization.

pub fn forward(&self, emb: &[f32]) -> (Vec<f32>, Vec<f32>)

Forward pass: embedding → hidden (ReLU) → logits.

pub fn predict(&self, emb: &[f32]) -> usize

Predict class index.

pub fn predict_probs(&self, emb: &[f32]) -> Vec<f32>

Predict class probabilities via softmax.

pub fn train( &mut self, embeddings: &[Vec<f32>], labels: &[usize], epochs: usize, learning_rate: f32, class_weights: Option<&[f32]>, weight_decay: f32, ) -> f32

Train with online SGD + class weights + L2 regularization.

pub fn num_parameters(&self) -> usize

Total trainable parameters.