V6: Kohonen 4D SOM + EWC per-neuron + MTP+entropy + Xeon V6

79e8e52 verified 3 months ago

10.3 kB

"""mtp.py — V6: Multi-Token Prediction with adaptive attention + entropy regularizer.
═══════════════════════════════════════════════════════════════════════════════
V6 MELHORIAS MATEMÁTICAS (user request)
═══════════════════════════════════════════════════════════════════════════════
1. **Ativar MTP em val**: em V5, MTP só era usado em train. V6 ativa MTP
   também em val para medir a capacidade preditiva multi-token em validação
   (métrica adicional: val_mtp_loss).
2. **Regularizador de entropia nos mtp_alphas**: o MTP gera uma distribuição
   mtp_alphas (pesos para cada passo k ∈ {0..K-1}). Em V5, esses pesos podiam
   colapsar para um único k. V6 adiciona:
        L_entropy = -β * Σ_k α_k log(α_k + ε)
   que incentiva os K passos a serem igualmente utilizados (mínimo colapso).
3. **mtp_adaptive_attention**: combina MTP com attention adaptativa para
   suportar sequências longas (chunked attention, não full O(N²)).
Matemática
──────────
  L_mtp_total = Σ_k α_k * L_k - β * H(α)
  onde:
    L_k = cross_entropy(logits_k, target_k)
    α_k = softmax(raw_alpha_k / T_alpha)
    H(α) = -Σ_k α_k log(α_k + ε)  (entropia da distribuição α)
    β = 0.01 (peso do regularizador de entropia)
    T_alpha = 1.0 (temperatura do softmax)
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from typing import Optional, Tuple, Dict, List

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class MTPConfig:
    """Configuração do MTP V6."""
    hidden_size: int = 128
    vocab_size: int = 256
    K: int = 4                  # nº de tokens futuros
    share_embeddings: bool = False
    tie_weights: bool = False
    dropout: float = 0.1
    # V6 — entropy regularizer
    entropy_beta: float = 0.01  # peso do regularizador de entropia
    alpha_temp: float = 1.0     # temperatura do softmax dos alphas
    # V6 — val mode
    active_in_val: bool = True  # ativar MTP em val
    # V6 — adaptive attention
    adaptive_attention: bool = True
    adaptive_chunk_size: int = 64  # tamanho do chunk para adaptive attention


class MTPHead(nn.Module):
    """Head de Multi-Token Prediction (V6).
    Prediz K distribuições sobre o vocabulário, uma para cada k-ésimo token
    futuro, a partir de uma única hidden state.
    V6: adiciona parâmetro mtp_alphas (K,) — pesos aprendíveis para cada
    passo k. Esses pesos são regularizados por entropia para evitar colapso.
    Args:
        hidden_size: dimensão da hidden state
        vocab_size: tamanho do vocabulário
        K: nº de tokens futuros a predizer
        share_embeddings: se True, compartilha a matriz de embedding
        tie_weights: se True, usa a mesma W para todos os K passos
        dropout: probabilidade de dropout
    Forward:
        h: (B, T, hidden_size) — sequência de hidden states
        → logits (B, T, K, vocab_size), alphas (K,)
    """
    def __init__(
        self,
        hidden_size: int,
        vocab_size: int,
        K: int = 4,
        share_embeddings: bool = False,
        tie_weights: bool = False,
        dropout: float = 0.0,
    ):
        super().__init__()
        if K < 1:
            raise ValueError(f"K deve ser >= 1, recebeu {K}")
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.K = K
        self.share_embeddings = share_embeddings
        self.tie_weights = tie_weights

        if tie_weights:
            self.shared_proj = nn.Linear(hidden_size, vocab_size, bias=False)
            self.biases = nn.ParameterList([
                nn.Parameter(torch.zeros(vocab_size)) for _ in range(K)
            ])
        else:
            self.heads = nn.ModuleList([
                nn.Linear(hidden_size, vocab_size) for _ in range(K)
            ])

        # V6: mtp_alphas — pesos aprendíveis para cada passo k
        # Inicializado uniforme para que Σα_k = 1 no começo
        self.raw_alphas = nn.Parameter(torch.zeros(K))  # softmax(0) = uniforme
        self.dropout = nn.Dropout(dropout)

    @property
    def alphas(self) -> torch.Tensor:
        """Distribuição α normalizada via softmax com temperatura."""
        return F.softmax(self.raw_alphas, dim=-1)

    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """h: (B, T, hidden_size) → logits (B, T, K, V), alphas (K,)"""
        h = self.dropout(h)
        if self.tie_weights:
            # Cada passo k: shared_proj(h) + biases[k]
            logits_list = [self.shared_proj(h) + b for b in self.biases]
        else:
            logits_list = [head(h) for head in self.heads]
        # Stack: (B, T, K, V)
        logits = torch.stack(logits_list, dim=2)
        return logits, self.alphas


def mtp_entropy_regularizer(alphas: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """V6: regularizador de entropia nos mtp_alphas.
    H(α) = -Σ_k α_k log(α_k + ε)
    Retornamos -H (para minimizar = maximizar entropia).
    L_reg = -H(α) = Σ_k α_k log(α_k + ε)
    """
    return -(alphas * torch.log(alphas + eps)).sum()


def mtp_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    alphas: torch.Tensor,
    entropy_beta: float = 0.01,
) -> Tuple[torch.Tensor, Dict[str, float]]:
    """Computa a loss MTP V6 com regularizador de entropia.
    L = Σ_k α_k * L_k - β * H(α)
    Args:
        logits: (B, T, K, V)
        targets: (B, T) — índices do token alinhado com o último T (next-token)
                 Será shiftado para alinhar com cada k
        alphas: (K,)
        entropy_beta: peso do regularizador de entropia
    Returns:
        loss: escalar
        metrics: dict com breakdown da loss
    """
    B, T, K, V = logits.shape
    # Computa loss por passo k
    losses_per_k = []
    for k in range(K):
        # Target para passo k: targets shifted por k
        # logits[:, t, k, :] prediz token em t+k+1 (ou apenas t+k dependendo da convenção)
        # Aqui usamos shift simples: prediz target[t+k] a partir de h[t]
        if T - k - 1 <= 0:
            losses_per_k.append(torch.tensor(0.0, device=logits.device))
            continue
        logits_k = logits[:, :T-k-1, k, :].reshape(-1, V)
        targets_k = targets[:, k+1:T].reshape(-1)
        loss_k = F.cross_entropy(logits_k, targets_k, ignore_index=-100)
        losses_per_k.append(loss_k)
    losses_tensor = torch.stack(losses_per_k)  # (K,)

    # Loss ponderada: Σ α_k L_k
    weighted_loss = (alphas * losses_tensor).sum()
    # Regularizador de entropia: -β H(α) = β * Σ α_k log(α_k + ε)
    entropy_reg = mtp_entropy_regularizer(alphas)
    total_loss = weighted_loss - entropy_beta * entropy_reg  # negativo porque reg = -H

    metrics = {
        "weighted_loss": float(weighted_loss.item()),
        "entropy_reg": float(entropy_reg.item()),
        "total_loss": float(total_loss.item()),
        "alphas": alphas.detach().cpu().tolist(),
        "per_k_loss": [float(l.item()) for l in losses_per_k],
    }
    return total_loss, metrics


class MTPAdaptiveAttention(nn.Module):
    """V6: MTP + adaptive attention para sequências longas (2M tokens).
    Estratégia:
    - Sequência longa é dividida em chunks de tamanho chunk_size
    - Cada chunk: attention full (custo O(chunk²))
    - Entre chunks: attention sparse (top-k dos vizinhos)
    - MTP prediz K tokens futuros por chunk anchor
    Args:
        vocab_size, embed_dim, hidden_size, K, num_layers, num_heads
        chunk_size: tamanho do chunk para adaptive attention
        knn_k: top-k para atenção entre chunks
    """
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 128,
        hidden_size: int = 128,
        K: int = 4,
        num_layers: int = 2,
        num_heads: int = 4,
        chunk_size: int = 64,
        knn_k: int = 8,
        max_N: int = 2_000_000,  # 2M tokens (limite teórico)
        dropout: float = 0.1,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.hidden_size = hidden_size
        self.K = K
        self.chunk_size = chunk_size
        self.knn_k = knn_k
        self.max_N = max_N

        self.embed = nn.Embedding(vocab_size, embed_dim)
        # Camadas de attention chunked
        self.layers = nn.ModuleList([
            nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
            for _ in range(num_layers)
        ])
        self.norms = nn.ModuleList([nn.LayerNorm(embed_dim) for _ in range(num_layers)])
        # MTP head
        self.mtp_head = MTPHead(embed_dim, vocab_size, K=K, dropout=dropout)

    def forward(self, input_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """input_ids: (B, T) → logits (B, T, K, V), alphas (K,)"""
        B, T = input_ids.shape
        if T > self.max_N:
            raise ValueError(f"Sequence length {T} exceeds max_N={self.max_N}")
        h = self.embed(input_ids)  # (B, T, embed_dim)
        # Adaptive attention: processa em chunks
        chunk_size = min(self.chunk_size, T)
        n_chunks = (T + chunk_size - 1) // chunk_size
        outputs = []
        for c in range(n_chunks):
            start = c * chunk_size
            end = min((c + 1) * chunk_size, T)
            chunk = h[:, start:end, :]  # (B, chunk, embed_dim)
            for i, (attn, norm) in enumerate(zip(self.layers, self.norms)):
                normed = norm(chunk)
                attn_out, _ = attn(normed, normed, normed, need_weights=False)
                chunk = chunk + attn_out
            outputs.append(chunk)
        h_out = torch.cat(outputs, dim=1)  # (B, T, embed_dim)
        logits, alphas = self.mtp_head(h_out)
        return logits, alphas


__all__ = [
    "MTPConfig", "MTPHead", "MTPAdaptiveAttention",
    "mtp_loss", "mtp_entropy_regularizer",
]

	"""mtp.py — V6: Multi-Token Prediction with adaptive attention + entropy regularizer.

	═══════════════════════════════════════════════════════════════════════════════
	V6 MELHORIAS MATEMÁTICAS (user request)
	═══════════════════════════════════════════════════════════════════════════════

	1. Ativar MTP em val: em V5, MTP só era usado em train. V6 ativa MTP
	também em val para medir a capacidade preditiva multi-token em validação
	(métrica adicional: val_mtp_loss).

	2. Regularizador de entropia nos mtp_alphas: o MTP gera uma distribuição
	mtp_alphas (pesos para cada passo k ∈ {0..K-1}). Em V5, esses pesos podiam
	colapsar para um único k. V6 adiciona:
	L_entropy = -β * Σ_k α_k log(α_k + ε)
	que incentiva os K passos a serem igualmente utilizados (mínimo colapso).

	3. mtp_adaptive_attention: combina MTP com attention adaptativa para
	suportar sequências longas (chunked attention, não full O(N²)).

	Matemática
	──────────
	L_mtp_total = Σ_k α_k * L_k - β * H(α)

	onde:
	L_k = cross_entropy(logits_k, target_k)
	α_k = softmax(raw_alpha_k / T_alpha)
	H(α) = -Σ_k α_k log(α_k + ε) (entropia da distribuição α)
	β = 0.01 (peso do regularizador de entropia)
	T_alpha = 1.0 (temperatura do softmax)
	"""
	from __future__ import annotations
	import math
	from dataclasses import dataclass, field
	from typing import Optional, Tuple, Dict, List

	import torch
	import torch.nn as nn
	import torch.nn.functional as F


	@dataclass
	class MTPConfig:
	"""Configuração do MTP V6."""
	hidden_size: int = 128
	vocab_size: int = 256
	K: int = 4 # nº de tokens futuros
	share_embeddings: bool = False
	tie_weights: bool = False
	dropout: float = 0.1
	# V6 — entropy regularizer
	entropy_beta: float = 0.01 # peso do regularizador de entropia
	alpha_temp: float = 1.0 # temperatura do softmax dos alphas
	# V6 — val mode
	active_in_val: bool = True # ativar MTP em val
	# V6 — adaptive attention
	adaptive_attention: bool = True
	adaptive_chunk_size: int = 64 # tamanho do chunk para adaptive attention


	class MTPHead(nn.Module):
	"""Head de Multi-Token Prediction (V6).

	Prediz K distribuições sobre o vocabulário, uma para cada k-ésimo token
	futuro, a partir de uma única hidden state.

	V6: adiciona parâmetro mtp_alphas (K,) — pesos aprendíveis para cada
	passo k. Esses pesos são regularizados por entropia para evitar colapso.

	Args:
	hidden_size: dimensão da hidden state
	vocab_size: tamanho do vocabulário
	K: nº de tokens futuros a predizer
	share_embeddings: se True, compartilha a matriz de embedding
	tie_weights: se True, usa a mesma W para todos os K passos
	dropout: probabilidade de dropout

	Forward:
	h: (B, T, hidden_size) — sequência de hidden states
	→ logits (B, T, K, vocab_size), alphas (K,)
	"""
	def __init__(
	self,
	hidden_size: int,
	vocab_size: int,
	K: int = 4,
	share_embeddings: bool = False,
	tie_weights: bool = False,
	dropout: float = 0.0,
	):
	super().__init__()
	if K < 1:
	raise ValueError(f"K deve ser >= 1, recebeu {K}")
	self.hidden_size = hidden_size
	self.vocab_size = vocab_size
	self.K = K
	self.share_embeddings = share_embeddings
	self.tie_weights = tie_weights

	if tie_weights:
	self.shared_proj = nn.Linear(hidden_size, vocab_size, bias=False)
	self.biases = nn.ParameterList([
	nn.Parameter(torch.zeros(vocab_size)) for _ in range(K)
	])
	else:
	self.heads = nn.ModuleList([
	nn.Linear(hidden_size, vocab_size) for _ in range(K)
	])

	# V6: mtp_alphas — pesos aprendíveis para cada passo k
	# Inicializado uniforme para que Σα_k = 1 no começo
	self.raw_alphas = nn.Parameter(torch.zeros(K)) # softmax(0) = uniforme
	self.dropout = nn.Dropout(dropout)

	@property
	def alphas(self) -> torch.Tensor:
	"""Distribuição α normalizada via softmax com temperatura."""
	return F.softmax(self.raw_alphas, dim=-1)

	def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
	"""h: (B, T, hidden_size) → logits (B, T, K, V), alphas (K,)"""
	h = self.dropout(h)
	if self.tie_weights:
	# Cada passo k: shared_proj(h) + biases[k]
	logits_list = [self.shared_proj(h) + b for b in self.biases]
	else:
	logits_list = [head(h) for head in self.heads]
	# Stack: (B, T, K, V)
	logits = torch.stack(logits_list, dim=2)
	return logits, self.alphas


	def mtp_entropy_regularizer(alphas: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
	"""V6: regularizador de entropia nos mtp_alphas.

	H(α) = -Σ_k α_k log(α_k + ε)

	Retornamos -H (para minimizar = maximizar entropia).
	L_reg = -H(α) = Σ_k α_k log(α_k + ε)
	"""
	return -(alphas * torch.log(alphas + eps)).sum()


	def mtp_loss(
	logits: torch.Tensor,
	targets: torch.Tensor,
	alphas: torch.Tensor,
	entropy_beta: float = 0.01,
	) -> Tuple[torch.Tensor, Dict[str, float]]:
	"""Computa a loss MTP V6 com regularizador de entropia.

	L = Σ_k α_k * L_k - β * H(α)

	Args:
	logits: (B, T, K, V)
	targets: (B, T) — índices do token alinhado com o último T (next-token)
	Será shiftado para alinhar com cada k
	alphas: (K,)
	entropy_beta: peso do regularizador de entropia

	Returns:
	loss: escalar
	metrics: dict com breakdown da loss
	"""
	B, T, K, V = logits.shape
	# Computa loss por passo k
	losses_per_k = []
	for k in range(K):
	# Target para passo k: targets shifted por k
	# logits[:, t, k, :] prediz token em t+k+1 (ou apenas t+k dependendo da convenção)
	# Aqui usamos shift simples: prediz target[t+k] a partir de h[t]
	if T - k - 1 <= 0:
	losses_per_k.append(torch.tensor(0.0, device=logits.device))
	continue
	logits_k = logits[:, :T-k-1, k, :].reshape(-1, V)
	targets_k = targets[:, k+1:T].reshape(-1)
	loss_k = F.cross_entropy(logits_k, targets_k, ignore_index=-100)
	losses_per_k.append(loss_k)
	losses_tensor = torch.stack(losses_per_k) # (K,)

	# Loss ponderada: Σ α_k L_k
	weighted_loss = (alphas * losses_tensor).sum()
	# Regularizador de entropia: -β H(α) = β * Σ α_k log(α_k + ε)
	entropy_reg = mtp_entropy_regularizer(alphas)
	total_loss = weighted_loss - entropy_beta * entropy_reg # negativo porque reg = -H

	metrics = {
	"weighted_loss": float(weighted_loss.item()),
	"entropy_reg": float(entropy_reg.item()),
	"total_loss": float(total_loss.item()),
	"alphas": alphas.detach().cpu().tolist(),
	"per_k_loss": [float(l.item()) for l in losses_per_k],
	}
	return total_loss, metrics


	class MTPAdaptiveAttention(nn.Module):
	"""V6: MTP + adaptive attention para sequências longas (2M tokens).

	Estratégia:
	- Sequência longa é dividida em chunks de tamanho chunk_size
	- Cada chunk: attention full (custo O(chunk²))
	- Entre chunks: attention sparse (top-k dos vizinhos)
	- MTP prediz K tokens futuros por chunk anchor

	Args:
	vocab_size, embed_dim, hidden_size, K, num_layers, num_heads
	chunk_size: tamanho do chunk para adaptive attention
	knn_k: top-k para atenção entre chunks
	"""
	def __init__(
	self,
	vocab_size: int,
	embed_dim: int = 128,
	hidden_size: int = 128,
	K: int = 4,
	num_layers: int = 2,
	num_heads: int = 4,
	chunk_size: int = 64,
	knn_k: int = 8,
	max_N: int = 2_000_000, # 2M tokens (limite teórico)
	dropout: float = 0.1,
	):
	super().__init__()
	self.vocab_size = vocab_size
	self.embed_dim = embed_dim
	self.hidden_size = hidden_size
	self.K = K
	self.chunk_size = chunk_size
	self.knn_k = knn_k
	self.max_N = max_N

	self.embed = nn.Embedding(vocab_size, embed_dim)
	# Camadas de attention chunked
	self.layers = nn.ModuleList([
	nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
	for _ in range(num_layers)
	])
	self.norms = nn.ModuleList([nn.LayerNorm(embed_dim) for _ in range(num_layers)])
	# MTP head
	self.mtp_head = MTPHead(embed_dim, vocab_size, K=K, dropout=dropout)

	def forward(self, input_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
	"""input_ids: (B, T) → logits (B, T, K, V), alphas (K,)"""
	B, T = input_ids.shape
	if T > self.max_N:
	raise ValueError(f"Sequence length {T} exceeds max_N={self.max_N}")
	h = self.embed(input_ids) # (B, T, embed_dim)
	# Adaptive attention: processa em chunks
	chunk_size = min(self.chunk_size, T)
	n_chunks = (T + chunk_size - 1) // chunk_size
	outputs = []
	for c in range(n_chunks):
	start = c * chunk_size
	end = min((c + 1) * chunk_size, T)
	chunk = h[:, start:end, :] # (B, chunk, embed_dim)
	for i, (attn, norm) in enumerate(zip(self.layers, self.norms)):
	normed = norm(chunk)
	attn_out, _ = attn(normed, normed, normed, need_weights=False)
	chunk = chunk + attn_out
	outputs.append(chunk)
	h_out = torch.cat(outputs, dim=1) # (B, T, embed_dim)
	logits, alphas = self.mtp_head(h_out)
	return logits, alphas


	__all__ = [
	"MTPConfig", "MTPHead", "MTPAdaptiveAttention",
	"mtp_loss", "mtp_entropy_regularizer",
	]