Struct ProjectionHead 

pub struct ProjectionHead { /* private fields */ }

Two-layer MLP projection head used in contrastive learning.

Architecture: in_dim → hidden_dim (ReLU) → out_dim

Implementations

impl ProjectionHead

pub fn new(in_dim: usize, hidden_dim: usize, out_dim: usize, seed: u64) -> Self

Construct a new projection head with Xavier-uniform initialised weights.

Arguments

• in_dim – input dimension (encoder output size)
• hidden_dim – hidden layer dimension
• out_dim – projection output dimension
• seed – RNG seed

pub fn forward(&self, x: &Array2<f64>) -> Array2<f64>

Forward pass: x → W1 x + b1 → ReLU → W2 x + b2

Arguments

• x – input [batch × in_dim]

Returns

Projected representations [batch × out_dim]

pub fn in_dim(&self) -> usize

Input dimension.

pub fn out_dim(&self) -> usize

Output projection dimension.