osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
1059 lines
35 KiB
Rust
1059 lines
35 KiB
Rust
//! V-JEPA: Joint Embedding Predictive Architecture for Video
|
||
//!
|
||
//! Extends I-JEPA (jepa.rs) to video:
|
||
//! - 3D patch embeddings (temporal × spatial)
|
||
//! - Tube masking: random space-time tubes covering 90% of volume
|
||
//! - Temporal context masking strategy
|
||
//! - Neuro-JEPA extension: apply same paradigm to EEG/MEG signal segments
|
||
|
||
use super::jepa::{
|
||
BlockMaskConfig, EmaTargetEncoder, FeatureBank, JepaEvaluator, JepaLossResult, JepaPredictor,
|
||
jepa_loss,
|
||
};
|
||
|
||
// ============================================================================
|
||
// Video Patch Configuration
|
||
// ============================================================================
|
||
|
||
/// Configuration for 3D (video) patch embeddings
|
||
#[derive(Debug, Clone)]
|
||
pub struct VideoPatchConfig {
|
||
/// Spatial patch size in pixels (e.g. 16)
|
||
pub spatial_patch_size: usize,
|
||
/// Temporal patch size in frames (e.g. 2)
|
||
pub temporal_patch_size: usize,
|
||
/// Spatial resolution (height = width assumed square)
|
||
pub image_size: usize,
|
||
/// Number of input frames
|
||
pub num_frames: usize,
|
||
/// Embedding dimension
|
||
pub embed_dim: usize,
|
||
}
|
||
|
||
impl Default for VideoPatchConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
spatial_patch_size: 16,
|
||
temporal_patch_size: 2,
|
||
image_size: 224,
|
||
num_frames: 16,
|
||
embed_dim: 1024,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl VideoPatchConfig {
|
||
/// Spatial grid size (patches per dimension)
|
||
pub fn spatial_grid(&self) -> usize {
|
||
self.image_size / self.spatial_patch_size
|
||
}
|
||
|
||
/// Number of temporal segments
|
||
pub fn temporal_segments(&self) -> usize {
|
||
self.num_frames / self.temporal_patch_size
|
||
}
|
||
|
||
/// Total 3D patches: T × H × W
|
||
pub fn total_patches(&self) -> usize {
|
||
self.temporal_segments() * self.spatial_grid() * self.spatial_grid()
|
||
}
|
||
|
||
/// Convert flat patch index to (t, h, w) tuple
|
||
pub fn patch_idx_to_thw(&self, idx: usize) -> (usize, usize, usize) {
|
||
let sg = self.spatial_grid();
|
||
let spatial = sg * sg;
|
||
let t = idx / spatial;
|
||
let h = (idx % spatial) / sg;
|
||
let w = idx % sg;
|
||
(t, h, w)
|
||
}
|
||
|
||
/// Convert (t, h, w) to flat patch index
|
||
pub fn thw_to_patch_idx(&self, t: usize, h: usize, w: usize) -> usize {
|
||
let sg = self.spatial_grid();
|
||
t * sg * sg + h * sg + w
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// 3D Patch Embedder (CPU reference)
|
||
// ============================================================================
|
||
|
||
/// 3D patch embedder: video tensor → sequence of patch embeddings
|
||
/// Input: [T, H, W, C] video frames (frame-major)
|
||
/// Output: [T/Tp, H/Sp, W/Sp, embed_dim] patch sequence
|
||
#[derive(Debug, Clone)]
|
||
pub struct PatchEmbed3D {
|
||
pub config: VideoPatchConfig,
|
||
/// Projection weights [C * Tp * Sp * Sp, embed_dim]
|
||
proj_w: Vec<f32>,
|
||
proj_b: Vec<f32>,
|
||
/// Position embeddings [total_patches, embed_dim]
|
||
pos_embed: Vec<f32>,
|
||
}
|
||
|
||
impl PatchEmbed3D {
|
||
pub fn new(config: VideoPatchConfig, channels: usize) -> Self {
|
||
let tp = config.temporal_patch_size;
|
||
let sp = config.spatial_patch_size;
|
||
let d = config.embed_dim;
|
||
let patch_vol = channels * tp * sp * sp;
|
||
let n = config.total_patches();
|
||
|
||
let init = |size: usize, s: u64| -> Vec<f32> {
|
||
let mut v = vec![0.0f32; size];
|
||
let mut lcg = s.wrapping_mul(2654435761).wrapping_add(1);
|
||
let scale = (2.0f32 / size as f32).sqrt();
|
||
for x in v.iter_mut() {
|
||
lcg = lcg
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
let u = ((lcg >> 11) as f32) / (1u64 << 53) as f32;
|
||
*x = (u * 2.0 - 1.0) * scale;
|
||
}
|
||
v
|
||
};
|
||
|
||
Self {
|
||
config,
|
||
proj_w: init(patch_vol * d, 77),
|
||
proj_b: vec![0.0; d],
|
||
pos_embed: init(n * d, 78),
|
||
}
|
||
}
|
||
|
||
/// Embed a video frame sequence.
|
||
///
|
||
/// `video`: flat [num_frames * H * W * C] in frame-major order
|
||
/// Returns: [total_patches, embed_dim]
|
||
pub fn embed(&self, video: &[f32], channels: usize) -> Vec<f32> {
|
||
let cfg = &self.config;
|
||
let ts = cfg.temporal_segments();
|
||
let sg = cfg.spatial_grid();
|
||
let tp = cfg.temporal_patch_size;
|
||
let sp = cfg.spatial_patch_size;
|
||
let h = cfg.image_size;
|
||
let w = cfg.image_size;
|
||
let d = cfg.embed_dim;
|
||
let patch_vol = channels * tp * sp * sp;
|
||
let n = cfg.total_patches();
|
||
let mut out = vec![0.0f32; n * d];
|
||
|
||
for ti in 0..ts {
|
||
for hi in 0..sg {
|
||
for wi in 0..sg {
|
||
let patch_idx = ti * sg * sg + hi * sg + wi;
|
||
// Extract patch volume: tp frames × sp × sp × C
|
||
let mut patch = vec![0.0f32; patch_vol];
|
||
let mut pi = 0;
|
||
for tf in 0..tp {
|
||
let frame = ti * tp + tf;
|
||
for ph in 0..sp {
|
||
let row = hi * sp + ph;
|
||
for pw in 0..sp {
|
||
let col = wi * sp + pw;
|
||
for c in 0..channels {
|
||
let flat_idx = frame * h * w * channels
|
||
+ row * w * channels
|
||
+ col * channels
|
||
+ c;
|
||
if flat_idx < video.len() {
|
||
patch[pi] = video[flat_idx];
|
||
}
|
||
pi += 1;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// Linear projection
|
||
for j in 0..d {
|
||
let mut val = self.proj_b[j];
|
||
for k in 0..patch_vol {
|
||
val += patch[k] * self.proj_w[k * d + j];
|
||
}
|
||
// Add position embedding
|
||
val += self.pos_embed[patch_idx * d + j];
|
||
out[patch_idx * d + j] = val;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Tube Masking (V-JEPA)
|
||
// ============================================================================
|
||
|
||
/// V-JEPA tube masking: mask spatially-consistent tubes across time.
|
||
///
|
||
/// For each selected spatial position, ALL temporal frames at that
|
||
/// spatial location are masked. This prevents the model from inferring
|
||
/// masked patches from adjacent frames at the same spatial location.
|
||
#[derive(Debug, Clone)]
|
||
pub struct TubeMaskConfig {
|
||
/// Fraction of spatial positions to mask (V-JEPA uses ~0.90)
|
||
pub mask_ratio: f64,
|
||
/// Number of target tubes to sample (analogous to target blocks)
|
||
pub num_target_tubes: usize,
|
||
/// Spatial scale range for each tube block
|
||
pub spatial_scale_range: (f64, f64),
|
||
}
|
||
|
||
impl Default for TubeMaskConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
mask_ratio: 0.90,
|
||
num_target_tubes: 8,
|
||
spatial_scale_range: (0.15, 0.20),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Result of tube masking for a video
|
||
#[derive(Debug, Clone)]
|
||
pub struct TubeMaskResult {
|
||
/// Flat patch indices visible to context encoder [t, h, w] → flat
|
||
pub context_indices: Vec<usize>,
|
||
/// Flat patch indices for each target tube
|
||
pub target_tubes: Vec<Vec<usize>>,
|
||
/// Union of all target indices
|
||
pub all_target_indices: Vec<usize>,
|
||
/// Config reference
|
||
pub video_config: VideoPatchConfig,
|
||
}
|
||
|
||
impl TubeMaskResult {
|
||
pub fn num_context(&self) -> usize {
|
||
self.context_indices.len()
|
||
}
|
||
pub fn num_target_patches(&self) -> usize {
|
||
self.all_target_indices.len()
|
||
}
|
||
pub fn num_target_tubes(&self) -> usize {
|
||
self.target_tubes.len()
|
||
}
|
||
}
|
||
|
||
/// Video tube masking strategy
|
||
#[derive(Debug, Clone)]
|
||
pub struct TubeMaskStrategy {
|
||
config: TubeMaskConfig,
|
||
lcg: u64,
|
||
}
|
||
|
||
impl TubeMaskStrategy {
|
||
pub fn new(config: TubeMaskConfig) -> Self {
|
||
Self {
|
||
config,
|
||
lcg: 9876543210987654321,
|
||
}
|
||
}
|
||
|
||
fn lcg_next(&mut self) -> u64 {
|
||
self.lcg = self
|
||
.lcg
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
self.lcg
|
||
}
|
||
fn rand_f64(&mut self) -> f64 {
|
||
(self.lcg_next() >> 11) as f64 / (1u64 << 53) as f64
|
||
}
|
||
fn rand_range(&mut self, lo: f64, hi: f64) -> f64 {
|
||
lo + self.rand_f64() * (hi - lo)
|
||
}
|
||
fn rand_usize(&mut self, n: usize) -> usize {
|
||
(self.lcg_next() % n as u64) as usize
|
||
}
|
||
|
||
/// Sample a spatial block (set of spatial positions) and extend across all T
|
||
fn sample_tube(&mut self, video_cfg: &VideoPatchConfig) -> Vec<usize> {
|
||
let sg = video_cfg.spatial_grid();
|
||
let ts = video_cfg.temporal_segments();
|
||
let total_spatial = sg * sg;
|
||
|
||
let scale = self.rand_range(
|
||
self.config.spatial_scale_range.0,
|
||
self.config.spatial_scale_range.1,
|
||
);
|
||
let area = (scale * total_spatial as f64).ceil() as usize;
|
||
let block_h = (area as f64).sqrt().ceil() as usize;
|
||
let block_h = block_h.clamp(1, sg);
|
||
let block_w = (area as f64 / block_h as f64).ceil() as usize;
|
||
let block_w = block_w.clamp(1, sg);
|
||
|
||
let top = self.rand_usize((sg - block_h + 1).max(1));
|
||
let left = self.rand_usize((sg - block_w + 1).max(1));
|
||
|
||
let mut indices = Vec::with_capacity(ts * block_h * block_w);
|
||
for t in 0..ts {
|
||
for h in top..top + block_h {
|
||
for w in left..left + block_w {
|
||
indices.push(video_cfg.thw_to_patch_idx(t, h.min(sg - 1), w.min(sg - 1)));
|
||
}
|
||
}
|
||
}
|
||
indices.sort_unstable();
|
||
indices.dedup();
|
||
indices
|
||
}
|
||
|
||
/// Generate a tube mask for one video
|
||
pub fn generate(&mut self, video_cfg: &VideoPatchConfig, seed_offset: u64) -> TubeMaskResult {
|
||
self.lcg = self.lcg.wrapping_add(seed_offset.wrapping_mul(1234567891));
|
||
let total_patches = video_cfg.total_patches();
|
||
let cfg = &self.config.clone();
|
||
|
||
let mut target_tubes: Vec<Vec<usize>> = Vec::with_capacity(cfg.num_target_tubes);
|
||
let mut all_target_set: std::collections::BTreeSet<usize> = Default::default();
|
||
|
||
for _ in 0..cfg.num_target_tubes {
|
||
let tube = self.sample_tube(video_cfg);
|
||
for &idx in &tube {
|
||
all_target_set.insert(idx);
|
||
}
|
||
target_tubes.push(tube);
|
||
}
|
||
let all_target_indices: Vec<usize> = all_target_set.into_iter().collect();
|
||
|
||
// Context = non-target, subsampled to (1 - mask_ratio) fraction
|
||
let mut context_cands: Vec<usize> = (0..total_patches)
|
||
.filter(|p| !all_target_indices.contains(p))
|
||
.collect();
|
||
let keep = ((1.0 - cfg.mask_ratio) * context_cands.len() as f64).ceil() as usize;
|
||
let keep = keep.max(1).min(context_cands.len());
|
||
for i in 0..keep {
|
||
let j = i + self.rand_usize(context_cands.len() - i);
|
||
context_cands.swap(i, j);
|
||
}
|
||
let mut context_indices = context_cands[..keep].to_vec();
|
||
context_indices.sort_unstable();
|
||
|
||
TubeMaskResult {
|
||
context_indices,
|
||
target_tubes,
|
||
all_target_indices,
|
||
video_config: video_cfg.clone(),
|
||
}
|
||
}
|
||
|
||
/// Batch generation
|
||
pub fn generate_batch(
|
||
&mut self,
|
||
batch_size: usize,
|
||
video_cfg: &VideoPatchConfig,
|
||
) -> Vec<TubeMaskResult> {
|
||
(0..batch_size)
|
||
.map(|i| self.generate(video_cfg, i as u64))
|
||
.collect()
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// V-JEPA Configuration and Trainer
|
||
// ============================================================================
|
||
|
||
/// Full V-JEPA configuration
|
||
#[derive(Debug, Clone)]
|
||
pub struct VJepaConfig {
|
||
/// Video patch configuration
|
||
pub video: VideoPatchConfig,
|
||
/// Predictor depth (narrow transformer, same structure as I-JEPA)
|
||
pub predictor_depth: usize,
|
||
/// EMA tau schedule
|
||
pub ema_tau_start: f64,
|
||
pub ema_tau_end: f64,
|
||
/// Masking strategy
|
||
pub mask: TubeMaskConfig,
|
||
}
|
||
|
||
impl Default for VJepaConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
video: VideoPatchConfig::default(),
|
||
predictor_depth: 6,
|
||
ema_tau_start: 0.996,
|
||
ema_tau_end: 1.0,
|
||
mask: TubeMaskConfig::default(),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl VJepaConfig {
|
||
/// Small config for testing
|
||
pub fn small_test() -> Self {
|
||
Self {
|
||
video: VideoPatchConfig {
|
||
spatial_patch_size: 8,
|
||
temporal_patch_size: 2,
|
||
image_size: 32,
|
||
num_frames: 4,
|
||
embed_dim: 64,
|
||
},
|
||
predictor_depth: 2,
|
||
ema_tau_start: 0.996,
|
||
ema_tau_end: 1.0,
|
||
mask: TubeMaskConfig {
|
||
mask_ratio: 0.75,
|
||
num_target_tubes: 2,
|
||
spatial_scale_range: (0.15, 0.30),
|
||
},
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Minimal mock video encoder (CPU reference for testing)
|
||
#[derive(Debug, Clone)]
|
||
pub struct MockVideoEncoder {
|
||
pub embed_dim: usize,
|
||
pub total_patches: usize,
|
||
weights: Vec<f32>,
|
||
}
|
||
|
||
impl MockVideoEncoder {
|
||
pub fn new(embed_dim: usize, total_patches: usize) -> Self {
|
||
let mut lcg = 54321u64;
|
||
let weights: Vec<f32> = (0..total_patches * embed_dim)
|
||
.map(|_| {
|
||
lcg = lcg
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
((lcg >> 11) as f32 / (1u64 << 53) as f32) * 2.0 - 1.0
|
||
})
|
||
.collect();
|
||
Self {
|
||
embed_dim,
|
||
total_patches,
|
||
weights,
|
||
}
|
||
}
|
||
|
||
pub fn encode(&self, patch_indices: &[usize]) -> Vec<f32> {
|
||
let n = patch_indices.len();
|
||
let d = self.embed_dim;
|
||
let mut out = vec![0.0f32; n * d];
|
||
for (i, &pi) in patch_indices.iter().enumerate() {
|
||
let pi = pi.min(self.total_patches - 1);
|
||
for dd in 0..d {
|
||
out[i * d + dd] = self.weights[pi * d + dd];
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
pub fn l2_normalize(reps: &mut Vec<f32>, n: usize, d: usize) {
|
||
for i in 0..n {
|
||
let norm = (0..d)
|
||
.map(|dd| reps[i * d + dd].powi(2))
|
||
.sum::<f32>()
|
||
.sqrt()
|
||
.max(1e-8);
|
||
for dd in 0..d {
|
||
reps[i * d + dd] /= norm;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// V-JEPA training step metrics
|
||
#[derive(Debug, Clone)]
|
||
pub struct VJepaStepMetrics {
|
||
pub loss: f32,
|
||
pub tube_losses: Vec<f32>,
|
||
pub ema_tau: f64,
|
||
pub num_context_patches: usize,
|
||
pub num_target_patches: usize,
|
||
pub mask_ratio_actual: f64,
|
||
}
|
||
|
||
/// V-JEPA trainer (video analog of JepaTrainer)
|
||
#[derive(Debug, Clone)]
|
||
pub struct VJepaTrainer {
|
||
pub config: VJepaConfig,
|
||
pub context_encoder: MockVideoEncoder,
|
||
pub target_ema: Vec<f32>, // shadow weights of context_encoder
|
||
pub predictor: JepaPredictor,
|
||
pub mask_strategy: TubeMaskStrategy,
|
||
pub step: usize,
|
||
pub total_steps: usize,
|
||
}
|
||
|
||
impl VJepaTrainer {
|
||
pub fn new(config: VJepaConfig, total_steps: usize) -> Self {
|
||
let total_patches = config.video.total_patches();
|
||
let embed_dim = config.video.embed_dim;
|
||
let pred_dim = (embed_dim / 4).max(64);
|
||
let num_heads = (pred_dim / 32).max(1).min(8);
|
||
|
||
let context_encoder = MockVideoEncoder::new(embed_dim, total_patches);
|
||
let target_ema = context_encoder.weights.clone();
|
||
let predictor = JepaPredictor::new(
|
||
embed_dim,
|
||
pred_dim,
|
||
config.predictor_depth,
|
||
num_heads,
|
||
total_patches,
|
||
);
|
||
let mask_strategy = TubeMaskStrategy::new(config.mask.clone());
|
||
|
||
Self {
|
||
config,
|
||
context_encoder,
|
||
target_ema,
|
||
predictor,
|
||
mask_strategy,
|
||
step: 0,
|
||
total_steps,
|
||
}
|
||
}
|
||
|
||
fn tau(&self) -> f64 {
|
||
let p = self.step as f64 / self.total_steps.max(1) as f64;
|
||
self.config.ema_tau_start + (self.config.ema_tau_end - self.config.ema_tau_start) * p
|
||
}
|
||
|
||
fn target_encode(&self, patch_indices: &[usize]) -> Vec<f32> {
|
||
let n = patch_indices.len();
|
||
let d = self.config.video.embed_dim;
|
||
let np = self.config.video.total_patches();
|
||
let mut out = vec![0.0f32; n * d];
|
||
for (i, &pi) in patch_indices.iter().enumerate() {
|
||
let pi = pi.min(np - 1);
|
||
for dd in 0..d {
|
||
out[i * d + dd] = self.target_ema[pi * d + dd];
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
fn ema_update(&mut self, tau: f64) {
|
||
for (sw, &ow) in self
|
||
.target_ema
|
||
.iter_mut()
|
||
.zip(self.context_encoder.weights.iter())
|
||
{
|
||
*sw = (tau as f32) * *sw + (1.0 - tau as f32) * ow;
|
||
}
|
||
}
|
||
|
||
/// One V-JEPA training step
|
||
pub fn train_step(&mut self, batch_size: usize) -> VJepaStepMetrics {
|
||
let video_cfg = self.config.video.clone();
|
||
let embed_dim = video_cfg.embed_dim;
|
||
let masks = self.mask_strategy.generate_batch(batch_size, &video_cfg);
|
||
let tau = self.tau();
|
||
|
||
let mut total_loss = 0.0f32;
|
||
let num_tubes = self.config.mask.num_target_tubes;
|
||
let mut tube_loss_acc = vec![0.0f32; num_tubes];
|
||
let mut total_ctx = 0usize;
|
||
let mut total_tgt = 0usize;
|
||
|
||
for (bi, mask) in masks.iter().enumerate() {
|
||
let _ = bi;
|
||
let ctx_reps = self.context_encoder.encode(&mask.context_indices);
|
||
let all_target = &mask.all_target_indices;
|
||
let predicted = self
|
||
.predictor
|
||
.forward(&ctx_reps, &mask.context_indices, all_target);
|
||
let mut target_reps = self.target_encode(all_target);
|
||
let n_tgt = all_target.len();
|
||
MockVideoEncoder::l2_normalize(&mut target_reps, n_tgt, embed_dim);
|
||
let mut pred_n = predicted.clone();
|
||
MockVideoEncoder::l2_normalize(&mut pred_n, n_tgt, embed_dim);
|
||
|
||
let mut offsets = vec![0usize; mask.target_tubes.len()];
|
||
let mut cursor = 0usize;
|
||
for (k, tube) in mask.target_tubes.iter().enumerate() {
|
||
offsets[k] = cursor.min(n_tgt);
|
||
cursor += tube.len();
|
||
}
|
||
|
||
let loss_result = jepa_loss(
|
||
&pred_n,
|
||
&target_reps,
|
||
embed_dim,
|
||
&mask.target_tubes,
|
||
&offsets,
|
||
);
|
||
total_loss += loss_result.loss;
|
||
for (k, &bl) in loss_result.block_losses.iter().enumerate() {
|
||
if k < tube_loss_acc.len() {
|
||
tube_loss_acc[k] += bl;
|
||
}
|
||
}
|
||
total_ctx += mask.num_context();
|
||
total_tgt += mask.num_target_patches();
|
||
}
|
||
|
||
self.ema_update(tau);
|
||
self.step += 1;
|
||
|
||
let bs = batch_size.max(1) as f32;
|
||
let avg_ctx = total_ctx / batch_size.max(1);
|
||
let total_p = video_cfg.total_patches();
|
||
VJepaStepMetrics {
|
||
loss: total_loss / bs,
|
||
tube_losses: tube_loss_acc.iter().map(|&v| v / bs).collect(),
|
||
ema_tau: tau,
|
||
num_context_patches: avg_ctx,
|
||
num_target_patches: total_tgt / batch_size.max(1),
|
||
mask_ratio_actual: 1.0 - (avg_ctx as f64 / total_p as f64),
|
||
}
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Neuro-JEPA: Signal-domain JEPA for EEG/MEG
|
||
// ============================================================================
|
||
|
||
/// Configuration for applying JEPA to multi-channel neural time series
|
||
#[derive(Debug, Clone)]
|
||
pub struct NeuroJepaConfig {
|
||
/// Number of EEG/MEG channels
|
||
pub num_channels: usize,
|
||
/// Segment length in samples per patch
|
||
pub segment_len: usize,
|
||
/// Total number of time samples in a recording window
|
||
pub total_samples: usize,
|
||
/// Embedding dimension for each channel-segment patch
|
||
pub embed_dim: usize,
|
||
/// Fraction of segments to mask as target
|
||
pub mask_ratio: f64,
|
||
/// EMA tau
|
||
pub ema_tau: f64,
|
||
}
|
||
|
||
impl Default for NeuroJepaConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
num_channels: 64, // typical EEG montage
|
||
segment_len: 64, // ~250ms at 256Hz
|
||
total_samples: 1024, // ~4s window at 256Hz
|
||
embed_dim: 256,
|
||
mask_ratio: 0.75,
|
||
ema_tau: 0.996,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl NeuroJepaConfig {
|
||
/// Number of time segments (patches in time dimension)
|
||
pub fn num_time_segments(&self) -> usize {
|
||
self.total_samples / self.segment_len
|
||
}
|
||
/// Total number of patches = channels × time_segments
|
||
pub fn total_patches(&self) -> usize {
|
||
self.num_channels * self.num_time_segments()
|
||
}
|
||
}
|
||
|
||
/// Mask result for neural signal JEPA
|
||
#[derive(Debug, Clone)]
|
||
pub struct NeuroMaskResult {
|
||
/// Patch indices (channel × time) visible to context encoder
|
||
pub context_indices: Vec<usize>,
|
||
/// Target patch indices
|
||
pub target_indices: Vec<usize>,
|
||
}
|
||
|
||
/// Neuro-JEPA masking: independent channel-time tube masking
|
||
///
|
||
/// Masks entire time-segments of selected channels (tube in channel-time space)
|
||
#[derive(Debug, Clone)]
|
||
pub struct NeuroMaskStrategy {
|
||
lcg: u64,
|
||
}
|
||
|
||
impl NeuroMaskStrategy {
|
||
pub fn new() -> Self {
|
||
Self { lcg: 1111111111 }
|
||
}
|
||
|
||
fn lcg_next(&mut self) -> u64 {
|
||
self.lcg = self
|
||
.lcg
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
self.lcg
|
||
}
|
||
fn rand_usize(&mut self, n: usize) -> usize {
|
||
(self.lcg_next() % n as u64) as usize
|
||
}
|
||
|
||
/// Generate channel-tube masks: select `mask_fraction * num_channels` channels,
|
||
/// mask ALL time segments for those channels (tube masking in channel dimension).
|
||
pub fn generate(&mut self, cfg: &NeuroJepaConfig, seed_offset: u64) -> NeuroMaskResult {
|
||
self.lcg = self.lcg.wrapping_add(seed_offset.wrapping_mul(987654321));
|
||
let ts = cfg.num_time_segments();
|
||
let nc = cfg.num_channels;
|
||
let total = cfg.total_patches();
|
||
|
||
// Select channels to mask (without replacement)
|
||
let num_mask_channels = ((cfg.mask_ratio * nc as f64).ceil() as usize).min(nc);
|
||
let mut channel_order: Vec<usize> = (0..nc).collect();
|
||
for i in 0..num_mask_channels {
|
||
let j = i + self.rand_usize(nc - i);
|
||
channel_order.swap(i, j);
|
||
}
|
||
let masked_channels: std::collections::BTreeSet<usize> =
|
||
channel_order[..num_mask_channels].iter().cloned().collect();
|
||
|
||
let mut context_indices: Vec<usize> = Vec::new();
|
||
let mut target_indices: Vec<usize> = Vec::new();
|
||
for c in 0..nc {
|
||
for t in 0..ts {
|
||
let idx = c * ts + t;
|
||
if idx < total {
|
||
if masked_channels.contains(&c) {
|
||
target_indices.push(idx);
|
||
} else {
|
||
context_indices.push(idx);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
NeuroMaskResult {
|
||
context_indices,
|
||
target_indices,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Default for NeuroMaskStrategy {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Tests — Batch 23 (V-JEPA + Neuro-JEPA)
|
||
// ============================================================================
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
// ── Video patch config ────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_video_patch_config_totals() {
|
||
let cfg = VideoPatchConfig {
|
||
spatial_patch_size: 16,
|
||
temporal_patch_size: 2,
|
||
image_size: 64,
|
||
num_frames: 8,
|
||
embed_dim: 128,
|
||
};
|
||
assert_eq!(cfg.spatial_grid(), 4);
|
||
assert_eq!(cfg.temporal_segments(), 4);
|
||
assert_eq!(cfg.total_patches(), 64); // 4 × 4 × 4
|
||
}
|
||
|
||
#[test]
|
||
fn test_thw_roundtrip() {
|
||
let cfg = VideoPatchConfig::default();
|
||
let idx = cfg.thw_to_patch_idx(2, 3, 5);
|
||
let (t, h, w) = cfg.patch_idx_to_thw(idx);
|
||
assert_eq!((t, h, w), (2, 3, 5));
|
||
}
|
||
|
||
#[test]
|
||
fn test_patch_embed_3d_output_shape() {
|
||
let cfg = VideoPatchConfig {
|
||
spatial_patch_size: 4,
|
||
temporal_patch_size: 2,
|
||
image_size: 8,
|
||
num_frames: 4,
|
||
embed_dim: 32,
|
||
};
|
||
let total = cfg.total_patches(); // (4/2)×(8/4)×(8/4) = 2×2×2 = 8
|
||
let embed = PatchEmbed3D::new(cfg.clone(), 3);
|
||
let video = vec![0.0f32; 4 * 8 * 8 * 3]; // [T, H, W, C]
|
||
let out = embed.embed(&video, 3);
|
||
assert_eq!(out.len(), total * 32);
|
||
}
|
||
|
||
#[test]
|
||
fn test_patch_embed_nonzero() {
|
||
let cfg = VideoPatchConfig {
|
||
spatial_patch_size: 4,
|
||
temporal_patch_size: 1,
|
||
image_size: 8,
|
||
num_frames: 2,
|
||
embed_dim: 16,
|
||
};
|
||
let embed = PatchEmbed3D::new(cfg.clone(), 1);
|
||
let video: Vec<f32> = (0..2 * 8 * 8).map(|i| i as f32 * 0.01).collect();
|
||
let out = embed.embed(&video, 1);
|
||
let any_nonzero = out.iter().any(|&v| v.abs() > 1e-8);
|
||
assert!(
|
||
any_nonzero,
|
||
"non-zero input should produce non-zero embeddings"
|
||
);
|
||
}
|
||
|
||
// ── Tube masking ──────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_tube_mask_basic() {
|
||
let video_cfg = VJepaConfig::small_test().video;
|
||
let mask_cfg = TubeMaskConfig {
|
||
mask_ratio: 0.75,
|
||
num_target_tubes: 2,
|
||
..Default::default()
|
||
};
|
||
let mut strategy = TubeMaskStrategy::new(mask_cfg);
|
||
let result = strategy.generate(&video_cfg, 0);
|
||
assert!(!result.context_indices.is_empty());
|
||
assert!(!result.all_target_indices.is_empty());
|
||
assert_eq!(result.target_tubes.len(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_tube_mask_no_overlap() {
|
||
let video_cfg = VJepaConfig::small_test().video;
|
||
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
|
||
let result = strategy.generate(&video_cfg, 1);
|
||
for &ci in &result.context_indices {
|
||
assert!(
|
||
!result.all_target_indices.contains(&ci),
|
||
"context patch {ci} appears in target"
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_tube_mask_indices_in_range() {
|
||
let video_cfg = VJepaConfig::small_test().video;
|
||
let total = video_cfg.total_patches();
|
||
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
|
||
let result = strategy.generate(&video_cfg, 0);
|
||
for &ci in &result.context_indices {
|
||
assert!(ci < total);
|
||
}
|
||
for &ti in &result.all_target_indices {
|
||
assert!(ti < total);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_tube_spans_all_time_segments() {
|
||
let cfg = VideoPatchConfig {
|
||
spatial_patch_size: 4,
|
||
temporal_patch_size: 1,
|
||
image_size: 8,
|
||
num_frames: 4,
|
||
embed_dim: 32,
|
||
};
|
||
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig {
|
||
mask_ratio: 0.9,
|
||
num_target_tubes: 1,
|
||
spatial_scale_range: (0.3, 0.4),
|
||
});
|
||
let result = strategy.generate(&cfg, 0);
|
||
// Target tubes should span multiple time steps (temporal consistency)
|
||
let ts = cfg.temporal_segments();
|
||
let sg = cfg.spatial_grid();
|
||
if !result.all_target_indices.is_empty() {
|
||
// All time segments should appear somewhere in target
|
||
let time_steps_covered: std::collections::BTreeSet<usize> = result
|
||
.all_target_indices
|
||
.iter()
|
||
.map(|&idx| idx / (sg * sg))
|
||
.collect();
|
||
assert!(
|
||
time_steps_covered.len() <= ts,
|
||
"temporal coverage must be within bounds"
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_tube_mask_batch() {
|
||
let video_cfg = VJepaConfig::small_test().video;
|
||
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
|
||
let batch = strategy.generate_batch(4, &video_cfg);
|
||
assert_eq!(batch.len(), 4);
|
||
for mask in &batch {
|
||
assert!(!mask.context_indices.is_empty());
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_tube_mask_vary_across_batch() {
|
||
// Use a larger grid so masks are unlikely to collide
|
||
let video_cfg = VideoPatchConfig {
|
||
spatial_patch_size: 4,
|
||
temporal_patch_size: 2,
|
||
image_size: 32,
|
||
num_frames: 8,
|
||
embed_dim: 64,
|
||
};
|
||
let mut strategy = TubeMaskStrategy::new(TubeMaskConfig::default());
|
||
let batch = strategy.generate_batch(8, &video_cfg);
|
||
// At least some pairs in the batch should differ
|
||
let any_differ = batch
|
||
.windows(2)
|
||
.any(|w| w[0].all_target_indices != w[1].all_target_indices);
|
||
assert!(
|
||
any_differ,
|
||
"different seeds should produce different masks in at least one pair"
|
||
);
|
||
}
|
||
|
||
// ── V-JEPA trainer ────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_vjepa_trainer_step_basic() {
|
||
let cfg = VJepaConfig::small_test();
|
||
let mut trainer = VJepaTrainer::new(cfg, 100);
|
||
let metrics = trainer.train_step(2);
|
||
assert!(metrics.loss.is_finite());
|
||
assert!(metrics.ema_tau >= 0.996 && metrics.ema_tau <= 1.0);
|
||
assert!(metrics.num_context_patches >= 1);
|
||
assert!(metrics.num_target_patches >= 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vjepa_mask_ratio_actual() {
|
||
let cfg = VJepaConfig::small_test();
|
||
let mut trainer = VJepaTrainer::new(cfg, 100);
|
||
let metrics = trainer.train_step(4);
|
||
assert!(metrics.mask_ratio_actual >= 0.0 && metrics.mask_ratio_actual <= 1.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vjepa_step_increments() {
|
||
let cfg = VJepaConfig::small_test();
|
||
let mut trainer = VJepaTrainer::new(cfg, 100);
|
||
trainer.train_step(2);
|
||
assert_eq!(trainer.step, 1);
|
||
trainer.train_step(2);
|
||
assert_eq!(trainer.step, 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vjepa_ema_anneals() {
|
||
let cfg = VJepaConfig::small_test();
|
||
let mut trainer = VJepaTrainer::new(cfg, 5);
|
||
let m0 = trainer.train_step(2);
|
||
for _ in 0..4 {
|
||
trainer.train_step(2);
|
||
}
|
||
let m4 = trainer.train_step(2);
|
||
assert!(m4.ema_tau >= m0.ema_tau, "tau must anneal upward");
|
||
}
|
||
|
||
#[test]
|
||
fn test_vjepa_tube_losses_count() {
|
||
let mut cfg = VJepaConfig::small_test();
|
||
cfg.mask.num_target_tubes = 3;
|
||
let mut trainer = VJepaTrainer::new(cfg, 10);
|
||
let metrics = trainer.train_step(2);
|
||
assert_eq!(metrics.tube_losses.len(), 3);
|
||
}
|
||
|
||
// ── Neuro-JEPA ───────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_neuro_config_patches() {
|
||
let cfg = NeuroJepaConfig::default();
|
||
assert_eq!(cfg.num_time_segments(), 16); // 1024 / 64
|
||
assert_eq!(cfg.total_patches(), 64 * 16); // 64 channels × 16 segments
|
||
}
|
||
|
||
#[test]
|
||
fn test_neuro_mask_basic() {
|
||
let cfg = NeuroJepaConfig::default();
|
||
let mut strategy = NeuroMaskStrategy::new();
|
||
let result = strategy.generate(&cfg, 0);
|
||
assert!(!result.context_indices.is_empty());
|
||
assert!(!result.target_indices.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_neuro_mask_no_overlap() {
|
||
let cfg = NeuroJepaConfig::default();
|
||
let mut strategy = NeuroMaskStrategy::new();
|
||
let result = strategy.generate(&cfg, 0);
|
||
for &ci in &result.context_indices {
|
||
assert!(!result.target_indices.contains(&ci));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_neuro_mask_covers_all_patches() {
|
||
let cfg = NeuroJepaConfig {
|
||
num_channels: 4,
|
||
segment_len: 16,
|
||
total_samples: 64,
|
||
..Default::default()
|
||
};
|
||
let mut strategy = NeuroMaskStrategy::new();
|
||
let result = strategy.generate(&cfg, 0);
|
||
let total = cfg.total_patches();
|
||
let covered = result.context_indices.len() + result.target_indices.len();
|
||
assert_eq!(
|
||
covered, total,
|
||
"all patches must appear in context or target"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_neuro_mask_ratio_respected() {
|
||
let cfg = NeuroJepaConfig {
|
||
num_channels: 8,
|
||
mask_ratio: 0.5,
|
||
..Default::default()
|
||
};
|
||
let mut strategy = NeuroMaskStrategy::new();
|
||
let result = strategy.generate(&cfg, 0);
|
||
let tgt_channels = result.target_indices.len() / cfg.num_time_segments();
|
||
let expected = ((0.5 * cfg.num_channels as f64).ceil() as usize).min(cfg.num_channels);
|
||
assert_eq!(tgt_channels, expected, "masked channels should match ratio");
|
||
}
|
||
|
||
#[test]
|
||
fn test_neuro_mask_tube_structure() {
|
||
// Verify that masked channels are fully masked across all time segments
|
||
let cfg = NeuroJepaConfig {
|
||
num_channels: 4,
|
||
segment_len: 4,
|
||
total_samples: 16,
|
||
mask_ratio: 0.5,
|
||
..Default::default()
|
||
};
|
||
let ts = cfg.num_time_segments();
|
||
let mut strategy = NeuroMaskStrategy::new();
|
||
let result = strategy.generate(&cfg, 0);
|
||
|
||
// Each target channel should have ALL its time segments in the target
|
||
let mut channel_times: std::collections::HashMap<usize, usize> = Default::default();
|
||
for &ti in &result.target_indices {
|
||
let channel = ti / ts;
|
||
*channel_times.entry(channel).or_insert(0) += 1;
|
||
}
|
||
for (&_ch, &count) in &channel_times {
|
||
assert_eq!(
|
||
count, ts,
|
||
"tube masking: a masked channel must have ALL time segments masked"
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_neuro_vary_across_seeds() {
|
||
let cfg = NeuroJepaConfig::default();
|
||
let mut strategy = NeuroMaskStrategy::new();
|
||
let r0 = strategy.generate(&cfg, 0);
|
||
let r1 = strategy.generate(&cfg, 999);
|
||
assert_ne!(
|
||
r0.target_indices, r1.target_indices,
|
||
"different seeds should produce different masks"
|
||
);
|
||
}
|
||
}
|