v0.8.0: análise/prova/teste das 4 propostas Qwen3.8 (r-33) — Newton analítico no verificador (convexidade em gamma=1/T), árvore pós-ordem iterativa bit-a-bit com detecção de ciclos, Cayley exato no FNH (4,1× mais rápido), MoE vetorizada com gradiente RESTAURADO (bug v080-1) — 85/85 gates, 99/99 cargo
28430a8 verified | //! Quantizador GSQ por nível (Teos A.1/A.4/D.4) — o coração da ponte | |
| //! NCM×KHTST: codebook Stiefel com seleção Gumbel-Softmax (gradiente exato, | |
| //! sem STE) no treino e NN duro (argmax) em eval. | |
| //! | |
| //! Dois modos (r-01 §4.4): "gsq_rco" (padrão, Teo D.4 — relaxamento | |
| //! diferenciável com τ-annealing) e "ste_ema" (Teo A.0 — STE + commitment + | |
| //! codebook por EMA). A escala R_ℓ é mantida por EMA (Cor. A.1.1). | |
| use candle_core::{CpuStorage, CustomOp3, Device, Layout, Result, Shape, Tensor, Var, D}; | |
| use crate::stiefel; | |
| pub enum ModoQuant { | |
| GsqRco, | |
| SteEma, | |
| } | |
| impl ModoQuant { | |
| pub fn de_str(s: &str) -> ModoQuant { | |
| if s == "ste_ema" { ModoQuant::SteEma } else { ModoQuant::GsqRco } | |
| } | |
| } | |
| /// CustomOp3: (z, C, ruído_gumbel) → r_suave = R·softmax((zCᵀ/τ + g))·C. | |
| /// O ruído entra como tensor EXPLÍCITO (constante no grafo) para que o | |
| /// backward recompute exatamente o mesmo y do forward (gradiente exato do | |
| /// relaxamento — Teo A.4(iii), sem re-amostragem no adjunto). | |
| pub struct GsqSuave { | |
| pub r: f64, | |
| pub tau: f64, | |
| } | |
| fn softmax_ultima_dim(v: &[f32], n: usize, k: usize) -> Vec<f32> { | |
| let mut out = vec![0f32; v.len()]; | |
| for i in 0..n { | |
| let linha = &v[i * k..(i + 1) * k]; | |
| let mx = linha.iter().cloned().fold(f32::NEG_INFINITY, f32::max); | |
| let mut soma = 0f32; | |
| for (j, &x) in linha.iter().enumerate() { | |
| let e = (x - mx).exp(); | |
| out[i * k + j] = e; | |
| soma += e; | |
| } | |
| for j in 0..k { | |
| out[i * k + j] /= soma.max(1e-20); | |
| } | |
| } | |
| out | |
| } | |
| impl CustomOp3 for GsqSuave { | |
| fn name(&self) -> &'static str { "gsq_suave_stiefel" } | |
| fn cpu_fwd( | |
| &self, | |
| s1: &CpuStorage, l1: &Layout, | |
| _s2: &CpuStorage, _l2: &Layout, | |
| s3: &CpuStorage, l3: &Layout, | |
| ) -> Result<(CpuStorage, Shape)> { | |
| if !(l1.is_contiguous() && _l2.is_contiguous() && l3.is_contiguous()) { | |
| return Err(candle_core::Error::Msg( | |
| "gsq_suave: tensors de entrada devem ser contíguos (chame .contiguous())".into(), | |
| )); | |
| } | |
| let z = s1.as_slice::<f32>()?; | |
| let (n, d) = (l1.shape().dims()[0], l1.shape().dims()[1]); | |
| let c = _s2.as_slice::<f32>()?; | |
| let (k, _dc) = (_l2.shape().dims()[0], _l2.shape().dims()[1]); | |
| let ruido = s3.as_slice::<f32>()?; | |
| // v0.8.0 (bug v080-2, pré-existente): a fórmula anterior | |
| // `z.len()/d*k*d/d*k` avaliava para n·k²·(…) sem sentido (320 ≠ 40 | |
| // no teste) e disparava em TODO build debug — os gates da v0.7.0 só | |
| // passavam em release (debug_assert compilado fora). Correta: n·k. | |
| debug_assert_eq!(ruido.len(), n * k); // (n, k) | |
| // logits = z·Cᵀ/τ + g | |
| let mut logits = vec![0f32; n * k]; | |
| for i in 0..n { | |
| for j in 0..k { | |
| let mut acc = 0f32; | |
| for t in 0..d { | |
| acc += z[i * d + t] * c[j * d + t]; | |
| } | |
| logits[i * k + j] = acc as f32 / self.tau as f32 + ruido[i * k + j]; | |
| } | |
| } | |
| let y = softmax_ultima_dim(&logits, n, k); | |
| // r = R · y·C | |
| let mut r = vec![0f32; n * d]; | |
| for i in 0..n { | |
| for t in 0..d { | |
| let mut acc = 0f32; | |
| for j in 0..k { | |
| acc += y[i * k + j] * c[j * d + t]; | |
| } | |
| r[i * d + t] = (self.r as f32) * acc; | |
| } | |
| } | |
| Ok((CpuStorage::F32(r), (n, d).into())) | |
| } | |
| fn bwd( | |
| &self, | |
| arg1: &Tensor, arg2: &Tensor, arg3: &Tensor, | |
| _res: &Tensor, grad_res: &Tensor, | |
| ) -> Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> { | |
| // recompute y (mesmo ruído — arg3): y = softmax(zCᵀ/τ + g) | |
| let logits = (&arg1.matmul(&arg2.transpose(0, 1)?)?.affine(1.0 / self.tau, 0.0)? + arg3)?; | |
| let y = candle_nn::ops::softmax_last_dim(&logits)?; | |
| // gy = R·(g_·Cᵀ) | |
| let gy = (&grad_res.matmul(&arg2.transpose(0, 1)?)? * self.r)?; | |
| // gl = y ⊙ (gy − Σ(y⊙gy)) — media como coluna p/ broadcast (n,1) | |
| let media = (&y * &gy)?.sum(D::Minus1)?; | |
| let media_c = media.reshape((media.dims()[0], 1))?; | |
| let gl = (&y * &gy.broadcast_sub(&media_c.broadcast_div(&ones_like(&media_c)?)?)?)?; | |
| // ∂L/∂z = (1/τ)·gl·C ; ∂L/∂C = R·yᵀ·g_ ; ruído constante | |
| let gz = gl.matmul(arg2)?.affine(1.0 / self.tau, 0.0)?; | |
| let gc = (&y.transpose(0, 1)?.matmul(grad_res)? * self.r)?; | |
| Ok((Some(gz), Some(gc), None)) | |
| } | |
| } | |
| fn ones_like(t: &Tensor) -> Result<Tensor> { | |
| Tensor::ones(t.shape(), t.dtype(), t.device()) | |
| } | |
| /// Quantizador de um nível ℓ do VQ residual. | |
| pub struct QuantizadorNivel { | |
| pub w: Var, // (K, d) bruto — codebook = R(Wᵀ)ᵀ | |
| pub r_escala: f64, // R_ℓ (EMA) | |
| pub tau: f64, // τ corrente (agenda analítica) | |
| pub modo: ModoQuant, | |
| pub ema: f64, // fator EMA (0,99) | |
| pub escala_modo: String, // "ema_projecao" | "otima_teo" | |
| pub iters_polar: usize, | |
| pub delta_ort_ultimo: f64, | |
| pub uso: Vec<f64>, // contagem de uso por código (telemetria/colapso) | |
| pub semente: u64, // RNG dedicado do nível (pareamento — Teo 31.7) | |
| chamadas: std::cell::Cell<u64>, // contador p/ ruído fresco por chamada | |
| } | |
| impl QuantizadorNivel { | |
| pub fn novo(k: usize, d: usize, modo: ModoQuant, ema: f64, escala_modo: &str, dev: &Device) -> Result<Self> { | |
| let w0 = stiefel::init_stiefel(k, d, dev)?; | |
| Ok(QuantizadorNivel { | |
| w: Var::from_tensor(&w0)?, | |
| r_escala: 1.0 / (k as f64).sqrt(), | |
| tau: 1.0, | |
| modo, | |
| ema, | |
| escala_modo: escala_modo.to_string(), | |
| iters_polar: 30, | |
| delta_ort_ultimo: 0.0, | |
| uso: vec![0.0; k], | |
| semente: 0xC0FFEE + k as u64 * 7919 + d as u64, | |
| chamadas: std::cell::Cell::new(0), | |
| }) | |
| } | |
| /// Codebook C = R(Wᵀ)ᵀ com linhas ortonormais (diferenciável — Teo 26.3). | |
| /// SEMPRE contíguo (index_select/arquivo exigem layout padrão). | |
| pub fn codebook(&self) -> Result<Tensor> { | |
| let c = stiefel::codebook_stiefel(self.w.as_tensor(), self.iters_polar)?; | |
| c.contiguous() | |
| } | |
| /// δ_ort do codebook corrente (V9). | |
| pub fn delta_ort(&self) -> Result<f64> { | |
| let c = self.codebook()?.detach(); | |
| stiefel::erro_ortogonalidade(&c.transpose(0, 1)?) | |
| } | |
| /// Amostra ruído de Gumbel (n, K) — fora do grafo (constante). | |
| pub fn ruido_gumbel(&self, n: usize, k: usize, dev: &Device) -> Result<Tensor> { | |
| use rand::prelude::*; | |
| let mut rng = StdRng::seed_from_u64(self.semente_ruido()); | |
| let mut v = Vec::with_capacity(n * k); | |
| for _ in 0..n * k { | |
| let u: f32 = rng.gen::<f32>().clamp(1e-9, 1.0 - 1e-9); | |
| v.push((-(-u.ln()).ln()) as f32); | |
| } | |
| Tensor::from_vec(v, (n, k), dev) | |
| } | |
| fn semente_ruido(&self) -> u64 { | |
| let c = self.chamadas.get(); | |
| self.chamadas.set(c + 1); | |
| self.semente.wrapping_add(c.wrapping_mul(0x9E3779B97F4A7C15)) | |
| } | |
| /// Quantização de TREINO (…, d) → (…, d): | |
| /// - gsq_rco: relaxamento suave diferenciável (gradiente exato); | |
| /// - ste_ema: STE (r = z + sg[r_duro − z]). | |
| pub fn quantizar_treino(&self, z: &Tensor) -> Result<Tensor> { | |
| let c = self.codebook()?.contiguous()?; | |
| match self.modo { | |
| ModoQuant::GsqRco => { | |
| let z2 = contiguo_2d(z)?; | |
| let (n, _d) = (z2.dims()[0], z2.dims()[1]); | |
| let k = c.dims()[0]; | |
| let ruido = self.ruido_gumbel(n, k, z.device())?; | |
| let r = Tensor::apply_op3(&z2, &c, &ruido, GsqSuave { r: self.r_escala, tau: self.tau.max(1e-3) })?; | |
| redimensionar_de_volta(&r, z) | |
| } | |
| ModoQuant::SteEma => { | |
| let r_duro = self.valores_duros_int(z, &self.codebook()?.contiguous()?)?.detach(); | |
| // STE: z + (r_duro − z) com r_duro constante no grafo | |
| Ok(z.add(&r_duro.sub(z)?)?) | |
| } | |
| } | |
| } | |
| fn valores_duros_int(&self, z: &Tensor, c: &Tensor) -> Result<Tensor> { | |
| // c deve vir contíguo do caller (custom op exige layout padrão) | |
| // ids = argmax ⟨z, u_k⟩ (Teo A.4(i): uma GEMM + top-1) | |
| let z2 = contiguo_2d(z)?; | |
| let logits = z2.matmul(&c.transpose(0, 1)?)?; | |
| let ids = logits.argmax(D::Minus1)?; | |
| let linhas = c.index_select(&ids, 0)?; | |
| let r = linhas.affine(self.r_escala, 0.0)?; | |
| redimensionar_de_volta(&r, z) | |
| } | |
| /// Quantização DURA (eval/telemetria): (ids, valores R·c_id). | |
| pub fn quantizar_duro(&self, z: &Tensor) -> Result<(Tensor, Tensor)> { | |
| let c = self.codebook()?.detach(); | |
| let z2 = contiguo_2d(z)?; | |
| let logits = z2.matmul(&c.transpose(0, 1)?)?; | |
| let ids = logits.argmax(D::Minus1)?; | |
| let valores = self.valores_duros_int(z, &c)?; | |
| Ok((ids, valores)) | |
| } | |
| /// Materializa códigos ids → vetores R·c_id (para reconstrução na amostragem). | |
| pub fn materializar(&self, ids: &Tensor) -> Result<Tensor> { | |
| let c = self.codebook()?.detach(); | |
| let linhas = c.index_select(&ids, 0)?; | |
| Ok(linhas.affine(self.r_escala, 0.0)?) | |
| } | |
| /// Atualiza R_ℓ por EMA da escala de projeção (Cor. A.1.1): | |
| /// "ema_projecao": R ← EMA(E[‖P_C z‖]); "otima_teo": R ← EMA(E[‖P_C z‖/√K]). | |
| pub fn atualizar_escala(&mut self, z: &Tensor) -> Result<()> { | |
| let c = self.codebook()?.detach(); | |
| let z2 = contiguo_2d(z)?; | |
| let cz = z2.matmul(&c.transpose(0, 1)?)?; | |
| let pz = cz.matmul(&c)?; | |
| let alvo = pz.sqr()?.sum(D::Minus1)?.sqrt()?.mean_all()?.to_scalar::<f32>()? as f64; | |
| let alvo = match self.escala_modo.as_str() { | |
| "otima_teo" => alvo / (self.uso.len() as f64).sqrt(), | |
| _ => alvo, | |
| }; | |
| self.r_escala = self.ema * self.r_escala + (1.0 - self.ema) * alvo.max(1e-4); | |
| Ok(()) | |
| } | |
| /// Registra uso de códigos (telemetria de colapso — R2 do n-04 §9). | |
| pub fn registrar_uso(&mut self, ids: &Tensor) -> Result<()> { | |
| let v = ids.flatten_all()?.to_vec1::<u32>()?; | |
| for id in v { | |
| if (id as usize) < self.uso.len() { | |
| self.uso[id as usize] += 1.0; | |
| } | |
| } | |
| Ok(()) | |
| } | |
| /// Perplexidade de uso do codebook (utilização efetiva de K — telemetria V2/V6). | |
| pub fn perplexidade_uso(&self) -> f64 { | |
| let total: f64 = self.uso.iter().sum(); | |
| if total <= 0.0 { | |
| return 0.0; | |
| } | |
| let h: f64 = self.uso.iter().map(|&u| { | |
| let p = u / total; | |
| if p > 0.0 { -p * p.ln() } else { 0.0 } | |
| }).sum(); | |
| h.exp() | |
| } | |
| /// Agenda analítica de τ (Teo 30.6): decaimento de τ0 até τ_min em t_alvo passos. | |
| pub fn agendar_tau(&mut self, passo: usize, tau0: f64, tau_min: f64, t_alvo: usize) { | |
| let t = passo.min(t_alvo) as f64; | |
| let tau = tau_min + (tau0 - tau_min) * (1.0 - t / t_alvo.max(1) as f64); | |
| self.tau = tau.max(tau_min); | |
| } | |
| } | |
| // ---------- utilidades de forma ---------- | |
| /// Achata (…, d) → (n, d) e devolve também a forma original. | |
| fn contiguo_2d(t: &Tensor) -> Result<Tensor> { | |
| let d = t.dim(D::Minus1)?; | |
| let n = t.elem_count() / d; | |
| t.reshape((n, d))?.contiguous() | |
| } | |
| fn redimensionar_de_volta(r: &Tensor, modelo: &Tensor) -> Result<Tensor> { | |
| r.reshape(modelo.shape()) | |
| } | |
| mod testes { | |
| use super::*; | |
| fn gsq_gradiente_vs_finitas_diferencas() -> Result<()> { | |
| // V-análogo do V8: gradiente exato do relaxamento vs. FD | |
| let dev = Device::Cpu; | |
| let mut q = QuantizadorNivel::novo(8, 6, ModoQuant::GsqRco, 0.99, "ema_projecao", &dev)?; | |
| q.tau = 1.0; | |
| q.r_escala = 0.7; | |
| let z = Tensor::randn(0f32, 1f32, (5, 6), &dev)?; | |
| let z2 = contiguo_2d(&z)?; | |
| let c = q.codebook()?.contiguous()?; | |
| let ruido = q.ruido_gumbel(5, 8, &dev)?; | |
| // valor da perda com pesos fixos | |
| let pesos = Tensor::randn(0f32, 1f32, (5, 6), &dev)?; | |
| let perda = |zz: &Tensor| -> Result<f64> { | |
| let r = Tensor::apply_op3(zz, &c, &ruido, GsqSuave { r: q.r_escala, tau: q.tau })?; | |
| let l = (&r * &pesos)?.sum_all()?; | |
| Ok(l.to_scalar::<f32>()? as f64) | |
| }; | |
| // gradiente via op (aplicado a Var) | |
| let zv = Var::from_tensor(&z2)?; | |
| let r = Tensor::apply_op3(zv.as_tensor(), &c, &ruido, GsqSuave { r: q.r_escala, tau: q.tau })?; | |
| let l = (&r * &pesos)?.sum_all()?; | |
| let gs = l.backward()?; | |
| let g = gs.get(&zv).ok_or_else(|| candle_core::Error::Msg("sem gradiente".into()))?; | |
| let gv = g.flatten_all()?.to_vec1::<f32>()?; | |
| // FD central | |
| let eps = 1e-3f32; | |
| let mut erro_max = 0f64; | |
| for i in 0..z2.elem_count() { | |
| let mut zp = z2.flatten_all()?.to_vec1::<f32>()?; | |
| let mut zm = zp.clone(); | |
| zp[i] += eps; | |
| zm[i] -= eps; | |
| let tp = Tensor::from_vec(zp, (5, 6), &dev)?; | |
| let tm = Tensor::from_vec(zm, (5, 6), &dev)?; | |
| let fd = (perda(&tp)? - perda(&tm)?) / (2.0 * eps as f64); | |
| erro_max = erro_max.max((fd - gv[i] as f64).abs()); | |
| } | |
| assert!(erro_max < 5e-3, "erro máx FD vs exato = {erro_max}"); | |
| Ok(()) | |
| } | |
| fn perplexidade_uso_reflete_histograma() { | |
| let mut q = QuantizadorNivel { | |
| w: Var::from_tensor(&Tensor::zeros((4, 4), candle_core::DType::F32, &Device::Cpu).unwrap()).unwrap(), | |
| r_escala: 1.0, tau: 1.0, modo: ModoQuant::GsqRco, ema: 0.99, | |
| escala_modo: "ema_projecao".into(), iters_polar: 10, delta_ort_ultimo: 0.0, | |
| uso: vec![10.0, 10.0, 10.0, 10.0], | |
| semente: 7, | |
| chamadas: std::cell::Cell::new(0), | |
| }; | |
| let p = q.perplexidade_uso(); | |
| assert!((p - 4.0).abs() < 1e-9); | |
| q.uso = vec![40.0, 0.0, 0.0, 0.0]; | |
| assert!(q.perplexidade_uso() < 1.01); | |
| } | |
| } | |
| mod testes_debug { | |
| use super::*; | |
| use candle_core::{Device, Var}; | |
| /// Compara o adjunto do GsqSuave contra implementação nativa (ops candle). | |
| fn gsq_bwd_vs_nativo() -> Result<()> { | |
| let dev = Device::Cpu; | |
| let (n, d, k) = (5usize, 6, 8); | |
| let c0 = stiefel::init_stiefel(k, d, &dev)?.contiguous()?; | |
| let zv = Var::from_tensor(&Tensor::randn(0f32, 1f32, (n, d), &dev)?)?; | |
| let ruido = Tensor::randn(0f32, 1f32, (n, k), &dev)?; | |
| let pesos = Tensor::randn(0f32, 1f32, (n, d), &dev)?; | |
| let (r_op, tau) = (0.7f64, 1.3f64); | |
| // caminho NATIVO | |
| let logits_n = (zv.as_tensor().matmul(&c0.transpose(0, 1)?)?.affine(1.0 / tau, 0.0)? + &ruido)?; | |
| let y_n = candle_nn::ops::softmax(&logits_n, candle_core::D::Minus1)?; | |
| let r_n = y_n.matmul(&c0)?.affine(r_op, 0.0)?; | |
| let l_n = (&r_n * &pesos)?.sum_all()?; | |
| let gs_n = l_n.backward()?; | |
| let g_n = gs_n.get(&zv).unwrap().flatten_all()?.to_vec1::<f32>()?; | |
| // caminho OP | |
| let r_o = Tensor::apply_op3(zv.as_tensor(), &c0, &ruido, GsqSuave { r: r_op, tau })?; | |
| let l_o = (&r_o * &pesos)?.sum_all()?; | |
| let gs_o = l_o.backward()?; | |
| let g_o = gs_o.get(&zv).unwrap().flatten_all()?.to_vec1::<f32>()?; | |
| // forward deve ser idêntico | |
| let dif_fwd = (&r_n - &r_o)?.abs()?.max_all()?.to_scalar::<f32>()?; | |
| assert!(dif_fwd < 1e-5, "forward difere: {dif_fwd}"); | |
| let mut erro = 0f64; | |
| for i in 0..g_n.len() { | |
| erro = erro.max((g_n[i] - g_o[i]).abs() as f64); | |
| } | |
| assert!(erro < 1e-4, "grad nativo vs op = {erro}\n n={g_n:?}\n o={g_o:?}"); | |
| Ok(()) | |
| } | |
| } | |