v0.7.0: FNH gate de punição + árvore de pesos indexada (arestas duplas) + quebra-cabeça do Pensamento + NLP/NLG + decodificador especulativo + auto-ajuste SOM limitado + correções v070-1/2/3 — 73/73 gates, 84/84 cargo, provas r-28..r-32
4dbc907 verified | //! Tronco híbrido 5:1 (port do `mamba3/tronco.py`): 5 camadas ADBA (atenção | |
| //! block-causal com melhorias) + 1 camada SSD (slot Mamba-3 causal) — Teo D.1 | |
| //! com a restrição honesta H_D1 (bidirecionalidade intra-nível nas camadas de | |
| //! atenção; a camada recorrente é causal). | |
| //! | |
| //! Sequência do gerador: s = [cond; lev_1; …; lev_L] com entradas | |
| //! ẑ_{ℓ−1}^{(n_ℓ)} + e_ℓ (reconstrução cumulativa interpolada + embedding de | |
| //! nível — Teo D.1). Treino: 1 passada com teacher forcing (D.1(iv)); | |
| //! inferência: bloco a bloco com KV-cache (D.1(iii), T2). | |
| use candle_core::{Device, Result, Tensor, Var}; | |
| use crate::atencao::{AncoraNivel, BlocoAdba}; | |
| pub use crate::atencao::CacheKv; | |
| use crate::camadas::{Embedding, Linear, RmsNorm}; | |
| use crate::config::ConfigKhtst; | |
| use crate::ssd::BlocoSsd; | |
| pub struct Tronco { | |
| pub blocos_atencao: Vec<BlocoAdba>, | |
| pub bloco_ssd: BlocoSsd, | |
| pub norm_final: RmsNorm, | |
| pub ancora: Option<AncoraNivel>, | |
| /// gamma FiLM por nível (2*L+2: [escala, deslocamento] × nível) — ADBA. | |
| pub film_gamma: Option<Var>, | |
| pub film_beta: Option<Var>, | |
| /// W_f do FiLM (r-04 §1: x̃ = x⊙(1+γ·tanh(W_f·h)) + β·(W_f·h)) — ZERO no | |
| /// nascimento: tanh(0)=0 e β·0=0 ⇒ identidade bit-a-bit (Teo R.2). | |
| /// v070-1 (auditoria): a v0.6.0 injetava o resumo do âncora DIRETO (projeção | |
| /// aleatória) — nascimento não-neutro nos níveis ≥ 2; correção restaura o | |
| /// enunciado R.2 independentemente de γ=1. | |
| pub film_w: Option<crate::camadas::Linear>, | |
| pub n_niveis: usize, | |
| } | |
| impl Tronco { | |
| pub fn novo(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> { | |
| let m = &cfg.modelo; | |
| let l = cfg.l_niveis(); | |
| let adba = &cfg.adba; | |
| let mut blocos = Vec::with_capacity(m.n_camadas - 1); | |
| for _ in 0..(m.n_camadas - 1) { | |
| blocos.push(BlocoAdba::novo( | |
| m.d_modelo, m.n_cabecas, m.d_ff, l, | |
| adba.ativo, adba.reponderacao_nivel, adba.temperatura_nivel, dev, | |
| )?); | |
| } | |
| let ancora = if adba.ativo && adba.gamma_ancora > 0.0 { | |
| Some(AncoraNivel::nova(m.d_modelo, l, dev)?) | |
| } else { | |
| None | |
| }; | |
| let (film_gamma, film_beta, film_w) = if adba.ativo { | |
| let g = Tensor::ones((l + 1, m.d_modelo), candle_core::DType::F32, dev)?; | |
| let b = Tensor::zeros((l + 1, m.d_modelo), candle_core::DType::F32, dev)?; | |
| let w = crate::camadas::Linear::zeros_sem_vies(m.d_modelo, m.d_modelo, dev)?; | |
| (Some(Var::from_tensor(&g)?), Some(Var::from_tensor(&b)?), Some(w)) | |
| } else { | |
| (None, None, None) | |
| }; | |
| Ok(Tronco { | |
| blocos_atencao: blocos, | |
| bloco_ssd: BlocoSsd::novo(m.d_modelo, m.d_ff, dev)?, | |
| norm_final: RmsNorm::novo(m.d_modelo, 1e-5, dev)?, | |
| ancora, | |
| film_gamma, | |
| film_beta, | |
| film_w, | |
| n_niveis: l, | |
| }) | |
| } | |
| /// Injeção FiLM do resumo-âncora nas posições do nível `nivel` | |
| /// (ADBA (i): o resumo depende APENAS de c_{<ℓ} — paralelismo preservado). | |
| /// r-04 §1: x̃ = x⊙(1+γ·tanh(W_f·h)) + β·(W_f·h) com W_f=0 no nascimento | |
| /// ⇒ identidade bit-a-bit para QUALQUER resumo (Teo R.2, v070-1). | |
| pub fn aplicar_film(&self, x: &Tensor, resumo: &Tensor, nivel: usize) -> Result<Tensor> { | |
| match (&self.film_gamma, &self.film_beta, &self.film_w) { | |
| (Some(g), Some(b), Some(wf)) => { | |
| let d = x.dims()[2]; | |
| let bsz = resumo.dims()[0].max(1); | |
| let gv = g.as_tensor().get(nivel)?.reshape((1, 1, d))?; | |
| let bv = b.as_tensor().get(nivel)?.reshape((1, 1, d))?; | |
| // r = W_f·resumo — 0 no nascimento (identidade garantida) | |
| let r = wf.fwd(resumo)?.reshape((bsz, 1, d))?; | |
| // x ⊙ (1 + γ·tanh(r)) + β·r | |
| let gate = gv.broadcast_mul(&r.tanh()?)?.affine(1.0, 1.0)?; // 1 + γ·tanh(r) | |
| let xg = x.broadcast_mul(&gate)?; | |
| let desloc = bv.broadcast_mul(&r)?; | |
| xg.broadcast_add(&desloc) | |
| } | |
| _ => Ok(x.clone()), | |
| } | |
| } | |
| /// Forward de treino: sequência completa s (b, T, d) com máscara aditiva | |
| /// (T, T) e mapa de níveis por posição (query e chave). 1 passada — D.1(iv). | |
| pub fn forward_treino( | |
| &self, | |
| s: &Tensor, | |
| mascara: &Tensor, | |
| niveis: &[usize], | |
| resumos_por_nivel: &mut Vec<Tensor>, | |
| ) -> Result<Tensor> { | |
| let b = s.dim(0)?; | |
| let mut x = s.clone(); | |
| for (i, bloco) in self.blocos_atencao.iter().enumerate() { | |
| let mut cache = CacheKv::default(); | |
| x = bloco.fwd(&x, Some(mascara), &mut cache, Some(niveis), Some(niveis))?; | |
| // após o 1º bloco: injeta o resumo-âncora por nível (ADBA) | |
| if i == 0 { | |
| if let Some(anc) = &self.ancora { | |
| resumos_por_nivel.clear(); | |
| let n_prefixo = x.dims()[1] - self.tamanhos_por_nivel(niveis).iter().sum::<usize>(); | |
| // prefixo entra inalterado (sem FiLM — nível 0) | |
| let mut partes = vec![x.narrow(1, 0, n_prefixo)?]; | |
| let mut inicio = n_prefixo; | |
| for (nivel, &n_n) in self.tamanhos_por_nivel(niveis).iter().enumerate() { | |
| let trecho = x.narrow(1, inicio, n_n)?; | |
| // resumo do CONTEXTO ANTERIOR (níveis < atual): usa as | |
| // posições já processadas (função de c_{<ℓ} apenas) | |
| if inicio == n_prefixo { | |
| resumos_por_nivel.push(Tensor::zeros((b, s.dims()[2]), candle_core::DType::F32, s.device())?); | |
| } else { | |
| let ctx = x.narrow(1, 0, inicio)?; | |
| let mut r_nivel = Vec::new(); | |
| for bi in 0..b { | |
| r_nivel.push(anc.resumir(&ctx.narrow(0, bi, 1)?, nivel)?); | |
| } | |
| resumos_por_nivel.push(Tensor::cat(&r_nivel, 0)?); | |
| } | |
| let r = resumos_por_nivel[nivel].clone(); | |
| partes.push(self.aplicar_film(&trecho, &r, nivel + 1)?); | |
| inicio += n_n; | |
| } | |
| x = Tensor::cat(&partes, 1)?; | |
| } | |
| } | |
| } | |
| let mut cache_ssd = None; | |
| x = self.bloco_ssd.fwd(&x, &mut cache_ssd)?; | |
| self.norm_final.fwd(&x) | |
| } | |
| /// Forward incremental (inferência): bloco novo (b, n_novo, d) + caches. | |
| /// `niveis_novos`/`niveis_cache` mapeiam posições para níveis. | |
| /// `film` = Some((resumo, nivel 1-indexado)) aplica o FiLM nas posições novas | |
| /// após o bloco 0 — PARIDADE com forward_treino (v070-2: a v0.6.0 não | |
| /// aplicava FiLM na inferência incremental — treino↔geração divergiam). | |
| /// `saida_b0` = Some(&mut Tensor) recebe os estados pós-bloco-0 das | |
| /// posições novas (para o âncora acumular contexto nos níveis seguintes). | |
| pub fn forward_nivel( | |
| &self, | |
| bloco_x: &Tensor, | |
| caches: &mut Vec<CacheKv>, | |
| mascara_nova: &Tensor, | |
| niveis_novos: &[usize], | |
| niveis_totais: &[usize], | |
| film: Option<(&Tensor, usize)>, | |
| mut saida_b0: Option<&mut Tensor>, | |
| ) -> Result<Tensor> { | |
| let mut x = bloco_x.clone(); | |
| for (i, bloco) in self.blocos_atencao.iter().enumerate() { | |
| if caches.len() <= i { | |
| caches.push(CacheKv::default()); | |
| } | |
| x = bloco.fwd(&x, Some(mascara_nova), &mut caches[i], Some(niveis_novos), Some(niveis_totais))?; | |
| if i == 0 { | |
| if let Some(slot) = saida_b0.as_deref_mut() { | |
| *slot = x.clone(); | |
| } | |
| if let Some((resumo, nivel)) = film { | |
| x = self.aplicar_film(&x, resumo, nivel)?; | |
| } | |
| } | |
| } | |
| let mut cache_ssd = None; | |
| x = self.bloco_ssd.fwd(&x, &mut cache_ssd)?; | |
| self.norm_final.fwd(&x) | |
| } | |
| fn tamanhos_por_nivel(&self, niveis: &[usize]) -> Vec<usize> { | |
| let max_n = niveis.iter().copied().max().unwrap_or(0); | |
| let mut out = vec![0usize; max_n]; | |
| for &n in niveis { | |
| if n >= 1 && n <= max_n { | |
| out[n - 1] += 1; | |
| } | |
| } | |
| out | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| let mut v: Vec<Var> = self.blocos_atencao.iter().flat_map(|b| b.vars()).collect(); | |
| v.extend(self.bloco_ssd.vars()); | |
| v.extend(self.norm_final.vars()); | |
| if let Some(a) = &self.ancora { | |
| v.extend(a.vars()); | |
| } | |
| if let Some(g) = &self.film_gamma { | |
| v.push(g.clone()); | |
| } | |
| if let Some(b) = &self.film_beta { | |
| v.push(b.clone()); | |
| } | |
| if let Some(wf) = &self.film_w { | |
| v.push(wf.w.clone()); | |
| } | |
| v | |
| } | |
| } | |
| /// Cabeças de conceitos por nível (via A de T1): linear d → K_ℓ. | |
| pub struct CabecaConceitos { | |
| pub cabecas: Vec<Linear>, | |
| } | |
| impl CabecaConceitos { | |
| pub fn nova(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> { | |
| let mut c = Vec::with_capacity(cfg.l_niveis()); | |
| for _ in 0..cfg.l_niveis() { | |
| c.push(Linear::sem_vies(cfg.modelo.d_modelo, cfg.ncm.k_por_nivel, 0.02, dev)?); | |
| } | |
| Ok(CabecaConceitos { cabecas: c }) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| self.cabecas.iter().flat_map(|c| c.vars()).collect() | |
| } | |
| } | |
| /// Embeddings de nível e_ℓ (T1 via A). | |
| pub struct EmbNiveis { | |
| pub tabela: Embedding, | |
| } | |
| impl EmbNiveis { | |
| pub fn nova(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> { | |
| Ok(EmbNiveis { tabela: Embedding::novo(cfg.l_niveis() + 1, cfg.modelo.d_modelo, dev)? }) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| self.tabela.vars() | |
| } | |
| } | |