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 15 days ago

Raw History Blame Contribute Delete

8.62 kB

//! ADBA — Atenção Dual-Bloco Adaptativa (APRIMORAMENTO 1 do usuário:
//! "aprimoramento de Attention Mechanism durante o processamento paralelo
//! de NCM"). Prova completa na memória r-04; síntese:
//!
//! (i)  [paralelismo preservado] o resumo-âncora por nível h_ℓ é função
//!      APENAS de c_{<ℓ} (computado do cache/entrada corrente), logo a
//!      fatoração p(c_ℓ|c_{<ℓ}) = Π_i p(c_{ℓ,i}|c_{<ℓ}) e NFE = L
//!      permanecem EXATOS (Teo B.2/D.1);
//! (ii) [expressividade ⊇ NCM] com γ=0, w_ℓ'=1, τ_ℓ=1 a ADBA reduz-se à
//!      atenção block-causal do NCM (classe estritamente contida);
//! (iii)[planejamento reforçado] o âncora agrega o estado global dos níveis
//!      anteriores num vetor por nível — injetado por FiLM — fortalecendo a
//!      hipótese H2 do Teo C.2 (condicionais consistentes dado o plano).
//!
//! Componentes: (a) atenção multi-cabeça com máscara aditiva arbitrária +
//! KV-cache; (b) reponderação aprendível w_{ℓ'} dos logits por nível-fonte;
//! (c) temperatura aprendível τ_ℓ por nível; (d) injeção FiLM do resumo.

use candle_core::{Result, Tensor, Var, D};
use crate::camadas::{Linear, RmsNorm, Swiglu};

pub struct CabecaAdba {
    pub q: Linear,
    pub k: Linear,
    pub v: Linear,
    pub o: Linear,
    /// Reponderação de logits por nível-fonte (len = nº de níveis + 1 p/ prefixo).
    pub w_nivel: Option<Var>,
    /// Temperatura aprendível por nível (log-espaço, neutro = 0).
    pub log_tau_nivel: Option<Var>,
    pub n_cabecas: usize,
    pub d_cabeca: usize,
}

#[derive(Default, Clone)]
pub struct CacheKv {
    pub k: Option<Tensor>,
    pub v: Option<Tensor>,
}

impl CabecaAdba {
    pub fn nova(d_modelo: usize, n_cabecas: usize, n_niveis: usize, adba_ativo: bool, reponderar: bool, tau_nivel: bool, dev: &candle_core::Device) -> Result<Self> {
        let d_cabeca = d_modelo / n_cabecas;
        let desvio = 1.0 / ((d_modelo * 2) as f64).sqrt();
        let mk = || Linear::sem_vies(d_modelo, d_modelo, desvio, dev);
        let w_nivel = if adba_ativo && reponderar {
            Some(Var::from_tensor(&Tensor::ones((n_niveis + 1,), candle_core::DType::F32, dev)?)?)
        } else {
            None
        };
        let log_tau_nivel = if adba_ativo && tau_nivel {
            Some(Var::from_tensor(&Tensor::zeros((n_niveis + 1,), candle_core::DType::F32, dev)?)?)
        } else {
            None
        };
        Ok(CabecaAdba {
            q: mk()?, k: mk()?, v: mk()?, o: mk()?,
            w_nivel, log_tau_nivel,
            n_cabecas, d_cabeca,
        })
    }

    fn dividir_cabecas(&self, x: &Tensor) -> Result<Tensor> {
        let (b, n, _) = (x.dim(0)?, x.dim(1)?, x.dim(2)?);
        x.reshape((b, n, self.n_cabecas, self.d_cabeca))?
            .transpose(1, 2)?
            .contiguous()
    }

    fn juntar_cabecas(&self, x: &Tensor) -> Result<Tensor> {
        let b = x.dim(0)?;
        let n = x.dim(2)?;
        x.transpose(1, 2)?.contiguous()?
            .reshape((b, n, self.n_cabecas * self.d_cabeca))
    }

    /// Atenção com máscara aditiva (T_k × T_k) e offsets de nível para a
    /// reponderação/temperatura da ADBA. `niveis_query`/`niveis_chave` dizem
    /// a qual nível (0=prefixo, 1..L) pertence cada posição.
    pub fn attention(
        &self,
        x: &Tensor,
        mascara: Option<&Tensor>,
        cache: &mut CacheKv,
        niveis_query: Option<&[usize]>,
        niveis_chave: Option<&[usize]>,
    ) -> Result<Tensor> {
        let b = x.dim(0)?;
        let q = self.dividir_cabecas(&self.q.fwd(x)?)?;
        let k = self.dividir_cabecas(&self.k.fwd(x)?)?;
        let v = self.dividir_cabecas(&self.v.fwd(x)?)?;
        let (k, v) = match (&cache.k, &cache.v) {
            (Some(pk), Some(pv)) => (Tensor::cat(&[pk, &k], 2)?, Tensor::cat(&[pv, &v], 2)?),
            _ => (k, v),
        };
        cache.k = Some(k.clone());
        cache.v = Some(v.clone());
        let escala = 1.0 / (self.d_cabeca as f64).sqrt();
        let mut scores = q.matmul(&k.transpose(2, 3)?)?.affine(escala, 0.0)?;
        // ADBA (ii): reponderação por nível-fonte — score += log w_{ℓ'}
        if let (Some(w), Some(nq), Some(nk)) = (&self.w_nivel, niveis_query, niveis_chave) {
            let t_k = scores.dim(3)?;
            let w_clamped = w.as_tensor().clamp(1e-3f32, 1e3f32)?;
            let log_w = (&w_clamped + 1e-8)?.log()?;
            let mut biases = vec![0f32; t_k];
            for (j, &nivel) in nk.iter().enumerate().take(t_k) {
                biases[j] = log_w.get(nivel)?.to_scalar::<f32>()?;
            }
            let _ = nq;
            let bias = Tensor::from_vec(biases, (t_k,), x.device())?;
            scores = scores.broadcast_add(&bias.reshape((1, 1, 1, t_k))?)?;
        }
        // ADBA (iii): temperatura por nível do QUERY
        if let (Some(lt), Some(nq), _) = (&self.log_tau_nivel, niveis_query, niveis_chave) {
            let t_q = scores.dim(2)?;
            let tau = lt.as_tensor().affine(0.1, 0.0)?.exp()?; // suave (±~10% por unidade)
            let mut fatores = vec![1f32; t_q];
            for (i, &nivel) in nq.iter().enumerate().take(t_q) {
                fatores[i] = tau.get(nivel)?.to_scalar::<f32>()?;
            }
            let f = Tensor::from_vec(fatores, (t_q,), x.device())?
                .reshape((1, 1, t_q, 1))?;
            scores = scores.broadcast_div(&f)?;
        }
        if let Some(m) = mascara {
            scores = scores.broadcast_add(m)?;
        }
        let attn = candle_nn::ops::softmax(&scores, D::Minus1)?;
        let ctx = attn.matmul(&v)?;
        self.juntar_cabecas(&ctx)
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v: Vec<Var> = [self.q.vars(), self.k.vars(), self.v.vars(), self.o.vars()].concat();
        if let Some(w) = &self.w_nivel { v.push(w.clone()); }
        if let Some(t) = &self.log_tau_nivel { v.push(t.clone()); }
        v
    }
}

/// Bloco ADBA completo: pré-norm + atenção + resíduo + pré-norm + SwiGLU.
pub struct BlocoAdba {
    pub norm1: RmsNorm,
    pub attn: CabecaAdba,
    pub norm2: RmsNorm,
    pub ffn: Swiglu,
}

impl BlocoAdba {
    pub fn novo(d_modelo: usize, n_cabecas: usize, d_ff: usize, n_niveis: usize, adba_ativo: bool, reponderar: bool, tau_nivel: bool, dev: &candle_core::Device) -> Result<Self> {
        Ok(BlocoAdba {
            norm1: RmsNorm::novo(d_modelo, 1e-5, dev)?,
            attn: CabecaAdba::nova(d_modelo, n_cabecas, n_niveis, adba_ativo, reponderar, tau_nivel, dev)?,
            norm2: RmsNorm::novo(d_modelo, 1e-5, dev)?,
            ffn: Swiglu::novo(d_modelo, d_ff, dev)?,
        })
    }

    pub fn fwd(
        &self,
        x: &Tensor,
        mascara: Option<&Tensor>,
        cache: &mut CacheKv,
        niveis_query: Option<&[usize]>,
        niveis_chave: Option<&[usize]>,
    ) -> Result<Tensor> {
        let h = self.attn.attention(&self.norm1.fwd(x)?, mascara, cache, niveis_query, niveis_chave)?;
        let h = self.attn.o.fwd(&h)?;
        let x = x.add(&h)?;
        let y = self.ffn.fwd(&self.norm2.fwd(&x)?)?;
        x.add(&y)
    }

    pub fn vars(&self) -> Vec<Var> {
        [self.norm1.vars(), self.attn.vars(), self.norm2.vars(), self.ffn.vars()].concat()
    }
}

/// Resumo-âncora por nível (ADBA i): h_ℓ = Attn(q_ℓ, K/V = c_{<ℓ}) — pooling
/// sobre as posições dos níveis anteriores (função de c_{<ℓ} APENAS).
pub struct AncoraNivel {
    pub q: Var, // (d,) vetor de consulta aprendível por nível — (L, d)
    pub proj: Linear,
}

impl AncoraNivel {
    pub fn nova(d: usize, n_niveis: usize, dev: &candle_core::Device) -> Result<Self> {
        let q = Tensor::randn(0f32, 0.02, (n_niveis, d), dev)?;
        Ok(AncoraNivel { q: Var::from_tensor(&q)?, proj: Linear::novo(d, d, 0.02, dev)? })
    }

    /// contexto (1, n_ctx, d) das posições dos níveis < ℓ → resumo (1, d).
    pub fn resumir(&self, contexto: &Tensor, nivel: usize) -> Result<Tensor> {
        let d = self.q.as_tensor().dim(1)?;
        let q = self.q.as_tensor().get(nivel)?.reshape((1, 1, d))?;
        // atenção de cabeça única: softmax(q·Kᵀ/√d)·V
        let escala = 1.0 / (d as f64).sqrt();
        let scores = q.matmul(&contexto.transpose(1, 2)?)?.affine(escala, 0.0)?;
        let pesos = candle_nn::ops::softmax(&scores, D::Minus1)?;
        let resumo = pesos.matmul(contexto)?.reshape((1, d))?;
        self.proj.fwd(&resumo)
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = vec![self.q.clone()];
        v.extend(self.proj.vars());
        v
    }
}

	//! ADBA — Atenção Dual-Bloco Adaptativa (APRIMORAMENTO 1 do usuário:
	//! "aprimoramento de Attention Mechanism durante o processamento paralelo
	//! de NCM"). Prova completa na memória r-04; síntese:
	//!
	//! (i) [paralelismo preservado] o resumo-âncora por nível h_ℓ é função
	//! APENAS de c_{<ℓ} (computado do cache/entrada corrente), logo a
	//! fatoração p(c_ℓ\|c_{<ℓ}) = Π_i p(c_{ℓ,i}\|c_{<ℓ}) e NFE = L
	//! permanecem EXATOS (Teo B.2/D.1);
	//! (ii) [expressividade ⊇ NCM] com γ=0, w_ℓ'=1, τ_ℓ=1 a ADBA reduz-se à
	//! atenção block-causal do NCM (classe estritamente contida);
	//! (iii)[planejamento reforçado] o âncora agrega o estado global dos níveis
	//! anteriores num vetor por nível — injetado por FiLM — fortalecendo a
	//! hipótese H2 do Teo C.2 (condicionais consistentes dado o plano).
	//!
	//! Componentes: (a) atenção multi-cabeça com máscara aditiva arbitrária +
	//! KV-cache; (b) reponderação aprendível w_{ℓ'} dos logits por nível-fonte;
	//! (c) temperatura aprendível τ_ℓ por nível; (d) injeção FiLM do resumo.

	use candle_core::{Result, Tensor, Var, D};
	use crate::camadas::{Linear, RmsNorm, Swiglu};

	pub struct CabecaAdba {
	pub q: Linear,
	pub k: Linear,
	pub v: Linear,
	pub o: Linear,
	/// Reponderação de logits por nível-fonte (len = nº de níveis + 1 p/ prefixo).
	pub w_nivel: Option<Var>,
	/// Temperatura aprendível por nível (log-espaço, neutro = 0).
	pub log_tau_nivel: Option<Var>,
	pub n_cabecas: usize,
	pub d_cabeca: usize,
	}

	#[derive(Default, Clone)]
	pub struct CacheKv {
	pub k: Option<Tensor>,
	pub v: Option<Tensor>,
	}

	impl CabecaAdba {
	pub fn nova(d_modelo: usize, n_cabecas: usize, n_niveis: usize, adba_ativo: bool, reponderar: bool, tau_nivel: bool, dev: &candle_core::Device) -> Result<Self> {
	let d_cabeca = d_modelo / n_cabecas;
	let desvio = 1.0 / ((d_modelo * 2) as f64).sqrt();
	let mk = \|\| Linear::sem_vies(d_modelo, d_modelo, desvio, dev);
	let w_nivel = if adba_ativo && reponderar {
	Some(Var::from_tensor(&Tensor::ones((n_niveis + 1,), candle_core::DType::F32, dev)?)?)
	} else {
	None
	};
	let log_tau_nivel = if adba_ativo && tau_nivel {
	Some(Var::from_tensor(&Tensor::zeros((n_niveis + 1,), candle_core::DType::F32, dev)?)?)
	} else {
	None
	};
	Ok(CabecaAdba {
	q: mk()?, k: mk()?, v: mk()?, o: mk()?,
	w_nivel, log_tau_nivel,
	n_cabecas, d_cabeca,
	})
	}

	fn dividir_cabecas(&self, x: &Tensor) -> Result<Tensor> {
	let (b, n, _) = (x.dim(0)?, x.dim(1)?, x.dim(2)?);
	x.reshape((b, n, self.n_cabecas, self.d_cabeca))?
	.transpose(1, 2)?
	.contiguous()
	}

	fn juntar_cabecas(&self, x: &Tensor) -> Result<Tensor> {
	let b = x.dim(0)?;
	let n = x.dim(2)?;
	x.transpose(1, 2)?.contiguous()?
	.reshape((b, n, self.n_cabecas * self.d_cabeca))
	}

	/// Atenção com máscara aditiva (T_k × T_k) e offsets de nível para a
	/// reponderação/temperatura da ADBA. `niveis_query`/`niveis_chave` dizem
	/// a qual nível (0=prefixo, 1..L) pertence cada posição.
	pub fn attention(
	&self,
	x: &Tensor,
	mascara: Option<&Tensor>,
	cache: &mut CacheKv,
	niveis_query: Option<&[usize]>,
	niveis_chave: Option<&[usize]>,
	) -> Result<Tensor> {
	let b = x.dim(0)?;
	let q = self.dividir_cabecas(&self.q.fwd(x)?)?;
	let k = self.dividir_cabecas(&self.k.fwd(x)?)?;
	let v = self.dividir_cabecas(&self.v.fwd(x)?)?;
	let (k, v) = match (&cache.k, &cache.v) {
	(Some(pk), Some(pv)) => (Tensor::cat(&[pk, &k], 2)?, Tensor::cat(&[pv, &v], 2)?),
	_ => (k, v),
	};
	cache.k = Some(k.clone());
	cache.v = Some(v.clone());
	let escala = 1.0 / (self.d_cabeca as f64).sqrt();
	let mut scores = q.matmul(&k.transpose(2, 3)?)?.affine(escala, 0.0)?;
	// ADBA (ii): reponderação por nível-fonte — score += log w_{ℓ'}
	if let (Some(w), Some(nq), Some(nk)) = (&self.w_nivel, niveis_query, niveis_chave) {
	let t_k = scores.dim(3)?;
	let w_clamped = w.as_tensor().clamp(1e-3f32, 1e3f32)?;
	let log_w = (&w_clamped + 1e-8)?.log()?;
	let mut biases = vec![0f32; t_k];
	for (j, &nivel) in nk.iter().enumerate().take(t_k) {
	biases[j] = log_w.get(nivel)?.to_scalar::<f32>()?;
	}
	let _ = nq;
	let bias = Tensor::from_vec(biases, (t_k,), x.device())?;
	scores = scores.broadcast_add(&bias.reshape((1, 1, 1, t_k))?)?;
	}
	// ADBA (iii): temperatura por nível do QUERY
	if let (Some(lt), Some(nq), _) = (&self.log_tau_nivel, niveis_query, niveis_chave) {
	let t_q = scores.dim(2)?;
	let tau = lt.as_tensor().affine(0.1, 0.0)?.exp()?; // suave (±~10% por unidade)
	let mut fatores = vec![1f32; t_q];
	for (i, &nivel) in nq.iter().enumerate().take(t_q) {
	fatores[i] = tau.get(nivel)?.to_scalar::<f32>()?;
	}
	let f = Tensor::from_vec(fatores, (t_q,), x.device())?
	.reshape((1, 1, t_q, 1))?;
	scores = scores.broadcast_div(&f)?;
	}
	if let Some(m) = mascara {
	scores = scores.broadcast_add(m)?;
	}
	let attn = candle_nn::ops::softmax(&scores, D::Minus1)?;
	let ctx = attn.matmul(&v)?;
	self.juntar_cabecas(&ctx)
	}

	pub fn vars(&self) -> Vec<Var> {
	let mut v: Vec<Var> = [self.q.vars(), self.k.vars(), self.v.vars(), self.o.vars()].concat();
	if let Some(w) = &self.w_nivel { v.push(w.clone()); }
	if let Some(t) = &self.log_tau_nivel { v.push(t.clone()); }
	v
	}
	}

	/// Bloco ADBA completo: pré-norm + atenção + resíduo + pré-norm + SwiGLU.
	pub struct BlocoAdba {
	pub norm1: RmsNorm,
	pub attn: CabecaAdba,
	pub norm2: RmsNorm,
	pub ffn: Swiglu,
	}

	impl BlocoAdba {
	pub fn novo(d_modelo: usize, n_cabecas: usize, d_ff: usize, n_niveis: usize, adba_ativo: bool, reponderar: bool, tau_nivel: bool, dev: &candle_core::Device) -> Result<Self> {
	Ok(BlocoAdba {
	norm1: RmsNorm::novo(d_modelo, 1e-5, dev)?,
	attn: CabecaAdba::nova(d_modelo, n_cabecas, n_niveis, adba_ativo, reponderar, tau_nivel, dev)?,
	norm2: RmsNorm::novo(d_modelo, 1e-5, dev)?,
	ffn: Swiglu::novo(d_modelo, d_ff, dev)?,
	})
	}

	pub fn fwd(
	&self,
	x: &Tensor,
	mascara: Option<&Tensor>,
	cache: &mut CacheKv,
	niveis_query: Option<&[usize]>,
	niveis_chave: Option<&[usize]>,
	) -> Result<Tensor> {
	let h = self.attn.attention(&self.norm1.fwd(x)?, mascara, cache, niveis_query, niveis_chave)?;
	let h = self.attn.o.fwd(&h)?;
	let x = x.add(&h)?;
	let y = self.ffn.fwd(&self.norm2.fwd(&x)?)?;
	x.add(&y)
	}

	pub fn vars(&self) -> Vec<Var> {
	[self.norm1.vars(), self.attn.vars(), self.norm2.vars(), self.ffn.vars()].concat()
	}
	}

	/// Resumo-âncora por nível (ADBA i): h_ℓ = Attn(q_ℓ, K/V = c_{<ℓ}) — pooling
	/// sobre as posições dos níveis anteriores (função de c_{<ℓ} APENAS).
	pub struct AncoraNivel {
	pub q: Var, // (d,) vetor de consulta aprendível por nível — (L, d)
	pub proj: Linear,
	}

	impl AncoraNivel {
	pub fn nova(d: usize, n_niveis: usize, dev: &candle_core::Device) -> Result<Self> {
	let q = Tensor::randn(0f32, 0.02, (n_niveis, d), dev)?;
	Ok(AncoraNivel { q: Var::from_tensor(&q)?, proj: Linear::novo(d, d, 0.02, dev)? })
	}

	/// contexto (1, n_ctx, d) das posições dos níveis < ℓ → resumo (1, d).
	pub fn resumir(&self, contexto: &Tensor, nivel: usize) -> Result<Tensor> {
	let d = self.q.as_tensor().dim(1)?;
	let q = self.q.as_tensor().get(nivel)?.reshape((1, 1, d))?;
	// atenção de cabeça única: softmax(q·Kᵀ/√d)·V
	let escala = 1.0 / (d as f64).sqrt();
	let scores = q.matmul(&contexto.transpose(1, 2)?)?.affine(escala, 0.0)?;
	let pesos = candle_nn::ops::softmax(&scores, D::Minus1)?;
	let resumo = pesos.matmul(contexto)?.reshape((1, d))?;
	self.proj.fwd(&resumo)
	}

	pub fn vars(&self) -> Vec<Var> {
	let mut v = vec![self.q.clone()];
	v.extend(self.proj.vars());
	v
	}
	}