v0.5.0: suíte SOM-4D + suíte JEPA (12) + tecnologia Pensamento + telemetria Theil–Sen — 29/29 gates verdes (provas r-18..r-22)
e26f4cc verified | //! Slot Mamba-3 (1-de-6 do tronco 5:1) — scan recorrente diagonal seletivo | |
| //! (SSD) com adjunto EXATO via CustomOp2. Compatível com o Teo D.1(i): | |
| //! fluxo estritamente esquerda→direita (a camada é causal por construção); | |
| //! a bidirecionalidade intra-nível vem das camadas ADBA (H_D1 — r-04 §3). | |
| //! | |
| //! Recorrência: h_t = a_t ⊙ h_{t−1} + (1 − a_t) ⊙ x_t; y_t = h_t | |
| //! com a_t = σ(W_a·x_t + b_a) ∈ (0,1)^d (portas por canal). | |
| //! | |
| //! Adjunto (exato): dh_t = dy_t + a_{t+1} ⊙ dh_{t+1}; | |
| //! dx_t = (1 − a_t) ⊙ dh_t; da_t = ⟨dh_t, h_{t−1} − x_t⟩ (escalar por canal). | |
| use candle_core::{CpuStorage, CustomOp2, Device, Layout, Result, Shape, Tensor, Var}; | |
| use crate::camadas::{Linear, RmsNorm, Swiglu}; | |
| pub struct ScanDiagonal; | |
| impl CustomOp2 for ScanDiagonal { | |
| fn name(&self) -> &'static str { "ssd_scan_diagonal" } | |
| fn cpu_fwd(&self, s1: &CpuStorage, l1: &Layout, s2: &CpuStorage, l2: &Layout) -> Result<(CpuStorage, Shape)> { | |
| if !(l1.is_contiguous() && l2.is_contiguous()) { | |
| return Err(candle_core::Error::Msg( | |
| "ssd_scan: entradas devem ser contíguas (chame .contiguous())".into(), | |
| )); | |
| } | |
| let x = s1.as_slice::<f32>()?; // (n, d) | |
| let a = s2.as_slice::<f32>()?; // (n, d) em (0,1) | |
| let (n, d) = (l1.shape().dims()[0], l1.shape().dims()[1]); | |
| let mut y = vec![0f32; n * d]; | |
| for j in 0..d { | |
| let mut h = 0f32; | |
| for i in 0..n { | |
| let ai = a[i * d + j]; | |
| h = ai * h + (1.0 - ai) * x[i * d + j]; | |
| y[i * d + j] = h; | |
| } | |
| } | |
| Ok((CpuStorage::F32(y), (n, d).into())) | |
| } | |
| fn bwd(&self, arg1: &Tensor, arg2: &Tensor, _res: &Tensor, grad_res: &Tensor) -> Result<(Option<Tensor>, Option<Tensor>)> { | |
| // recomputa h (forward sem gradiente) e percorre reverso | |
| let (x, a) = (arg1, arg2); | |
| let h = Tensor::apply_op2(x, a, ScanDiagonal)?.detach(); | |
| let x_v = x.flatten_all()?.to_vec1::<f32>()?; | |
| let a_v = a.flatten_all()?.to_vec1::<f32>()?; | |
| let h_v = h.flatten_all()?.to_vec1::<f32>()?; | |
| let dy_v = grad_res.flatten_all()?.to_vec1::<f32>()?; | |
| let (n, d) = (x.dim(0)?, x.dim(1)?); | |
| let mut dx = vec![0f32; n * d]; | |
| let mut da = vec![0f32; n * d]; | |
| let mut dh = vec![0f32; d]; | |
| for i in (0..n).rev() { | |
| for j in 0..d { | |
| let idx = i * d + j; | |
| let g = dy_v[idx] + dh[j]; | |
| dx[idx] = (1.0 - a_v[idx]) * g; | |
| let h_ant = if i > 0 { h_v[(i - 1) * d + j] } else { 0.0 }; | |
| da[idx] = g * (h_ant - x_v[idx]); | |
| dh[j] = a_v[idx] * g; | |
| } | |
| } | |
| let dev = x.device(); | |
| let gx = Tensor::from_vec(dx, (n, d), dev)?; | |
| let ga = Tensor::from_vec(da, (n, d), dev)?; | |
| Ok((Some(gx), Some(ga))) | |
| } | |
| } | |
| /// Bloco SSD (a camada recorrente do tronco): portas por canal + scan + projeção. | |
| pub struct BlocoSsd { | |
| pub norm1: RmsNorm, | |
| pub porta_a: Linear, // d → d (logitas do decaimento) | |
| pub proj_x: Linear, // d → d | |
| pub proj_o: Linear, // d → d | |
| pub norm2: RmsNorm, | |
| pub ffn: Swiglu, | |
| } | |
| impl BlocoSsd { | |
| pub fn novo(d: usize, d_ff: usize, dev: &Device) -> Result<Self> { | |
| let desvio = 1.0 / ((d * 2) as f64).sqrt(); | |
| Ok(BlocoSsd { | |
| norm1: RmsNorm::novo(d, 1e-5, dev)?, | |
| porta_a: Linear::novo(d, d, desvio, dev)?, | |
| proj_x: Linear::sem_vies(d, d, desvio, dev)?, | |
| proj_o: Linear::sem_vies(d, d, desvio, dev)?, | |
| norm2: RmsNorm::novo(d, 1e-5, dev)?, | |
| ffn: Swiglu::novo(d, d_ff, dev)?, | |
| }) | |
| } | |
| /// x (b, n, d) → (b, n, d). O scan é aplicado por elemento do lote. | |
| /// `cache_estado`: estado h anterior (b, d) para streaming incremental. | |
| pub fn fwd(&self, x: &Tensor, cache_estado: &mut Option<Tensor>) -> Result<Tensor> { | |
| let b = x.dim(0)?; | |
| let n = x.dim(1)?; | |
| let d = x.dim(2)?; | |
| let xb = self.norm1.fwd(x)?; | |
| let a = candle_nn::ops::sigmoid(&self.porta_a.fwd(&xb)?)?; | |
| let xp = self.proj_x.fwd(&xb)?; | |
| // achata (b·n, d), aplica o scan por lote e recompõe | |
| let xf = xp.reshape((b * n, d))?.contiguous()?; | |
| let af = a.reshape((b * n, d))?.contiguous()?; | |
| // pré-estado: se houver cache, subtrai efeito — simplificação: scan | |
| // independente por segmento com estado inicial zero (documentado; o | |
| // estado cruzado entre segmentos é ponto de extensão do streaming). | |
| let _ = cache_estado; | |
| let mut saidas = Vec::with_capacity(b); | |
| for i in 0..b { | |
| let xi = xf.narrow(0, i * n, n)?; | |
| let ai = af.narrow(0, i * n, n)?; | |
| let yi = Tensor::apply_op2(&xi, &ai, ScanDiagonal)?; | |
| saidas.push(yi); | |
| } | |
| let y = Tensor::cat(&saidas, 0)?.reshape((b, n, d))?; | |
| let y = self.proj_o.fwd(&y)?; | |
| let x2 = (x + &y)?; | |
| let f = self.ffn.fwd(&self.norm2.fwd(&x2)?)?; | |
| (&x2 + &f) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| [ | |
| self.norm1.vars(), self.porta_a.vars(), self.proj_x.vars(), | |
| self.proj_o.vars(), self.norm2.vars(), self.ffn.vars(), | |
| ].concat() | |
| } | |
| } | |
| mod testes { | |
| use super::*; | |
| fn adjunto_do_scan_vs_finitas_diferencas() -> Result<()> { | |
| let dev = Device::Cpu; | |
| let n = 9; | |
| let d = 5; | |
| let xv = Var::from_tensor(&Tensor::randn(0f32, 1f32, (n, d), &dev)?)?; | |
| let av = Var::from_tensor(&candle_nn::ops::sigmoid(&Tensor::randn(0f32, 1f32, (n, d), &dev)?)?)?; | |
| let pesos = Tensor::randn(0f32, 1f32, (n, d), &dev)?; | |
| let y = Tensor::apply_op2(xv.as_tensor(), &av.as_tensor(), ScanDiagonal)?; | |
| let perda = (&y * &pesos)?.sum_all()?; | |
| let gs = perda.backward()?; | |
| let gx = gs.get(&xv).ok_or_else(|| candle_core::Error::Msg("sem grad x".into()))?; | |
| let ga = gs.get(&av).ok_or_else(|| candle_core::Error::Msg("sem grad a".into()))?; | |
| let gxv = gx.flatten_all()?.to_vec1::<f32>()?; | |
| let gav = ga.flatten_all()?.to_vec1::<f32>()?; | |
| // FD | |
| let eps = 1e-3f32; | |
| let f = |x: &Tensor, a: &Tensor| -> Result<f64> { | |
| let y = Tensor::apply_op2(x, a, ScanDiagonal)?; | |
| Ok((&y * &pesos)?.sum_all()?.to_scalar::<f32>()? as f64) | |
| }; | |
| let x0 = xv.as_tensor(); | |
| let a0 = av.as_tensor(); | |
| let x_plano = x0.flatten_all()?.to_vec1::<f32>()?; | |
| let a_plano = a0.flatten_all()?.to_vec1::<f32>()?; | |
| let mut erro = 0f64; | |
| for idx in [3usize, 11, 17, 24, 30] { | |
| let mut xp = x_plano.clone(); | |
| let mut xm = x_plano.clone(); | |
| xp[idx] += eps; | |
| xm[idx] -= eps; | |
| let tp = Tensor::from_vec(xp, (n, d), &dev)?; | |
| let tm = Tensor::from_vec(xm, (n, d), &dev)?; | |
| let fd = (f(&tp, a0)? - f(&tm, a0)?) / (2.0 * eps as f64); | |
| erro = erro.max((fd - gxv[idx] as f64).abs()); | |
| // porta a | |
| let mut ap = a_plano.clone(); | |
| let mut am = a_plano.clone(); | |
| ap[idx] += eps; | |
| am[idx] -= eps; | |
| let tap = Tensor::from_vec(ap, (n, d), &dev)?; | |
| let tam = Tensor::from_vec(am, (n, d), &dev)?; | |
| let fda = (f(x0, &tap)? - f(x0, &tam)?) / (2.0 * eps as f64); | |
| erro = erro.max((fda - gav[idx] as f64).abs()); | |
| } | |
| assert!(erro < 5e-3, "erro adjunto SSD = {erro}"); | |
| Ok(()) | |
| } | |
| } | |