v0.9.0: análise/prova/teste das críticas Qwen3.8 lote 2 (r-34) — SIMD com despacho em runtime bit-a-bit (X-1), oráculo candle_nn::optim::AdamW (nosso 1,53× mais rápido), Thompson Sampling (regret 0,0019 vs UCB1 0,0156), GNG por entropia com piso (T-G1'), JEPA de fração adaptativa (bit-a-bit com 1 braço), configs tipadas, entrada numpy (2008× na fronteira), avaliação com dados reais — 104/104 gates, 111/111 cargo
d444358 verified | //! Suíte JEPA v0.5.0 — núcleo unificado + 12 variantes + SonicMoE + ReplaySSM. | |
| //! | |
| //! **Núcleo (todas as variantes compartilham):** | |
| //! - codificador `E_θ` (MLp SiLU), preditor `P_φ` (condicionado por posição/ | |
| //! ação/edição), codificador-alvo EMA `E_ξ` com `ξ ← m·ξ + (1−m)·θ` | |
| //! (gradiente cortado — alvo nunca recebe gradiente); | |
| //! - regularizador variância-covariância (VICReg-lite) anti-colapso: | |
| //! `V(z) = Σ_j max(0, γ − std(z_j))²` + fora-diagonal da covariância. | |
| //! - perda: `L = Σ_b ‖P(z_ctx, pos_b) − sg(ẑ_b)‖²/B + λ_vc·V`. | |
| //! | |
| //! **Prova J-1 (impossibilidade de colapso, em r-20):** se todas as | |
| //! incorporações colapsam para constante, `std(z_j) = 0 ∀j` ⇒ `V = d·γ² > 0`; | |
| //! logo o colapso total NÃO é minimizador com `λ_vc > 0`. | |
| //! | |
| //! **Prova J-2 (EMA = stop-grad + ponto fixo):** `∂L/∂ξ = 0` (detach); se | |
| //! `θ_t → θ*` então `ξ_t → θ*` (recursão linear contratante com `m < 1`). | |
| //! | |
| //! **Prova J-3 (decomposição de informação):** com `P` suficientemente | |
| //! rico, `min_P L = Σ_b Var(ẑ_b | z_ctx)` — a perda retém só o componentes | |
| //! previsíveis do alvo dado o contexto. | |
| //! | |
| //! **Variantes** (esquema de máscara + adaptador de modalidade): | |
| //! Image-JEPA (blocos 2D) · V-JEPA 2 (tubos 3D + cond. de ação) · MC-JEPA | |
| //! (cruz-modal, dois codificadores) · Audio-JEPA (bandas de espectrograma) · | |
| //! Point-JEPA (grupos FPS) · 3D-JEPA (voxels) · Brain-JEPA (canal×tempo) · | |
| //! T-JEPA (série temporal canal-independente) · LeJEPA (edição-condicionado) · | |
| //! BJEPA (bidirecional com consistência de ciclo). | |
| //! | |
| //! **SonicMoE** — MoE com roteamento top-k, capacidade por especialista | |
| //! (tile-aware) e perda auxiliar de balanceamento (argumento Switch: | |
| //! `L_aux = E·Σ f_e·p_e ≥ 1` com igualdade sse uniforme). | |
| //! | |
| //! **ReplaySSM** — SSM de porta diagonal (`h_t = σ(a)⊙h_{t−1} + | |
| //! (1−σ(a))⊙(Bx_t)`, estabilidade herdada do Teo S-2') + replay de experiências | |
| //! para aprendizado continuado (interpolação de gradientes limita o | |
| //! esquecimento — Teo R-1). | |
| use candle_core::{Device, Result, Tensor, Var}; | |
| use crate::camadas::Linear; | |
| use crate::rng::RngDeterministico; | |
| use std::collections::HashMap; | |
| // --------------------------------------------------------------------------- | |
| // Parte 1 — núcleo | |
| // --------------------------------------------------------------------------- | |
| /// MLP 2 camadas com SiLU (d → h → d_saida). | |
| pub struct MlpJepa { | |
| pub l1: Linear, | |
| pub l2: Linear, | |
| pub d_entrada: usize, | |
| pub d_saida: usize, | |
| } | |
| impl MlpJepa { | |
| pub fn novo(d_entrada: usize, d_oculta: usize, d_saida: usize, dev: &Device) -> Result<Self> { | |
| let s1 = (1.0 / (d_entrada as f64).sqrt()) * 0.9; | |
| let s2 = (1.0 / (d_oculta as f64).sqrt()) * 0.9; | |
| Ok(Self { | |
| l1: Linear::sem_vies(d_entrada, d_oculta, s1, dev)?, | |
| l2: Linear::sem_vies(d_oculta, d_saida, s2, dev)?, | |
| d_entrada, | |
| d_saida, | |
| }) | |
| } | |
| /// v0.9.0 — `novo` com RNG determinístico (quita o débito r-33 §8.1). | |
| /// Mesma distribuição (N(0, s), s = 0,9/√d) — só a FONTE de aleatoriedade | |
| /// muda: `RngDeterministico` semeado em vez do RNG global do candle. | |
| pub fn novo_semeado( | |
| d_entrada: usize, | |
| d_oculta: usize, | |
| d_saida: usize, | |
| dev: &Device, | |
| rng: &mut RngDeterministico, | |
| ) -> Result<Self> { | |
| let s1 = (1.0 / (d_entrada as f64).sqrt()) * 0.9; | |
| let s2 = (1.0 / (d_oculta as f64).sqrt()) * 0.9; | |
| Ok(Self { | |
| l1: Linear::sem_vies_semeado(d_entrada, d_oculta, s1, dev, rng)?, | |
| l2: Linear::sem_vies_semeado(d_oculta, d_saida, s2, dev, rng)?, | |
| d_entrada, | |
| d_saida, | |
| }) | |
| } | |
| pub fn fwd(&self, x: &Tensor) -> Result<Tensor> { | |
| let h = candle_nn::ops::silu(&self.l1.fwd(x)?)?; | |
| self.l2.fwd(&h) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| [self.l1.vars(), self.l2.vars()].concat() | |
| } | |
| /// EMA: `self ← m·self + (1−m)·outro` (usado no codificador-alvo). | |
| pub fn ema_de(&mut self, outro: &MlpJepa, m: f64) -> Result<()> { | |
| for (a, b) in self.vars().iter().zip(outro.vars().iter()) { | |
| let novo = (a.as_tensor().affine(m, 0.0)? + b.as_tensor().affine(1.0 - m, 0.0)?)?; | |
| a.set(&novo)?; | |
| } | |
| Ok(()) | |
| } | |
| } | |
| /// Núcleo JEPA: codificador + preditor + alvo EMA + regularizador VC. | |
| pub struct NucleoJepa { | |
| pub encoder: MlpJepa, | |
| pub preditor: MlpJepa, | |
| pub alvo: MlpJepa, | |
| /// tabela de posições (blocos) — entrada extra do preditor. | |
| pub pos: Var, | |
| pub d_emb: usize, | |
| pub d_cond: usize, | |
| pub m: f64, | |
| pub lambda_vc: f64, | |
| pub gamma_vc: f64, | |
| } | |
| impl NucleoJepa { | |
| pub fn novo( | |
| d_patch: usize, | |
| d_emb: usize, | |
| d_cond: usize, | |
| n_pos: usize, | |
| m: f64, | |
| lambda_vc: f64, | |
| gamma_vc: f64, | |
| dev: &Device, | |
| ) -> Result<Self> { | |
| Self::novo_com_semente(d_patch, d_emb, d_cond, n_pos, m, lambda_vc, gamma_vc, 42, dev) | |
| } | |
| /// v0.9.0 — construção com INICIALIZAÇÃO DETERMINÍSTICA (quita o débito | |
| /// r-33 §8.1): pesos e tabela de posições nascem de | |
| /// `RngDeterministico::novo(semente ^ 0x5EED)` — fluxo SEPARADO do RNG de | |
| /// máscara (`semente`), que permanece intocado. `novo` (acima) delega com | |
| /// semente fixa 42 para compatibilidade das demos. | |
| pub fn novo_com_semente( | |
| d_patch: usize, | |
| d_emb: usize, | |
| d_cond: usize, | |
| n_pos: usize, | |
| m: f64, | |
| lambda_vc: f64, | |
| gamma_vc: f64, | |
| semente: u64, | |
| dev: &Device, | |
| ) -> Result<Self> { | |
| let mut rng = RngDeterministico::novo(semente ^ 0x5EED_u64); | |
| let encoder = MlpJepa::novo_semeado(d_patch, 2 * d_emb, d_emb, dev, &mut rng)?; | |
| let preditor = MlpJepa::novo_semeado(d_emb + d_cond + d_emb, 2 * d_emb, d_emb, dev, &mut rng)?; | |
| // alvo nasce IGUAL ao encoder (ponto de partida do EMA) | |
| let mut alvo = MlpJepa::novo_semeado(d_patch, 2 * d_emb, d_emb, dev, &mut rng)?; | |
| for (va, vb) in alvo.vars().iter().zip(encoder.vars().iter()) { | |
| va.set(vb.as_tensor())?; | |
| } | |
| let n_posicoes = n_pos.max(1); | |
| let pos_dados: Vec<f32> = (0..n_posicoes * d_emb) | |
| .map(|_| rng.normal(0.0, 0.02) as f32) | |
| .collect(); | |
| let pos_t = Tensor::from_vec(pos_dados, (n_posicoes, d_emb), dev)?; | |
| Ok(Self { | |
| encoder, | |
| preditor, | |
| alvo, | |
| pos: Var::from_tensor(&pos_t)?, | |
| d_emb, | |
| d_cond, | |
| m, | |
| lambda_vc, | |
| gamma_vc, | |
| }) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| let mut v = self.encoder.vars(); | |
| v.extend(self.preditor.vars()); | |
| v.push(self.pos.clone()); | |
| v // NOTA: alvo NÃO entra — é buffer EMA | |
| } | |
| pub fn ema_passo(&mut self) -> Result<()> { | |
| self.alvo.ema_de(&self.encoder, self.m) | |
| } | |
| /// Codifica um conjunto de patches (n, d_patch) → (n, d_emb). | |
| pub fn codificar(&self, patches: &Tensor) -> Result<Tensor> { | |
| self.encoder.fwd(patches) | |
| } | |
| /// Codifica com o alvo EMA — SEM gradiente (detach por construção: as | |
| /// vars do alvo não estão no otimizador e o tensor é desconectado). | |
| pub fn alvo_codificar(&self, patches: &Tensor) -> Result<Tensor> { | |
| Ok(self.alvo.fwd(patches)?.detach()) | |
| } | |
| /// Prediz o embedding do bloco `i_pos` a partir do contexto `z_ctx` | |
| /// (d_emb) e da condição `cond` (d_cond; zeros se None). | |
| pub fn predizer(&self, z_ctx: &Tensor, i_pos: usize, cond: Option<&Tensor>) -> Result<Tensor> { | |
| // z_ctx: (b, d_emb) ou (d_emb,) | |
| let b = z_ctx.dim(0).unwrap_or(1); | |
| let dev = z_ctx.device(); | |
| let pos = self.pos.as_tensor().get(i_pos.min(self.pos.dim(0)?.saturating_sub(1)))?; | |
| let pos_b = pos.broadcast_as((b, self.d_emb))?; | |
| let cond_b = match cond { | |
| Some(c) => { | |
| if c.dims().len() == 1 && b > 1 { | |
| c.broadcast_as((b, self.d_cond))? | |
| } else { | |
| c.clone() | |
| } | |
| } | |
| None => Tensor::zeros((b, self.d_cond), candle_core::DType::F32, dev)?, | |
| }; | |
| let entrada = Tensor::cat(&[z_ctx, &pos_b, &cond_b], candle_core::D::Minus1)?; | |
| self.preditor.fwd(&entrada) | |
| } | |
| /// Regularizador variância-covariância do lote de incorporações z (n, d). | |
| /// | |
| /// **J-1:** no colapso total `z_i = c ∀i` ⇒ `std_j = 0` ⇒ `V = d·γ²`. | |
| pub fn perda_vc(&self, z: &Tensor) -> Result<Tensor> { | |
| let n = z.dim(0)? as f64; | |
| let d = self.d_emb; | |
| let media = z.mean_keepdim(0)?; // (1, d) | |
| let centr = z.broadcast_sub(&media)?; | |
| let var = (centr.sqr()?.sum(0)? / n)?; // (d,) | |
| let std = var.sqrt()?; | |
| // dobradiça: max(0, γ − std)² | |
| let gamma_t = | |
| Tensor::new(self.gamma_vc as f32, z.device())?.broadcast_as(std.dims())?; | |
| let gap = gamma_t.broadcast_sub(&std)?.clamp(0.0, f64::INFINITY)?; | |
| let v = gap.sqr()?.sum_all()?; | |
| // covariância fora-diagonal: C = centrᵀ·centr/n; perda = ‖C − diag‖²_F/(2d) | |
| let cov = (centr.transpose(0, 1)?.matmul(¢r)? / n)?; // (d, d) | |
| let mask_diag = { | |
| let mut m = vec![0f32; d * d]; | |
| for i in 0..d { | |
| m[i * d + i] = 1.0; | |
| } | |
| Tensor::from_vec(m, (d, d), z.device())? | |
| }; | |
| let fora = ((&cov * &mask_diag)?.sqr()?.sum_all()? / (2.0 * d.max(1) as f64))?; | |
| let v_norm = (v / d.max(1) as f64)?; | |
| Ok((v_norm + fora)?) | |
| } | |
| /// Perda JEPA padrão para UM lote: contexto (número variável de patches), | |
| /// lista de blocos-alvo (posições), condição opcional. | |
| /// | |
| /// `z_ctx` = incorporação agregada do contexto (b, d_emb); | |
| /// `alvos`: para cada elemento do lote, Vec<(i_pos, patches_alvo)>. | |
| pub fn perda_lote( | |
| &self, | |
| z_ctx: &Tensor, | |
| alvos: &[Vec<(usize, Tensor)>], | |
| cond: Option<&Tensor>, | |
| z_ctx_todos: &Tensor, | |
| ) -> Result<(Tensor, f64, f64)> { | |
| let b = z_ctx.dim(0)?; | |
| let dev = z_ctx.device(); | |
| let mut perdas: Vec<Tensor> = Vec::new(); | |
| let mut n_alvos = 0usize; | |
| for (i, blocos) in alvos.iter().enumerate() { | |
| for (i_pos, patches_alvo) in blocos { | |
| // ẑ_b = média das incorporações-alvo dos patches do bloco | |
| let zt = self.alvo_codificar(patches_alvo)?; // (m, d_emb) | |
| let zt_m = zt.mean(0)?; // (d_emb,) | |
| let zt_b = zt_m.broadcast_as((1, self.d_emb))?; | |
| let zc = z_ctx.get(i)?.broadcast_as((1, self.d_emb))?; | |
| let c_i = cond.map(|c| c.get(i)?.broadcast_as((1, self.d_cond))).transpose()?; | |
| let pred = self.predizer(&zc, *i_pos, c_i.as_ref())?; | |
| let dif = (&pred - &zt_b)?; | |
| let l = (dif.sqr()?.sum_all()? / self.d_emb as f64)?; | |
| perdas.push(l); | |
| n_alvos += 1; | |
| } | |
| } | |
| let l_pred = if perdas.is_empty() { | |
| Tensor::zeros((), candle_core::DType::F32, &dev)? | |
| } else { | |
| (Tensor::stack(&perdas, 0)?.sum_all()? / n_alvos.max(1) as f64)? | |
| }; | |
| let l_vc = self.perda_vc(z_ctx_todos)?; | |
| let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(self.lambda_vc, 0.0)?)?; | |
| let lp = l_pred.to_scalar::<f32>()? as f64; | |
| let lv = l_vc.to_scalar::<f32>()? as f64; | |
| Ok((total, lp, lv)) | |
| } | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Parte 2 — máscaras | |
| // --------------------------------------------------------------------------- | |
| /// Máscara de blocos 2D (Image-JEPA): grade (h, w) de patches; sorteia | |
| /// `n_blocos` retângulos aleatórios com fração alvo `frac` de patches. | |
| pub fn mascara_blocos_2d(h: usize, w: usize, n_blocos: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> { | |
| let total = h * w; | |
| let alvo = (total as f64 * frac) as usize; | |
| let mut masc = vec![false; total]; | |
| let mut coberto = 0usize; | |
| for _ in 0..n_blocos.max(1) { | |
| if coberto >= alvo { | |
| break; | |
| } | |
| let bh = 1 + rng.inteiro(h.min(4)); | |
| let bw = 1 + rng.inteiro(w.min(4)); | |
| let y0 = rng.inteiro(h + 1 - bh); | |
| let x0 = rng.inteiro(w + 1 - bw); | |
| for y in y0..y0 + bh { | |
| for x in x0..x0 + bw { | |
| let i = y * w + x; | |
| if !masc[i] { | |
| masc[i] = true; | |
| coberto += 1; | |
| } | |
| } | |
| } | |
| } | |
| masc | |
| } | |
| /// Máscara de tubos 3D (V-JEPA 2): grade (t, h, w); tubos = retângulos que | |
| /// atravessam o tempo. | |
| pub fn mascara_tubos_3d(t: usize, h: usize, w: usize, n_tubos: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> { | |
| let total = t * h * w; | |
| let alvo = (total as f64 * frac) as usize; | |
| let mut masc = vec![false; total]; | |
| let mut coberto = 0usize; | |
| for _ in 0..n_tubos.max(1) { | |
| if coberto >= alvo { | |
| break; | |
| } | |
| let bh = 1 + rng.inteiro(h.min(3)); | |
| let bw = 1 + rng.inteiro(w.min(3)); | |
| let y0 = rng.inteiro(h + 1 - bh); | |
| let x0 = rng.inteiro(w + 1 - bw); | |
| for tt in 0..t { | |
| for y in y0..y0 + bh { | |
| for x in x0..x0 + bw { | |
| let i = (tt * h + y) * w + x; | |
| if !masc[i] { | |
| masc[i] = true; | |
| coberto += 1; | |
| } | |
| } | |
| } | |
| } | |
| } | |
| masc | |
| } | |
| /// Máscara de bandas (Audio-JEPA): grade (f, tau) — bandas de frequência | |
| /// horizontais com extensão temporal. | |
| pub fn mascara_bandas(f: usize, tau: usize, n_bandas: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> { | |
| mascara_blocos_2d(f, tau, n_bandas, frac, rng) // blocos na grade (f, tau) | |
| } | |
| /// Máscara de blocos 3D (3D-JEPA): grade (x, y, z) de voxels. | |
| pub fn mascara_blocos_3d(x: usize, y: usize, z: usize, n_blocos: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> { | |
| mascara_tubos_3d(x, y, z, n_blocos, frac, rng) | |
| } | |
| /// Máscara de segmentos canal×tempo (Brain-JEPA): canais inteiros OU | |
| /// janelas temporais (aleatório por chamada). | |
| pub fn mascara_segmentos(c: usize, t: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> { | |
| let total = c * t; | |
| let alvo = (total as f64 * frac) as usize; | |
| let mut masc = vec![false; total]; | |
| if rng.inteiro(2) == 0 { | |
| // canais inteiros | |
| let mut coberto = 0; | |
| let mut canais: Vec<usize> = (0..c).collect(); | |
| // embaralhar | |
| for i in (1..canais.len()).rev() { | |
| let j = rng.inteiro(i + 1); | |
| canais.swap(i, j); | |
| } | |
| for &ch in canais.iter() { | |
| if coberto >= alvo { | |
| break; | |
| } | |
| for tt in 0..t { | |
| masc[ch * t + tt] = true; | |
| coberto += 1; | |
| } | |
| } | |
| } else { | |
| // janelas temporais | |
| let mut coberto = 0; | |
| let mut t0 = 0; | |
| while coberto < alvo && t0 < t { | |
| let len = 1 + rng.inteiro((t / 4).max(1)); | |
| for ch in 0..c { | |
| for tt in t0..(t0 + len).min(t) { | |
| masc[ch * t + tt] = true; | |
| coberto += 1; | |
| } | |
| } | |
| t0 += len + 1; | |
| } | |
| } | |
| masc | |
| } | |
| /// Máscara 1D (T-JEPA): intervalos contíguos em n posições. | |
| pub fn mascara_intervalos_1d(n: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> { | |
| let alvo = (n as f64 * frac) as usize; | |
| let mut masc = vec![false; n]; | |
| let mut coberto = 0; | |
| while coberto < alvo { | |
| let len = 1 + rng.inteiro((n / 8).max(1)); | |
| let i0 = rng.inteiro(n); | |
| for i in i0..(i0 + len).min(n) { | |
| if !masc[i] { | |
| masc[i] = true; | |
| coberto += 1; | |
| } | |
| } | |
| } | |
| masc | |
| } | |
| /// Amostragem dos pontos mais distantes (FPS) para Point-JEPA. | |
| pub fn pontos_mais_distantes(pontos: &[Vec<f64>], k: usize, rng: &mut RngDeterministico) -> Vec<usize> { | |
| let n = pontos.len(); | |
| if n == 0 || k == 0 { | |
| return vec![]; | |
| } | |
| let k = k.min(n); | |
| let mut centros = vec![rng.inteiro(n)]; | |
| let mut dist_min: Vec<f64> = pontos | |
| .iter() | |
| .map(|p| distancia_eucl(p, &pontos[centros[0]])) | |
| .collect(); | |
| while centros.len() < k { | |
| let prox = dist_min | |
| .iter() | |
| .enumerate() | |
| .filter(|(i, _)| !centros.contains(i)) | |
| .max_by(|a, b| a.1.partial_cmp(b.1).unwrap()) | |
| .map(|(i, _)| i) | |
| .unwrap_or(0); | |
| centros.push(prox); | |
| for (i, p) in pontos.iter().enumerate() { | |
| let d = distancia_eucl(p, &pontos[prox]); | |
| if d < dist_min[i] { | |
| dist_min[i] = d; | |
| } | |
| } | |
| } | |
| centros | |
| } | |
| fn distancia_eucl(a: &[f64], b: &[f64]) -> f64 { | |
| a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f64>().sqrt() | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Parte 3 — adaptadores de modalidade | |
| // --------------------------------------------------------------------------- | |
| /// (C, H, W) → patches (H/ph · W/pw, C·ph·pw) — Image-JEPA. | |
| pub fn em_patches_2d(img: &Tensor, ph: usize, pw: usize) -> Result<Tensor> { | |
| let (c, h, w) = (img.dim(0)?, img.dim(1)?, img.dim(2)?); | |
| let dev = img.device(); | |
| let mut linhas: Vec<Tensor> = Vec::new(); | |
| for y in (0..h).step_by(ph) { | |
| for x in (0..w).step_by(pw) { | |
| let hh = ph.min(h - y); | |
| let ww = pw.min(w - x); | |
| let patch = img.narrow(0, 0, c)?.narrow(1, y, hh)?.narrow(2, x, ww)?; | |
| // padding para (C, ph, pw) se borda | |
| let patch = if hh < ph || ww < pw { | |
| let mut buf = patch.flatten_all()?.to_vec1::<f32>()?; | |
| let falta = c * ph * pw - buf.len(); | |
| buf.extend(std::iter::repeat(0f32).take(falta)); | |
| Tensor::from_vec(buf, (c * ph * pw,), &dev)? | |
| } else { | |
| patch.flatten_all()? | |
| }; | |
| linhas.push(patch); | |
| } | |
| } | |
| Ok(Tensor::stack(&linhas, 0)?) | |
| } | |
| /// (T, H, W) → tubos (T/pt · H/ph · W/pw, pt·ph·pw) — V-JEPA 2. | |
| pub fn em_tubos_3d(vid: &Tensor, pt: usize, ph: usize, pw: usize) -> Result<Tensor> { | |
| let (t, h, w) = (vid.dim(0)?, vid.dim(1)?, vid.dim(2)?); | |
| let dev = vid.device(); | |
| let mut linhas: Vec<Tensor> = Vec::new(); | |
| for tt in (0..t).step_by(pt) { | |
| let tl = pt.min(t - tt); | |
| for y in (0..h).step_by(ph) { | |
| for x in (0..w).step_by(pw) { | |
| let hh = ph.min(h - y); | |
| let ww = pw.min(w - x); | |
| let tubo = vid | |
| .narrow(0, tt, tl)? | |
| .narrow(1, y, hh)? | |
| .narrow(2, x, ww)? | |
| .flatten_all()?; | |
| let mut buf = tubo.to_vec1::<f32>()?; | |
| let falta = pt * ph * pw - buf.len(); | |
| buf.extend(std::iter::repeat(0f32).take(falta)); | |
| linhas.push(Tensor::from_vec(buf, (pt * ph * pw,), &dev)?); | |
| } | |
| } | |
| } | |
| Ok(Tensor::stack(&linhas, 0)?) | |
| } | |
| /// (F, τ) → patches (F/pf · τ/pt, pf·pt) — Audio-JEPA. | |
| pub fn em_patches_espectrograma(spec: &Tensor, pf: usize, pt: usize) -> Result<Tensor> { | |
| let (f, tau) = (spec.dim(0)?, spec.dim(1)?); | |
| let dev = spec.device(); | |
| let mut linhas: Vec<Tensor> = Vec::new(); | |
| for i in (0..f).step_by(pf) { | |
| for j in (0..tau).step_by(pt) { | |
| let ff = pf.min(f - i); | |
| let tt = pt.min(tau - j); | |
| let p = spec.narrow(0, i, ff)?.narrow(1, j, tt)?.flatten_all()?; | |
| let mut buf = p.to_vec1::<f32>()?; | |
| let falta = pf * pt - buf.len(); | |
| buf.extend(std::iter::repeat(0f32).take(falta)); | |
| linhas.push(Tensor::from_vec(buf, (pf * pt,), &dev)?); | |
| } | |
| } | |
| Ok(Tensor::stack(&linhas, 0)?) | |
| } | |
| /// (C, T) → patches canal-independentes (C · T/pt, pt) — T-JEPA/Brain-JEPA. | |
| pub fn em_patches_serie(ts: &Tensor, pt: usize) -> Result<Tensor> { | |
| let (c, t) = (ts.dim(0)?, ts.dim(1)?); | |
| let dev = ts.device(); | |
| let mut linhas: Vec<Tensor> = Vec::new(); | |
| for ch in 0..c { | |
| for j in (0..t).step_by(pt) { | |
| let tt = pt.min(t - j); | |
| let p = ts.narrow(0, ch, 1)?.narrow(1, j, tt)?.flatten_all()?; | |
| let mut buf = p.to_vec1::<f32>()?; | |
| let falta = pt - buf.len(); | |
| buf.extend(std::iter::repeat(0f32).take(falta)); | |
| linhas.push(Tensor::from_vec(buf, (pt,), &dev)?); | |
| } | |
| } | |
| Ok(Tensor::stack(&linhas, 0)?) | |
| } | |
| /// (X, Y, Z) → voxels (X/px · Y/py · Z/pz, px·py·pz) — 3D-JEPA. | |
| pub fn em_voxels(vol: &Tensor, px: usize, py: usize, pz: usize) -> Result<Tensor> { | |
| let (x, y, z) = (vol.dim(0)?, vol.dim(1)?, vol.dim(2)?); | |
| let dev = vol.device(); | |
| let mut linhas: Vec<Tensor> = Vec::new(); | |
| for i in (0..x).step_by(px) { | |
| for j in (0..y).step_by(py) { | |
| for k in (0..z).step_by(pz) { | |
| let lx = px.min(x - i); | |
| let ly = py.min(y - j); | |
| let lz = pz.min(z - k); | |
| let p = vol.narrow(0, i, lx)?.narrow(1, j, ly)?.narrow(2, k, lz)?.flatten_all()?; | |
| let mut buf = p.to_vec1::<f32>()?; | |
| let falta = px * py * pz - buf.len(); | |
| buf.extend(std::iter::repeat(0f32).take(falta)); | |
| linhas.push(Tensor::from_vec(buf, (px * py * pz,), &dev)?); | |
| } | |
| } | |
| } | |
| Ok(Tensor::stack(&linhas, 0)?) | |
| } | |
| /// Nuvem de pontos (n, d) → grupos FPS (k, m·d): k centros por FPS, cada | |
| /// grupo = os m vizinhos mais próximos do centro (concatenados) — Point-JEPA. | |
| pub fn em_grupos_pontos(pc: &Tensor, k: usize, m: usize, rng: &mut RngDeterministico) -> Result<Tensor> { | |
| let n = pc.dim(0)?; | |
| let d = pc.dim(1)?; | |
| let coords = pc.to_vec2::<f32>()?; | |
| let coords_f: Vec<Vec<f64>> = coords.iter().map(|r| r.iter().map(|&v| v as f64).collect()).collect(); | |
| let centros = pontos_mais_distantes(&coords_f, k, rng); | |
| let dev = pc.device(); | |
| let mut linhas: Vec<Tensor> = Vec::new(); | |
| for &ci in ¢ros { | |
| // m vizinhos mais próximos do centro ci (incluindo ele próprio) | |
| let mut ord: Vec<usize> = (0..n).collect(); | |
| ord.sort_by(|&a, &b| { | |
| let da = distancia_eucl(&coords_f[a], &coords_f[ci]); | |
| let db = distancia_eucl(&coords_f[b], &coords_f[ci]); | |
| da.partial_cmp(&db).unwrap() | |
| }); | |
| let mut buf: Vec<f32> = Vec::with_capacity(m * d); | |
| for &idx in ord.iter().take(m) { | |
| for j in 0..d { | |
| buf.push(coords[idx][j]); | |
| } | |
| } | |
| let falta = m * d - buf.len(); | |
| buf.extend(std::iter::repeat(0f32).take(falta)); | |
| linhas.push(Tensor::from_vec(buf, (m * d,), &dev)?); | |
| } | |
| Ok(Tensor::stack(&linhas, 0)?) | |
| } | |
| // --------------------------------------------------------------------------- | |
| // 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_com_semente( | |
| d_patch, | |
| cfg.d_emb, | |
| if cond.is_some() { cond.unwrap().dim(1)? } else { 0 }, | |
| n, | |
| cfg.m, | |
| cfg.lambda_vc, | |
| cfg.gamma_vc, | |
| cfg.semente, | |
| 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() | |
| } | |
| // --------------------------------------------------------------------------- | |
| // v0.9.0 — fração de máscara ADAPTATIVA por Thompson Sampling | |
| // (crítica Qwen3.8 C; provas T-J1/T-J2 em r-34 §7) | |
| // --------------------------------------------------------------------------- | |
| /// Relatório do treino com fração adaptativa. | |
| pub struct RelatorioJepaAdaptativa { | |
| pub base: RelatorioJepa, | |
| pub fracs: Vec<f64>, | |
| /// (época, braço, frac, perda_pred) por época — trilha auditável. | |
| pub escolhas: Vec<(usize, usize, f64, f64)>, | |
| /// nº de épocas por braço. | |
| pub jogadas_por_braco: Vec<usize>, | |
| /// médias posteriores finais (estimativa de qualidade relativa). | |
| pub posterior: Vec<f64>, | |
| } | |
| /// **T-J1 (conexão Information Bottleneck — o porquê de a fração importar):** | |
| /// o JEPA minimiza `E‖s(y) − f(s(x_ctx))‖²` com alvo mascarado. Sob ruído | |
| /// gaussiano homocedástico, minimizar o MSE equivale a maximizar um limite | |
| /// inferior de `I(Z; Y | X_ctx)` (o preditor ótimo é `E[s(Y)|X_ctx]` e o | |
| /// resíduo é o ruído epistêmico — princípio InfoMax). A fração `f` controla o | |
| /// TRADE-OFF: f→0 torna o alvo visível no contexto (atalho de identidade, | |
| /// informação trivial); f→1 destrói o contexto (erro de estimação domina). | |
| /// A perda preditiva em função de f é, portanto, em U — com ótimo dependente | |
| /// dos dados. NENHUM f fixo é ótimo universalmente (por isso a adaptação). | |
| /// | |
| /// **T-J2 (esquema e recompensa honesta):** um bandit Thompson (T-TS1) sobre | |
| /// `f ∈ fracs` escolhe a fração de cada época. A recompensa usa a PERDA | |
| /// PREDITIVA ISOLADA `l_p` (critério honesto herdado da v0.5.0 — a VC pode | |
| /// crescer legitimamente) e mede MELHORIA RELATIVA do próprio braço: | |
| /// `r = ½ + (l_prev − l)/(2·l_prev)` ∈ [0,1] — ½ = estagnação, >½ = | |
| /// melhorando. Escalar por l_prev torna r comparável entre regimes de perda | |
| /// diferentes; oscilações simétricas não enganam o posterior em expectativa | |
| /// (r médio = ½). O RNG do bandit é um FLUXO SEPARADO (semente derivada) — | |
| /// o fluxo de máscara (cfg.semente) permanece intocado: com um único braço, | |
| /// o treino é BIT-A-BIT idêntico ao `treinar_variante` com fração fixa | |
| /// (teste `jepa_adaptativa_um_braco_e_fixa`). | |
| pub fn treinar_variante_adaptativa( | |
| patches: &Tensor, | |
| grid: &[usize], | |
| fracs: &[f64], | |
| cfg: &CfgJepa, | |
| variante: VarianteJepa, | |
| cond: Option<&Tensor>, | |
| dev: &Device, | |
| ) -> Result<RelatorioJepaAdaptativa> { | |
| if fracs.is_empty() { | |
| candle_core::bail!("treinar_variante_adaptativa: fracs não pode ser vazio"); | |
| } | |
| let fracs: Vec<f64> = fracs.iter().cloned().map(|f| f.clamp(0.05, 0.95)).collect(); | |
| // fluxo de RNG SEPARADO do bandit — preserva o fluxo de máscara | |
| let mut bandit = crate::controlador::ThompsonSampling::novo( | |
| fracs.iter().map(|f| format!("frac={f:.2}")).collect(), | |
| cfg.semente ^ 0xBADC0DEu64, | |
| ); | |
| let mut ultima_perda_braco = vec![f64::NAN; fracs.len()]; | |
| let mut escolhas: Vec<(usize, usize, f64, f64)> = Vec::with_capacity(cfg.epocas); | |
| // ---- corpo: réplica fiel do loop de `treinar_variante` (mesma ordem de | |
| // ops; única diferença: a fração vem do bandit por época) ---- | |
| let n = patches.dim(0)?; | |
| let d_patch = patches.dim(1)?; | |
| let mut rng = RngDeterministico::novo(cfg.semente); | |
| let mut nucleo = NucleoJepa::novo_com_semente( | |
| d_patch, | |
| cfg.d_emb, | |
| if cond.is_some() { cond.unwrap().dim(1)? } else { 0 }, | |
| n, | |
| cfg.m, | |
| cfg.lambda_vc, | |
| cfg.gamma_vc, | |
| cfg.semente, | |
| dev, | |
| )?; | |
| let otimo_vars = nucleo.vars(); | |
| let mut otimizador = crate::treino::AdamWManual::novo(otimo_vars, cfg.lr, 0.01); | |
| 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 { | |
| // === escolha do braço (fração) desta época === | |
| let k = bandit.escolher(); | |
| let frac_braco = fracs[k]; | |
| let (ctx_idx, blocos) = loop { | |
| let mut masc = mascara_por_variante(variante, grid, frac_braco, &mut rng); | |
| 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); | |
| } | |
| }; | |
| 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)?; | |
| let z_ctx = z_ctx_todos.mean_keepdim(0)?; | |
| 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)?)]); | |
| } | |
| 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; | |
| } | |
| // === recompensa do braço (T-J2): melhoria relativa da perda PREDITIVA === | |
| if lp.is_finite() { | |
| let anterior = ultima_perda_braco[k]; | |
| let r = if anterior.is_finite() && anterior > 0.0 { | |
| (0.5 + (anterior - lp) / (2.0 * anterior)).clamp(0.0, 1.0) | |
| } else { | |
| 0.5 // primeira medida do braço: neutra | |
| }; | |
| bandit.registrar_continuo(k, r); | |
| ultima_perda_braco[k] = lp; | |
| escolhas.push((epoca, k, frac_braco, lp)); | |
| } | |
| // decaimento cosseno do lr + clipping (idêntico ao fixo) | |
| 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); | |
| let jogadas_por_braco = bandit.jogadas.iter().map(|&c| c as usize).collect(); | |
| let posterior = bandit.medias_posteriori(); | |
| Ok(RelatorioJepaAdaptativa { | |
| base: 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, | |
| }, | |
| fracs, | |
| escolhas, | |
| jogadas_por_braco, | |
| posterior, | |
| }) | |
| } | |
| 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) | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Parte 5 — SonicMoE (IO/tile-aware) | |
| // --------------------------------------------------------------------------- | |
| /// MoE top-k com capacidade por especialista e perda de balanceamento. | |
| /// | |
| /// **Prova SM-1 (limite de capacidade):** o roteamento tile-aware mantém no | |
| /// máximo `C` tokens por especialista por ladrilho — o drop é determinístico | |
| /// (ordem de chegada) e o nº de tokens descartados é `Σ_e max(0, load_e − C)`. | |
| /// | |
| /// **Prova SM-2 (balanceamento):** `L_aux = E·Σ_e f_e·p_e`; com roteamento | |
| /// auto-consistente (`f = p`), `L_aux = E·Σ p_e² ≥ 1` com igualdade sse | |
| /// `p_e = 1/E` (Cauchy–Schchwarz) — o mínimo força o roteador uniforme. | |
| pub struct SonicMoe { | |
| pub roteador: Linear, // d → E | |
| pub especialistas: Vec<MlpJepa>, // E experts d → d | |
| pub n_top: usize, | |
| pub capacidade: usize, | |
| pub lambda_aux: f64, | |
| } | |
| pub struct SaidaMoe { | |
| pub saida: Tensor, | |
| pub perda_aux: Tensor, | |
| pub carga: Vec<usize>, // tokens efetivamente processados por especialista | |
| pub descartados: usize, // tokens com especialistas cortados pela capacidade | |
| } | |
| impl SonicMoe { | |
| pub fn novo(d: usize, n_expertos: usize, n_top: usize, capacidade: usize, dev: &Device) -> Result<Self> { | |
| let roteador = Linear::sem_vies(d, n_expertos, 0.1, dev)?; | |
| let especialistas = (0..n_expertos) | |
| .map(|_| MlpJepa::novo(d, 2 * d, d, dev)) | |
| .collect::<Result<Vec<_>>>()?; | |
| Ok(Self { roteador, especialistas, n_top, capacidade, lambda_aux: 0.01 }) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| let mut v = self.roteador.vars(); | |
| for e in &self.especialistas { | |
| v.extend(e.vars()); | |
| } | |
| v | |
| } | |
| /// Forward tile-aware: (n, d) → (n, d). | |
| pub fn fwd(&self, x: &Tensor) -> Result<SaidaMoe> { | |
| let n = x.dim(0)?; | |
| let d = x.dim(1)?; | |
| let e = self.especialistas.len(); | |
| let dev = x.device(); | |
| let logits = self.roteador.fwd(x)?; // (n, E) | |
| let probs = candle_nn::ops::softmax(&logits, candle_core::D::Minus1)?; | |
| // top-k por token (não diferenciável — pesos do top-k são dif. via probs) | |
| let probs_v = probs.to_vec2::<f32>()?; | |
| let mut carga = vec![0usize; e]; | |
| let mut descartados = 0usize; | |
| // acumulador de saída (n, d) — construído por soma de contribuições | |
| let mut saida = Tensor::zeros((n, d), candle_core::DType::F32, &dev)?; | |
| let mut peso_total = Tensor::zeros((n, 1), candle_core::DType::F32, &dev)?; | |
| // por especialista (tile): tokens cujo top-k contém e e com carga < C | |
| for ei in 0..e { | |
| // candidatos: tokens ordenados por prob do especialista ei | |
| let mut cand: Vec<(usize, f32)> = (0..n) | |
| .map(|i| (i, probs_v[i][ei])) | |
| .collect(); | |
| cand.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); | |
| let mut escolhidos: Vec<usize> = Vec::new(); | |
| for (i, _) in cand { | |
| if carga[ei] >= self.capacidade { | |
| break; | |
| } | |
| // verifica se ei está no top-k do token i | |
| let mut ordem: Vec<(f32, usize)> = | |
| (0..e).map(|k| (probs_v[i][k], k)).collect(); | |
| ordem.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); | |
| let top: Vec<usize> = ordem.iter().rev().take(self.n_top).map(|(_, k)| *k).collect(); | |
| if top.contains(&ei) { | |
| escolhidos.push(i); | |
| carga[ei] += 1; | |
| } | |
| } | |
| if escolhidos.is_empty() { | |
| continue; | |
| } | |
| // tile do especialista: uma única passagem em lote | |
| let idx: Vec<u32> = escolhidos.iter().map(|&i| i as u32).collect(); | |
| let n_e = idx.len(); | |
| let idx_t = Tensor::from_vec(idx, (n_e,), &dev)?; | |
| let tile = x.index_select(&idx_t, 0)?; | |
| let y = self.especialistas[ei].fwd(&tile)?; // (m, d) | |
| // pesos: prob do especialista ei para cada token escolhido | |
| let w: Vec<f32> = escolhidos.iter().map(|&i| probs_v[i][ei]).collect(); | |
| let n_w = escolhidos.len(); | |
| let w_t = Tensor::from_vec(w.clone(), (n_w, 1), &dev)?; | |
| let contrib = y.broadcast_mul(&w_t)?; | |
| // scatter-add: acumula nas posições | |
| let mut buf = saida.to_vec2::<f32>()?; | |
| let contrib_v = contrib.to_vec2::<f32>()?; | |
| for (k, &i) in escolhidos.iter().enumerate() { | |
| for j in 0..d { | |
| buf[i][j] += contrib_v[k][j]; | |
| } | |
| } | |
| saida = Tensor::from_vec( | |
| buf.into_iter().flatten().collect::<Vec<f32>>(), | |
| (n, d), | |
| &dev, | |
| )?; | |
| let mut pt = peso_total.to_vec2::<f32>()?; | |
| for (k, &i) in escolhidos.iter().enumerate() { | |
| pt[i][0] += w[k]; | |
| } | |
| peso_total = Tensor::from_vec( | |
| pt.into_iter().flatten().collect::<Vec<f32>>(), | |
| (n, 1), | |
| &dev, | |
| )?; | |
| } | |
| // tokens sem especialista (todos cortados) → descartados; saída 0 | |
| let pt_v = peso_total.to_vec2::<f32>()?; | |
| for i in 0..n { | |
| if pt_v[i][0] == 0.0 { | |
| descartados += 1; | |
| } | |
| } | |
| // renormaliza pelo peso efetivo (evita viés da capacidade) | |
| let denom = peso_total.clamp(1e-6, f64::INFINITY)?; | |
| let saida = saida.broadcast_div(&denom)?; | |
| // perda auxiliar: E·Σ f_e·p_e (f = fração de tokens cujo top-1 é e) | |
| let mut f = vec![0f32; e]; | |
| for i in 0..n { | |
| let (mut melhor, mut mp) = (0usize, f32::NEG_INFINITY); | |
| for k in 0..e { | |
| if probs_v[i][k] > mp { | |
| mp = probs_v[i][k]; | |
| melhor = k; | |
| } | |
| } | |
| f[melhor] += 1.0 / n as f32; | |
| } | |
| let f_t = Tensor::from_vec(f, (e,), &dev)?; | |
| let p_media = probs.mean(0)?; // (E,) | |
| let perda_aux = ((&f_t * &p_media)?.sum_all()? * (e as f64))?; | |
| Ok(SaidaMoe { saida, perda_aux, carga, descartados }) | |
| } | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Parte 6 — ReplaySSM | |
| // --------------------------------------------------------------------------- | |
| /// SSM de porta diagonal com replay para aprendizado continuado. | |
| /// | |
| /// Recorrência: `h_t = σ(a) ⊙ h_{t−1} + (1 − σ(a)) ⊙ (B x_t)`. | |
| /// | |
| /// **Prova R-1 (estabilidade):** `|σ(a)| < 1` ⇒ perturbações de `h_0` | |
| /// decaem geometricamente (mesma estrutura do Teo S-2' do SSD-Δ). | |
| /// | |
| /// **Prova R-2 (limite de esquecimento):** o passo de treino mistura | |
| /// gradientes `(1−α)·∇L_novo + α·∇L_replay`; para `L_antiga` μ-fortemente | |
| /// convexa com gradientes novos limitados por `G`, o afastamento do ótimo | |
| /// antigo por passo é `≤ (1−α)·η·G` (sem replay: `≤ η·G`) — o replay | |
| /// contrai o drift por um fator `1/(1−α)` no regime linear. | |
| pub struct ReplaySsm { | |
| pub a: Var, // (d,) — porta (sigmoid no uso) | |
| pub b: Var, // (d_in, d) — projeção de entrada | |
| pub d: usize, | |
| pub alpha_replay: f64, | |
| pub buffer: Vec<Tensor>, | |
| pub capacidade_buffer: usize, | |
| } | |
| impl ReplaySsm { | |
| pub fn novo(d_in: usize, d: usize, alpha_replay: f64, capacidade_buffer: usize, dev: &Device) -> Result<Self> { | |
| let a_t = Tensor::randn(0f32, 0.5, (d,), dev)?; | |
| let b_t = Tensor::randn(0f32, 0.1, (d_in, d), dev)?; | |
| Ok(Self { | |
| a: Var::from_tensor(&a_t)?, | |
| b: Var::from_tensor(&b_t)?, | |
| d, | |
| alpha_replay, | |
| buffer: Vec::new(), | |
| capacidade_buffer, | |
| }) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| vec![self.a.clone(), self.b.clone()] | |
| } | |
| /// Scan: x (T, d_in) → h (T, d). Diferenciável (candle autograd). | |
| pub fn scan(&self, x: &Tensor) -> Result<Tensor> { | |
| let t = x.dim(0)?; | |
| let dev = x.device(); | |
| let porta = candle_nn::ops::sigmoid(&self.a.as_tensor().broadcast_as((1, self.d))?)?; // (1, d) | |
| let um_menos = porta.affine(-1.0, 1.0)?; // 1 − σ(a) | |
| let entrada = x.matmul(&self.b.as_tensor())?; // (T, d) | |
| let mut h = Tensor::zeros((1, self.d), candle_core::DType::F32, &dev)?; | |
| let mut saidas: Vec<Tensor> = Vec::with_capacity(t); | |
| for i in 0..t { | |
| let xt = entrada.narrow(0, i, 1)?; | |
| h = ((&porta * &h)? + (&um_menos * &xt)?)?; | |
| saidas.push(h.clone()); | |
| } | |
| Ok(Tensor::stack(&saidas, 0)?) | |
| } | |
| /// Empurra uma janela de treino no buffer (FIFO). | |
| pub fn guardar(&mut self, x: &Tensor) { | |
| if self.buffer.len() >= self.capacidade_buffer { | |
| self.buffer.remove(0); | |
| } | |
| self.buffer.push(x.detach()); | |
| } | |
| /// Um passo de treino no lote novo (com mistura de replay). | |
| /// Devolve (perda_total, perda_novo, perda_replay). | |
| pub fn passo_treino( | |
| &mut self, | |
| x: &Tensor, | |
| alvo: &Tensor, | |
| otimizador: &mut crate::treino::AdamWManual, | |
| ) -> Result<(f64, f64, f64)> { | |
| // perda nova | |
| let h = self.scan(x)?; | |
| let dif = (&h - alvo)?; | |
| let perda_novo = dif.sqr()?.mean_all()?; | |
| // replay | |
| let mut perda_total = perda_novo.clone(); | |
| let mut p_r = 0.0; | |
| if !self.buffer.is_empty() { | |
| let idx = self.buffer.len() - 1; | |
| let antigo = self.buffer[idx].clone(); | |
| let hr = self.scan(&antigo)?; | |
| // alvo do replay: manter a dinâmica aprendida estável (magnitudes | |
| // finitas) — auto-supervisão de estabilidade | |
| let dr = hr.sqr()?.mean_all()?; | |
| perda_total = (perda_total.affine(1.0 - self.alpha_replay, 0.0)? | |
| + dr.affine(self.alpha_replay, 0.0)?)?; | |
| p_r = dr.to_scalar::<f32>()? as f64; | |
| } | |
| let pn = perda_novo.to_scalar::<f32>()? as f64; | |
| otimizador.passo(crate::treino::grads_achatados(&perda_total, &self.vars())?); | |
| self.guardar(x); | |
| let pt = perda_total.to_scalar::<f32>()? as f64; | |
| Ok((pt, pn, p_r)) | |
| } | |
| } | |
| // --------------------------------------------------------------------------- | |
| // 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_com_semente( | |
| d_patch, cfg.d_emb, d_cond, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, 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_com_semente(d_a, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, dev)?; | |
| let mut na = NucleoJepa::novo_com_semente(d_b, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, 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_com_semente( | |
| d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, dev, | |
| )?; | |
| // preditor anti-causal: núcleo separado | |
| let mut ninv = NucleoJepa::novo_com_semente( | |
| d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, 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" | |
| } | |
| mod testes { | |
| use super::*; | |
| fn tensor_2d(dados: &[Vec<f64>], dev: &Device) -> Tensor { | |
| let flat: Vec<f32> = dados.iter().flat_map(|v| v.iter().map(|&x| x as f32)).collect(); | |
| Tensor::from_vec(flat, (dados.len(), dados[0].len()), dev).unwrap() | |
| } | |
| /// T-J3 (falsificação da neutralidade): com UM único braço (frac fixa), | |
| /// o treino adaptativo deve ser BIT-A-BIT idêntico ao `treinar_variante` | |
| /// com a mesma fração — o RNG do bandit é fluxo separado e não toca o | |
| /// fluxo de máscara. Qualquer vazamento de aleatoriedade quebra a igualdade. | |
| fn jepa_adaptativa_um_braco_e_fixa() { | |
| let dev = Device::Cpu; | |
| let mut rng = RngDeterministico::novo(3); | |
| let dados: Vec<Vec<f64>> = (0..96) | |
| .map(|_| (0..12).map(|_| rng.normal(0.0, 1.0)).collect()) | |
| .collect(); | |
| let t = tensor_2d(&dados, &dev); | |
| let mut cfg = CfgJepa::default(); | |
| cfg.epocas = 12; | |
| cfg.d_emb = 16; | |
| cfg.frac_mascara = 0.35; | |
| let grid = [12usize, 8]; | |
| let var = VarianteJepa::ImageJepa; | |
| // fixo | |
| let rel_fixo = treinar_variante( | |
| &t, | |
| &|_n, r| mascara_por_variante(var, &grid, 0.35, r), | |
| &cfg, | |
| var, | |
| None, | |
| &dev, | |
| ) | |
| .unwrap(); | |
| // adaptativo com um braço = mesma fração | |
| let rel_adapt = treinar_variante_adaptativa( | |
| &t, | |
| &grid, | |
| &[0.35], | |
| &cfg, | |
| var, | |
| None, | |
| &dev, | |
| ) | |
| .unwrap(); | |
| assert_eq!( | |
| rel_fixo.perda_pred_final.to_bits(), | |
| rel_adapt.base.perda_pred_final.to_bits(), | |
| "perda preditiva divergiu: fixa {} vs adapt {}", | |
| rel_fixo.perda_pred_final, | |
| rel_adapt.base.perda_pred_final | |
| ); | |
| assert_eq!(rel_fixo.perda_final.to_bits(), rel_adapt.base.perda_final.to_bits()); | |
| assert_eq!(rel_adapt.jogadas_por_braco, vec![cfg.epocas]); | |
| // determinismo do adaptativo multi-braço | |
| let roda = || { | |
| treinar_variante_adaptativa(&t, &grid, &[0.2, 0.35, 0.5], &cfg, var, None, &dev) | |
| .unwrap() | |
| }; | |
| let a1 = roda(); | |
| let a2 = roda(); | |
| assert_eq!(a1.base.perda_pred_final.to_bits(), a2.base.perda_pred_final.to_bits()); | |
| assert_eq!(a1.escolhas.len(), a2.escolhas.len()); | |
| // multi-braço joga todos os braços ao menos uma vez (exploração TS) | |
| assert!(a1.jogadas_por_braco.iter().all(|&j| j > 0), "braços não explorados: {:?}", a1.jogadas_por_braco); | |
| } | |
| fn teste_mascaras_validas() { | |
| // fração respeitada ±, particionamento contexto/alvo não-vazio | |
| let mut rng = RngDeterministico::novo(1); | |
| let m2 = mascara_blocos_2d(8, 8, 4, 0.3, &mut rng); | |
| assert_eq!(m2.len(), 64); | |
| let frac = m2.iter().filter(|&&b| b).count() as f64 / 64.0; | |
| assert!(frac >= 0.15 && frac <= 0.55, "frac 2d = {frac}"); | |
| assert!(m2.iter().any(|&b| !b), "contexto vazio"); | |
| let m3 = mascara_tubos_3d(4, 6, 6, 3, 0.3, &mut rng); | |
| assert_eq!(m3.len(), 144); | |
| let f3 = m3.iter().filter(|&&b| b).count() as f64 / 144.0; | |
| assert!(f3 >= 0.15 && f3 <= 0.6, "frac 3d = {f3}"); | |
| let ms = mascara_segmentos(4, 32, 0.25, &mut rng); | |
| assert_eq!(ms.len(), 128); | |
| assert!(ms.iter().any(|&b| b) && ms.iter().any(|&b| !b)); | |
| let m1 = mascara_intervalos_1d(64, 0.3, &mut rng); | |
| let f1 = m1.iter().filter(|&&b| b).count() as f64 / 64.0; | |
| assert!(f1 >= 0.25 && f1 <= 0.45, "frac 1d = {f1}"); | |
| } | |
| fn teste_fps_cobre() { | |
| let mut rng = RngDeterministico::novo(2); | |
| let pontos: Vec<Vec<f64>> = (0..40) | |
| .map(|i| { | |
| let a = i as f64 * 0.1; | |
| vec![a.cos() * 3.0, a.sin() * 3.0, (i % 7) as f64 * 0.3] | |
| }) | |
| .collect(); | |
| let centros = pontos_mais_distantes(&pontos, 6, &mut rng); | |
| assert_eq!(centros.len(), 6); | |
| // centros distintos | |
| let mut s = centros.clone(); | |
| s.sort(); | |
| s.dedup(); | |
| assert_eq!(s.len(), 6); | |
| } | |
| fn teste_j1_colapso_impossivel() { | |
| // No colapso total (z constante), a perda VC = γ² > 0 — verificando a | |
| // fórmula implementada com um lote constante. | |
| let dev = Device::Cpu; | |
| let nucleo = NucleoJepa::novo(4, 8, 0, 4, 0.99, 1.0, 1.0, &dev).unwrap(); | |
| let z = Tensor::from_vec(vec![0.5f32; 8 * 8], (8, 8), &dev).unwrap(); | |
| let lvc = nucleo.perda_vc(&z).unwrap(); | |
| let v = lvc.to_scalar::<f32>().unwrap() as f64; | |
| // std=0 ⇒ dobradiça = γ = 1 ⇒ V/d = 1; cov = 0 | |
| assert!(v > 0.9, "perda VC no colapso = {v} (deveria ser ~1)"); | |
| } | |
| fn teste_j2_ema_ponto_fixo() { | |
| // encoder congelado ⇒ após muitas passadas EMA, alvo = encoder | |
| let dev = Device::Cpu; | |
| let mut nucleo = NucleoJepa::novo(4, 8, 0, 4, 0.5, 0.1, 1.0, &dev).unwrap(); | |
| for _ in 0..40 { | |
| nucleo.ema_passo().unwrap(); | |
| } | |
| let a = nucleo.alvo.vars()[0].as_tensor().flatten_all().unwrap().to_vec1::<f32>().unwrap(); | |
| let b = nucleo.encoder.vars()[0].as_tensor().flatten_all().unwrap().to_vec1::<f32>().unwrap(); | |
| let dif: f32 = a.iter().zip(&b).map(|(x, y)| (x - y).abs()).fold(0.0, f32::max); | |
| assert!(dif < 1e-5, "EMA não convergiu: dif máx = {dif}"); | |
| } | |
| fn teste_jepa_aprende_latente_previsivel() { | |
| // J-3: com alvo previsível (função determinística do contexto), a | |
| // perda pred must cair substancialmente e o std das incorporações | |
| // NÃO colapsa (J-1). | |
| let dev = Device::Cpu; | |
| let mut rng = RngDeterministico::novo(3); | |
| // patches = 64 vetores 8-D em 4 grupos (estrutura previsível) | |
| let mut dados: Vec<Vec<f64>> = Vec::new(); | |
| for g in 0..4 { | |
| let c = g as f64 * 2.0; | |
| for _ in 0..16 { | |
| dados.push((0..8).map(|k| c + rng.normal(0.0, 0.3) + (k as f64) * 0.05).collect()); | |
| } | |
| } | |
| let patches = tensor_2d(&dados, &dev); | |
| let cfg = CfgJepa { d_emb: 16, epocas: 40, lr: 5e-3, ..Default::default() }; | |
| let rel = treinar_variante( | |
| &patches, | |
| &|n, r| mascara_intervalos_1d(n, 0.3, r), | |
| &cfg, | |
| VarianteJepa::TJepa, | |
| None, | |
| &dev, | |
| ) | |
| .unwrap(); | |
| assert!( | |
| rel.perda_final < rel.perda_inicial * 0.8, | |
| "perda deveria cair: {} → {}", | |
| rel.perda_inicial, | |
| rel.perda_final | |
| ); | |
| assert!( | |
| rel.std_emb_final > 0.05, | |
| "sem colapso: std_final = {}", | |
| rel.std_emb_final | |
| ); | |
| } | |
| fn teste_adaptadores_formas() { | |
| let dev = Device::Cpu; | |
| let img = Tensor::randn(0f32, 1.0, (3, 16, 16), &dev).unwrap(); | |
| let p = em_patches_2d(&img, 4, 4).unwrap(); | |
| assert_eq!(p.dims().to_vec(), vec![16, 48]); | |
| let vid = Tensor::randn(0f32, 1.0, (4, 8, 8), &dev).unwrap(); | |
| let t3 = em_tubos_3d(&vid, 2, 4, 4).unwrap(); | |
| assert_eq!(t3.dims().to_vec(), vec![8, 32]); | |
| let spec = Tensor::randn(0f32, 1.0, (16, 32), &dev).unwrap(); | |
| let ps = em_patches_espectrograma(&spec, 4, 8).unwrap(); | |
| assert_eq!(ps.dims().to_vec(), vec![16, 32]); | |
| let ts = Tensor::randn(0f32, 1.0, (4, 32), &dev).unwrap(); | |
| let pt_ = em_patches_serie(&ts, 8).unwrap(); | |
| assert_eq!(pt_.dims().to_vec(), vec![16, 8]); | |
| let vol = Tensor::randn(0f32, 1.0, (8, 8, 8), &dev).unwrap(); | |
| let vx = em_voxels(&vol, 4, 4, 4).unwrap(); | |
| assert_eq!(vx.dims().to_vec(), vec![8, 64]); | |
| let pc = Tensor::randn(0f32, 1.0, (50, 3), &dev).unwrap(); | |
| let mut rng = RngDeterministico::novo(5); | |
| let gr = em_grupos_pontos(&pc, 6, 4, &mut rng).unwrap(); | |
| assert_eq!(gr.dims().to_vec(), vec![6, 12]); | |
| } | |
| fn teste_sonic_moe_capacidade_e_balanceamento() { | |
| // SM-1: carga por especialista ≤ capacidade; SM-2: aux ≥ ~1 quando | |
| // roteador é ~uniforme (início). | |
| let dev = Device::Cpu; | |
| let moe = SonicMoe::novo(8, 4, 2, 6, &dev).unwrap(); | |
| let x = Tensor::randn(0f32, 1.0, (20, 8), &dev).unwrap(); | |
| let s = moe.fwd(&x).unwrap(); | |
| for (e, c) in s.carga.iter().enumerate() { | |
| assert!(*c <= 6, "especialista {e} com carga {c} > capacidade 6"); | |
| } | |
| assert_eq!(s.saida.dims().to_vec(), vec![20, 8]); | |
| let aux = s.perda_aux.to_scalar::<f32>().unwrap() as f64; | |
| // com roteador quase uniforme, aux ≈ E·Σ (1/E)·(1/E) = 1 | |
| assert!(aux > 0.5 && aux < 3.0, "aux = {aux}"); | |
| // cada token processado por ≤ n_top especialistas | |
| let total_carga: usize = s.carga.iter().sum(); | |
| assert!(total_carga <= 20 * 2, "carga total {total_carga} > n·topk"); | |
| } | |
| fn teste_replay_ssm_estabilidade_e_replay() { | |
| // R-1: estados finitos e limitados; R-2: com replay α>0, após treinar | |
| // a tarefa nova, a dinâmica antiga (magnitude do estado) permanece | |
| // estável (não diverge). | |
| let dev = Device::Cpu; | |
| let mut ssm = ReplaySsm::novo(4, 6, 0.5, 8, &dev).unwrap(); | |
| let vars = ssm.vars(); | |
| let mut otim = crate::treino::AdamWManual::novo(vars, 1e-2, 0.0); | |
| // tarefa "antiga": sinal constante 1.0 | |
| let x_antigo = Tensor::from_vec(vec![1.0f32; 16 * 4], (16, 4), &dev).unwrap(); | |
| for _ in 0..30 { | |
| let h = ssm.scan(&x_antigo).unwrap(); | |
| let alvo = h.detach(); | |
| ssm.passo_treino(&x_antigo, &alvo, &mut otim).unwrap(); | |
| } | |
| // estado após treino antigo: finito | |
| let h1 = ssm.scan(&x_antigo).unwrap(); | |
| let m1 = h1.abs().unwrap().max_all().unwrap().to_scalar::<f32>().unwrap(); | |
| assert!(m1.is_finite() && m1 < 50.0, "estado explodiu: {m1}"); | |
| // tarefa "nova": sinal −1.0 — treino com replay | |
| let x_novo = Tensor::from_vec(vec![-1.0f32; 16 * 4], (16, 4), &dev).unwrap(); | |
| for _ in 0..30 { | |
| let h = ssm.scan(&x_novo).unwrap(); | |
| let alvo = h.detach(); | |
| ssm.passo_treino(&x_novo, &alvo, &mut otim).unwrap(); | |
| } | |
| // a dinâmica antiga continua estável (sem explosão) | |
| let h2 = ssm.scan(&x_antigo).unwrap(); | |
| let m2 = h2.abs().unwrap().max_all().unwrap().to_scalar::<f32>().unwrap(); | |
| assert!(m2.is_finite() && m2 < 50.0, "estado antigo explodiu: {m2}"); | |
| // portas dentro de (0,1) | |
| let a = ssm.a.as_tensor().to_vec1::<f32>().unwrap(); | |
| for v in a { | |
| assert!(v.is_finite()); | |
| } | |
| } | |
| } | |