v0.10.0: análise/prova/teste das críticas de engenharia Qwen3.8 lote 3 (r-35) — divisão SRP dos monólitos jepa/pensamento/moe/python em 26 submódulos (API estável via re-exports; 'Amdahl de compilação' refutada por experimento controlado, builds frios neutros), checkpoints atômicos SHA-256 puro FIPS 180-4 + fsync + rename + manifesto (CP-1..CP-3; correções da proposta: sem destruir estado anterior, manifesto com barreira, sem bincode/unwrap), TemperaturaInversa T-V3 (fronteira tipada da convexidade N-2; alegação compile-time corrigida), trait CamadaKhtst+CamadaHomeostatica (premissa do controlador refutada no fonte; Vec<Var>; adaptador MoeComPapel bit-a-bit), CLI Typer 75 (fachada aditiva com resume verificado por hash) — 123/123 cargo, 107/107 gates, treino real Wikipédia PT-BR com checkpoints íntegros
06e4b6f verified | //! `jepa::variantes` — extraído do monólito na v0.10.0 | |
| //! (divisão SRP, crítica Qwen3.8 lote 3 nº 1; r-35 §2). | |
| use super::*; | |
| // --------------------------------------------------------------------------- | |
| // Parte 4 — variantes e treino unificado | |
| // --------------------------------------------------------------------------- | |
| pub enum VarianteJepa { | |
| ImageJepa, | |
| VJepa2, | |
| McJepa, | |
| AudioJepa, | |
| PointJepa, | |
| Jepa3D, | |
| BrainJepa, | |
| TJepa, | |
| LeJepa, | |
| BJepa, | |
| } | |
| impl VarianteJepa { | |
| pub fn nome(&self) -> &'static str { | |
| match self { | |
| Self::ImageJepa => "Image-JEPA", | |
| Self::VJepa2 => "V-JEPA 2", | |
| Self::McJepa => "MC-JEPA", | |
| Self::AudioJepa => "Audio-JEPA", | |
| Self::PointJepa => "Point-JEPA", | |
| Self::Jepa3D => "3D-JEPA", | |
| Self::BrainJepa => "Brain-JEPA", | |
| Self::TJepa => "T-JEPA", | |
| Self::LeJepa => "LeJEPA", | |
| Self::BJepa => "BJEPA", | |
| } | |
| } | |
| } | |
| pub struct CfgJepa { | |
| pub d_emb: usize, | |
| pub epocas: usize, | |
| pub lr: f64, | |
| pub m: f64, // EMA | |
| pub lambda_vc: f64, | |
| pub gamma_vc: f64, | |
| pub frac_mascara: f64, | |
| pub semente: u64, | |
| /// tamanho do bloco-alvo (nº de patches por bloco). | |
| pub tam_bloco: usize, | |
| } | |
| impl Default for CfgJepa { | |
| fn default() -> Self { | |
| Self { | |
| d_emb: 32, | |
| epocas: 60, | |
| lr: 1.5e-3, | |
| m: 0.99, | |
| lambda_vc: 0.5, | |
| gamma_vc: 1.0, | |
| frac_mascara: 0.35, | |
| semente: 42, | |
| tam_bloco: 2, | |
| } | |
| } | |
| } | |
| /// Relatório de treino de uma variante. | |
| pub struct RelatorioJepa { | |
| pub variante: String, | |
| pub perda_inicial: f64, | |
| pub perda_final: f64, | |
| /// perda PREDITIVA isolada (sem o regularizador VC) — critério de | |
| /// aprendizado honesto (a VC pode crescer legitimamente ao empurrar o | |
| /// std das incorporações para cima). | |
| pub perda_pred_inicial: f64, | |
| pub perda_pred_final: f64, | |
| /// desvio-padrão médio por dimensão das incorporações (monitor de | |
| /// colapso: → 0 = colapso). | |
| pub std_emb_final: f64, | |
| pub std_emb_inicial: f64, | |
| pub n_patches: usize, | |
| pub d_patch: usize, | |
| pub epocas: usize, | |
| } | |
| /// Treina uma variante JEPA sobre patches já extraídos `(n, d_patch)` com | |
| /// um esquema de máscara fornecido (padrão: intervalos 1D). | |
| /// | |
| /// Máscara=true ⇒ patch é ALVO (bloco); false ⇒ CONTEXTO. | |
| pub fn treinar_variante( | |
| patches: &Tensor, | |
| mascara: &dyn Fn(usize, &mut RngDeterministico) -> Vec<bool>, | |
| cfg: &CfgJepa, | |
| variante: VarianteJepa, | |
| cond: Option<&Tensor>, | |
| dev: &Device, | |
| ) -> Result<RelatorioJepa> { | |
| let n = patches.dim(0)?; | |
| let d_patch = patches.dim(1)?; | |
| let mut rng = RngDeterministico::novo(cfg.semente); | |
| let mut nucleo = NucleoJepa::novo( | |
| d_patch, | |
| cfg.d_emb, | |
| if cond.is_some() { cond.unwrap().dim(1)? } else { 0 }, | |
| n, | |
| cfg.m, | |
| cfg.lambda_vc, | |
| cfg.gamma_vc, | |
| dev, | |
| )?; | |
| let otimo_vars = nucleo.vars(); | |
| let mut otimizador = crate::treino::AdamWManual::novo(otimo_vars, cfg.lr, 0.01); | |
| // métricas iniciais | |
| let z0 = nucleo.codificar(patches)?; | |
| let std0 = std_por_dim(&z0)?; | |
| let mut perda_inicial = f64::NAN; | |
| let mut perda_final = f64::NAN; | |
| let mut pred_inicial = f64::NAN; | |
| let mut pred_final = f64::NAN; | |
| let mut vc_final = f64::NAN; | |
| for epoca in 0..cfg.epocas { | |
| // máscara fresca a cada época (com proteção contra máscara degenerada: | |
| // se contexto ou alvo ficar vazio, repete com fração menor — a época | |
| // NUNCA é pulada, a perda sempre é computada) | |
| let (ctx_idx, blocos) = loop { | |
| let mut masc = mascara(n, &mut rng); | |
| // defensivo: máscara menor que n => o resto é contexto | |
| if masc.len() < n { | |
| masc.resize(n, false); | |
| } | |
| masc.truncate(n); | |
| let ctx: Vec<u32> = (0..n).filter(|&i| !masc[i]).map(|i| i as u32).collect(); | |
| let alvos: Vec<usize> = (0..n).filter(|&i| masc[i]).collect(); | |
| if !ctx.is_empty() && !alvos.is_empty() { | |
| break (ctx, alvos); | |
| } | |
| }; | |
| // agrupa alvos em blocos de tam_bloco | |
| let grupos: Vec<Vec<usize>> = blocos | |
| .chunks(cfg.tam_bloco.max(1)) | |
| .map(|c| c.to_vec()) | |
| .collect(); | |
| if ctx_idx.is_empty() || grupos.is_empty() { | |
| continue; | |
| } | |
| let dev2 = patches.device(); | |
| let n_ctx = ctx_idx.len(); | |
| let idx_t = Tensor::from_vec(ctx_idx, (n_ctx,), &dev2)?; | |
| let ctx_patches = patches.index_select(&idx_t, 0)?; | |
| let z_ctx_todos = nucleo.codificar(&ctx_patches)?; // (n_ctx, d_emb) | |
| let z_ctx = z_ctx_todos.mean_keepdim(0)?; // (1, d_emb) | |
| // alvos por grupo | |
| let mut alvos: Vec<Vec<(usize, Tensor)>> = Vec::new(); | |
| for g in &grupos { | |
| let idx: Vec<u32> = g.iter().map(|&i| i as u32).collect(); | |
| let n_g = idx.len(); | |
| let it = Tensor::from_vec(idx, (n_g,), &dev2)?; | |
| alvos.push(vec![(g[0], patches.index_select(&it, 0)?)]); | |
| } | |
| // condição por amostra — broadcast (o lote é 1 aqui: o conjunto todo) | |
| let (perda, lp, lv) = nucleo.perda_lote(&z_ctx, &[alvos_concat(&alvos)], cond, &z_ctx_todos)?; | |
| let p = perda.to_scalar::<f32>()? as f64; | |
| if p.is_finite() { | |
| if epoca == 0 { | |
| perda_inicial = p; | |
| pred_inicial = lp; | |
| } | |
| perda_final = p; | |
| pred_final = lp; | |
| vc_final = lv; | |
| } | |
| // decaimento cosseno do lr + clipping de gradiente (estabilidade) | |
| let frac = epoca as f64 / cfg.epocas.max(1) as f64; | |
| let lr_t = cfg.lr * 0.5 * (1.0 + (std::f64::consts::PI * frac).cos()); | |
| otimizador.set_lr(lr_t); | |
| let mut grads = crate::treino::grads_achatados(&perda, &nucleo.vars())?; | |
| crate::treino::AdamWManual::clipar(&mut grads, 1.0); | |
| otimizador.passo(grads); | |
| nucleo.ema_passo()?; | |
| } | |
| let z1 = nucleo.codificar(patches)?; | |
| let std1 = std_por_dim(&z1)?; | |
| let _ = (pred_inicial, vc_final); | |
| Ok(RelatorioJepa { | |
| variante: variante.nome().to_string(), | |
| perda_inicial, | |
| perda_final, | |
| perda_pred_inicial: pred_inicial, | |
| perda_pred_final: pred_final, | |
| std_emb_final: std1, | |
| std_emb_inicial: std0, | |
| n_patches: n, | |
| d_patch, | |
| epocas: cfg.epocas, | |
| }) | |
| } | |
| fn alvos_concat(alvos: &[Vec<(usize, Tensor)>]) -> Vec<(usize, Tensor)> { | |
| alvos.iter().flatten().cloned().collect() | |
| } | |
| fn std_por_dim(z: &Tensor) -> Result<f64> { | |
| let n = z.dim(0)? as f64; | |
| let media = z.mean_keepdim(0)?; | |
| let centr = z.broadcast_sub(&media)?; | |
| let var = (centr.sqr()?.sum(0)? / n)?; | |
| Ok(var.sqrt()?.mean_all()?.to_scalar::<f32>()? as f64) | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Variantes especiais: LeJEPA (edição-condicionado), MC-JEPA (cruz-modal), | |
| // BJEPA (bidirecional com consistência de ciclo) | |
| // --------------------------------------------------------------------------- | |
| /// **LeJEPA** — predição latente condicionada por EDIÇÃO: contexto = E(x), | |
| /// condição = instrução de edição `e` (vetor de baixa dimensão), alvo = | |
| /// `E_ξ(edit_e(x))`. A condição é a INSTRUÇÃO (não o resultado — sem | |
| /// vazamento: `e` não contém `edit(x)`). | |
| pub fn treinar_lejepa( | |
| dados: &Tensor, // (n, d_patch) — originais | |
| edits: &Tensor, // (n, d_cond) — instruções de edição | |
| editados: &Tensor, // (n, d_patch) — resultados das edições | |
| cfg: &CfgJepa, | |
| dev: &Device, | |
| ) -> Result<RelatorioJepa> { | |
| let n = dados.dim(0)?; | |
| let d_patch = dados.dim(1)?; | |
| let d_cond = edits.dim(1)?; | |
| let mut rng = RngDeterministico::novo(cfg.semente); | |
| let mut nucleo = NucleoJepa::novo( | |
| d_patch, cfg.d_emb, d_cond, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev, | |
| )?; | |
| let mut otim = crate::treino::AdamWManual::novo(nucleo.vars(), cfg.lr, 0.01); | |
| let mut perda_inicial = f64::NAN; | |
| let mut perda_final = f64::NAN; | |
| let z0 = nucleo.codificar(dados)?; | |
| let std0 = std_por_dim(&z0)?; | |
| for epoca in 0..cfg.epocas { | |
| // subamostra do lote | |
| let idx = rng.amostra_k_de_n(n.min(64), n); | |
| let sel: Vec<u32> = idx.iter().map(|&i| i as u32).collect(); | |
| let it = Tensor::from_vec(sel.clone(), (sel.len(),), dev)?; | |
| let xb = dados.index_select(&it, 0)?; | |
| let eb = edits.index_select(&it, 0)?; | |
| let yb = editados.index_select(&it, 0)?; | |
| let z_ctx = nucleo.codificar(&xb)?; // (m, d_emb) | |
| let mut perdas: Vec<Tensor> = Vec::new(); | |
| for i in 0..xb.dim(0)? { | |
| let zc = z_ctx.get(i)?.broadcast_as((1, cfg.d_emb))?; | |
| let e_i = eb.get(i)?.broadcast_as((1, d_cond))?; | |
| let yt = yb.get(i)?.broadcast_as((1, d_patch))?; | |
| let zt = nucleo | |
| .alvo_codificar(&yt)? | |
| .mean(0)? | |
| .broadcast_as((1, cfg.d_emb))?; | |
| let pred = nucleo.predizer(&zc, 0, Some(&e_i))?; | |
| let dif = (&pred - &zt)?; | |
| perdas.push((dif.sqr()?.sum_all()? / cfg.d_emb as f64)?); | |
| } | |
| let l_pred = (Tensor::stack(&perdas, 0)?.sum_all()? / perdas.len() as f64)?; | |
| let l_vc = nucleo.perda_vc(&z_ctx)?; | |
| let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(cfg.lambda_vc, 0.0)?)?; | |
| let p = total.to_scalar::<f32>()? as f64; | |
| if p.is_finite() { | |
| if epoca == 0 { | |
| perda_inicial = p; | |
| } | |
| perda_final = p; | |
| } | |
| let mut grads = crate::treino::grads_achatados(&total, &nucleo.vars())?; | |
| crate::treino::AdamWManual::clipar(&mut grads, 1.0); | |
| otim.passo(grads); | |
| nucleo.ema_passo()?; | |
| } | |
| let z1 = nucleo.codificar(dados)?; | |
| Ok(RelatorioJepa { | |
| variante: "LeJEPA".into(), | |
| perda_inicial, | |
| perda_final, | |
| perda_pred_inicial: perda_inicial, | |
| perda_pred_final: perda_final, | |
| std_emb_final: std_por_dim(&z1)?, | |
| std_emb_inicial: std0, | |
| n_patches: n, | |
| d_patch, | |
| epocas: cfg.epocas, | |
| }) | |
| } | |
| /// **MC-JEPA** — cruz-modal: dois codificadores, predição nos DOIS sentidos | |
| /// (A→B e B→A) com preditores separados e alvos EMA dos dois lados. | |
| pub fn treinar_mcjepa( | |
| mod_a: &Tensor, // (n, d_a) | |
| mod_b: &Tensor, // (n, d_b) | |
| cfg: &CfgJepa, | |
| dev: &Device, | |
| ) -> Result<RelatorioJepa> { | |
| let n = mod_a.dim(0)?; | |
| let d_a = mod_a.dim(1)?; | |
| let d_b = mod_b.dim(1)?; | |
| // núcleo A→B e núcleo B→A (compartilham apenas a estrutura) | |
| let mut nb = NucleoJepa::novo(d_a, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev)?; | |
| let mut na = NucleoJepa::novo(d_b, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev)?; | |
| let mut otim = crate::treino::AdamWManual::novo( | |
| [nb.vars(), na.vars()].concat(), | |
| cfg.lr, | |
| 0.01, | |
| ); | |
| let mut rng = RngDeterministico::novo(cfg.semente); | |
| let mut perda_inicial = f64::NAN; | |
| let mut perda_final = f64::NAN; | |
| for epoca in 0..cfg.epocas { | |
| let idx = rng.amostra_k_de_n(n.min(64), n); | |
| let sel: Vec<u32> = idx.iter().map(|&i| i as u32).collect(); | |
| let it = Tensor::from_vec(sel, (idx.len(),), dev)?; | |
| let xa = mod_a.index_select(&it, 0)?; | |
| let xb = mod_b.index_select(&it, 0)?; | |
| // z_a = E_a(x_a) prevê ẑ_b; z_b = E_b(x_b) prevê ẑ_a | |
| let z_a = nb.codificar(&xa)?; // (m, d_emb) | |
| let z_b = na.codificar(&xb)?; | |
| let t_a = na.alvo_codificar(&xa)?; // alvo EMA do lado A | |
| let t_b = nb.alvo_codificar(&xb)?; // alvo EMA do lado B | |
| let m_ = xa.dim(0)?; | |
| let mut perdas: Vec<Tensor> = Vec::new(); | |
| for i in 0..m_ { | |
| // A → B | |
| let zc = z_a.get(i)?.broadcast_as((1, cfg.d_emb))?; | |
| let zt = t_b.get(i)?.broadcast_as((1, cfg.d_emb))?; | |
| let pred = nb.predizer(&zc, 0, None)?; | |
| let dif = (&pred - &zt)?; | |
| perdas.push((dif.sqr()?.sum_all()? / cfg.d_emb as f64)?); | |
| // B → A | |
| let zc2 = z_b.get(i)?.broadcast_as((1, cfg.d_emb))?; | |
| let zt2 = t_a.get(i)?.broadcast_as((1, cfg.d_emb))?; | |
| let pred2 = na.predizer(&zc2, 0, None)?; | |
| let dif2 = (&pred2 - &zt2)?; | |
| perdas.push((dif2.sqr()?.sum_all()? / cfg.d_emb as f64)?); | |
| } | |
| let l_pred = (Tensor::stack(&perdas, 0)?.sum_all()? / perdas.len() as f64)?; | |
| let l_vc = (nb.perda_vc(&z_a)? + na.perda_vc(&z_b)?)?; | |
| let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(cfg.lambda_vc, 0.0)?)?; | |
| let p = total.to_scalar::<f32>()? as f64; | |
| if p.is_finite() { | |
| if epoca == 0 { | |
| perda_inicial = p; | |
| } | |
| perda_final = p; | |
| } | |
| let mut grads = crate::treino::grads_achatados(&total, &[nb.vars(), na.vars()].concat())?; | |
| crate::treino::AdamWManual::clipar(&mut grads, 1.0); | |
| otim.passo(grads); | |
| nb.ema_passo()?; | |
| na.ema_passo()?; | |
| } | |
| let z1 = nb.codificar(mod_a)?; | |
| Ok(RelatorioJepa { | |
| variante: "MC-JEPA".into(), | |
| perda_inicial, | |
| perda_final, | |
| perda_pred_inicial: perda_inicial, | |
| perda_pred_final: perda_final, | |
| std_emb_final: std_por_dim(&z1)?, | |
| std_emb_inicial: std_por_dim(&nb.codificar(mod_a)?)?, | |
| n_patches: n, | |
| d_patch: d_a.max(d_b), | |
| epocas: cfg.epocas, | |
| }) | |
| } | |
| /// **BJEPA** — bidirecional: preditor causal (z_t → z̃_{t+1}), anti-causal | |
| /// (z_{t+1} → z̃_t) e consistência de ciclo | |
| /// `‖P_b(P_f(z_t)) − sg(z_t)‖²` — a volta deve reconstruir o presente. | |
| pub fn treinar_bjepa(patches: &Tensor, cfg: &CfgJepa, dev: &Device) -> Result<RelatorioJepa> { | |
| let n = patches.dim(0)?; | |
| let d_patch = patches.dim(1)?; | |
| let mut nucleo = NucleoJepa::novo( | |
| d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev, | |
| )?; | |
| // preditor anti-causal: núcleo separado | |
| let mut ninv = NucleoJepa::novo( | |
| d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev, | |
| )?; | |
| let mut otim = | |
| crate::treino::AdamWManual::novo([nucleo.vars(), ninv.vars()].concat(), cfg.lr, 0.01); | |
| let mut perda_inicial = f64::NAN; | |
| let mut perda_final = f64::NAN; | |
| for epoca in 0..cfg.epocas { | |
| let z = nucleo.codificar(patches)?; // (n, d_emb) | |
| let zt = nucleo.alvo_codificar(patches)?; // alvos EMA (stop-grad) | |
| let mut perdas: Vec<Tensor> = Vec::new(); | |
| for t in 0..n.saturating_sub(1) { | |
| let zc = z.get(t)?.broadcast_as((1, cfg.d_emb))?; | |
| let z_prox = zt.get(t + 1)?.broadcast_as((1, cfg.d_emb))?; | |
| // causal: z_t → z̃_{t+1} (pos 0 = "próximo") | |
| let pf = nucleo.predizer(&zc, 0, None)?; | |
| let dif = (&pf - &z_prox)?; | |
| perdas.push((dif.sqr()?.sum_all()? / cfg.d_emb as f64)?); | |
| // anti-causal: z_{t+1} → z̃_t (pos 1 = "anterior") | |
| let zb = z.get(t + 1)?.broadcast_as((1, cfg.d_emb))?; | |
| let z_ant = zt.get(t)?.broadcast_as((1, cfg.d_emb))?; | |
| let pb = ninv.predizer(&zb, 1, None)?; | |
| let dif2 = (&pb - &z_ant)?; | |
| perdas.push((dif2.sqr()?.sum_all()? / cfg.d_emb as f64)?); | |
| // ciclo: P_b(P_f(z_t)) ≈ sg(z_t) | |
| let pf_sg = nucleo.predizer(&zc, 0, None)?.detach(); | |
| let pciclo = ninv.predizer(&pf_sg.broadcast_as((1, cfg.d_emb))?, 1, None)?; | |
| let zt_sg = zt.get(t)?.broadcast_as((1, cfg.d_emb))?; | |
| let dif3 = (&pciclo - &zt_sg)?; | |
| perdas.push((dif3.sqr()?.sum_all()? / cfg.d_emb as f64)?); | |
| } | |
| let l_pred = (Tensor::stack(&perdas, 0)?.sum_all()? / perdas.len().max(1) as f64)?; | |
| let l_vc = nucleo.perda_vc(&z)?; | |
| let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(cfg.lambda_vc, 0.0)?)?; | |
| let p = total.to_scalar::<f32>()? as f64; | |
| if p.is_finite() { | |
| if epoca == 0 { | |
| perda_inicial = p; | |
| } | |
| perda_final = p; | |
| } | |
| let mut grads = crate::treino::grads_achatados(&total, &[nucleo.vars(), ninv.vars()].concat())?; | |
| crate::treino::AdamWManual::clipar(&mut grads, 1.0); | |
| otim.passo(grads); | |
| nucleo.ema_passo()?; | |
| ninv.ema_passo()?; | |
| } | |
| let z1 = nucleo.codificar(patches)?; | |
| Ok(RelatorioJepa { | |
| variante: "BJEPA".into(), | |
| perda_inicial, | |
| perda_final, | |
| perda_pred_inicial: perda_inicial, | |
| perda_pred_final: perda_final, | |
| std_emb_final: std_por_dim(&z1)?, | |
| std_emb_inicial: std_por_dim(&nucleo.codificar(patches)?)?, | |
| n_patches: n, | |
| d_patch, | |
| epocas: cfg.epocas, | |
| }) | |
| } | |
| /// Esquema de máscara por variante a partir da grade de patches. | |
| pub fn mascara_por_variante( | |
| variante: VarianteJepa, | |
| grid: &[usize], | |
| frac: f64, | |
| rng: &mut RngDeterministico, | |
| ) -> Vec<bool> { | |
| match variante { | |
| VarianteJepa::ImageJepa => mascara_blocos_2d(grid[0], grid.get(1).copied().unwrap_or(1), 4, frac, rng), | |
| VarianteJepa::VJepa2 | VarianteJepa::Jepa3D => mascara_tubos_3d( | |
| grid[0], | |
| grid.get(1).copied().unwrap_or(1), | |
| grid.get(2).copied().unwrap_or(1), | |
| 3, | |
| frac, | |
| rng, | |
| ), | |
| VarianteJepa::AudioJepa => mascara_bandas(grid[0], grid.get(1).copied().unwrap_or(1), 3, frac, rng), | |
| VarianteJepa::BrainJepa => mascara_segmentos(grid[0], grid.get(1).copied().unwrap_or(1), frac, rng), | |
| _ => mascara_intervalos_1d(grid.iter().product::<usize>().max(1), frac, rng), | |
| } | |
| } | |
| pub fn versao() -> &'static str { | |
| "jepa-0.5.0" | |
| } | |