Delete src/bigru_t/model/orq_cell.py

src/bigru_t/model/orq_cell.py

@@ -1,738 +0,0 @@
-"""orq_cell.py — V6: OrqCell Kohonen 4D Self-Organizing Map (SOM).
-
-═══════════════════════════════════════════════════════════════════════════════
-V6 — KOHONEN 4D SELF-ORGANIZING MAP
-═══════════════════════════════════════════════════════════════════════════════
-
-Reescrito completamente em V6 como um SOM 4D (Kohonen). A camada OrqCell agora
-é o ORQUESTRADOR AUTO-ORGANIZÁVEL do BiGRU_T — semelhante ao PGVector em
-espírito (busca por similaridade), mas que aprende a topologia dos dados ao
-longo do treinamento, em vez de ser uma simples memória vetorial estática.
-
-Mapeamento arquitetural (definido pelo usuário):
-
-  ┌─────────────────────────────────────────────────────────────────────────┐
-  │ Componente SOM               │ Implementação BiGRU_T                    │
-  ├─────────────────────────────────────────────────────────────────────────┤
-  │ Neurônios do mapa (8)        │ 8 BiGRU4 paralelas dentro de cada       │
-  │                              │ u8cell_T — cada BiGRU4 tem peso próprio │
-  │ Pesos dos neurônios          │ Parâmetros aprendíveis (cache + proj_in │
-  │                              │ + proj_out + TransformerUnit)           │
-  │ Dimensão semântica           │ TransformerUnit (atenção = função de    │
-  │ (função de vizinhança)       │ vizinhança aprendível)                  │
-  │ Mecanismo de atualização     │ Best-Matching Unit (BMU) + vizinhos     │
-  │                              │ puxados na direção do vetor de entrada │
-  │ Preservação de topologia     │ Função Gaussiana de decaimento          │
-  │                              │ Λ(d, σ) = exp(-d²/(2σ²))                │
-  │ Raio de vizinhança σ         │ Decai ao longo do treino (sigma_t =     │
-  │                              │ sigma_0 * (sigma_f/sigma_0)^(t/T))      │
-  │ Taxa de aprendizado α        │ Decai ao longo do treino (alpha_t =     │
-  │                              │ alpha_0 * (alpha_f/alpha_0)^(t/T))      │
-  │ Malha 4D                     │ 8 BiGRU4 dispostas em grid 2x2x2x2     │
-  │                              │ ( coordenadas (i,j,k,l), i,j,k,l∈{0,1}) │
-  └─────────────────────────────────────────────────────────────────────────┘
-
-COMO FUNCIONA
-─────────────
-
-1. Para cada vetor de entrada x (saída agregada dos u8cell_T), os 8 neurônios
-   BiGRU4 produzem 8 representações internas.
-2. O TransformerUnit (dimensão semântica) calcula scores de similaridade
-   aprendíveis entre x e cada neurônio.
-3. O neurônio com maior score é o BMU (Best-Matching Unit).
-4. A função Gaussiana Λ(d_bmu, σ) pondera cada neurônio pela distância 4D até
-   o BMU no grid 2x2x2x2.
-5. Os pesos dos neurônios são atualizados (em treino) na direção de x:
-      w_i ← w_i + α · Λ(d_i_bmu, σ) · (x - w_i)
-   onde α é a taxa de aprendizado e σ é o raio de vizinhança atuais.
-6. Em EVAL mode, o mapa é congelado — apenas a busca BMU é executada.
-7. A saída agregada combina:
-   - A representação do BMU (classe vencedora)
-   - Média ponderada por Λ (que preserva topologia)
-   - Cache aprendível cross-attention (compatibilidade V5)
-
-COMO LIDA COM DADOS NÃO-PREPARADOS PARA KOHONEN 4D
-──────────────────────────────────────────────────
-
-O usuário observou que os datasets não são preparados para essa arquitetura
-Kohonen 4D. Portanto:
-
-  a) Toda vez que um tensor/vetor de entrada NÃO tem a dimensão esperada para
-     um neurônio BiGRU4 (semântica + pesos), o OrqCell atribui valores
-     aleatórios aos parâmetros faltantes e prossegue. Isto é feito via
-     `_ensure_neuron_compat()` que verifica shape em runtime e faz pad/resize.
-  b) Em val mode (ou quando há punição), o OrqCell primeiro consulta o mapa
-     para ver se a informação já está presente (`map_contains_info()`).
-     Se não houver, atribui valores aleatórios controlados aos parâmetros
-     faltantes (magnitudes pequenas para não perturbar o mapa já treinado).
-  c) O BMU e os vizinhos são atualizados apenas em train mode. Em eval, o
-     mapa �� read-only.
-
-MONITORAMENTO V6
-────────────────
-
-A cada step de treino, expõe para o MetricsMonitor:
-  - sigma_t: raio de vizinhança atual
-  - alpha_t: taxa de aprendizado atual
-  - bmu_id: índice do neurônio vencedor (média do batch)
-  - bmu_diversity: nº de BMUs distintos no batch (alta = map cobrindo bem)
-  - neighborhood_spread: média ponderada de Λ (alta = update amplo)
-  - map_entropy: entropia da distribuição de BMUs (alta = uso balanceado)
-"""
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass, field
-from typing import Optional, Tuple, List, Dict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-# ============================================================================
-# Config
-# ============================================================================
-
-@dataclass
-class Kohonen4DConfig:
-    """Configuração do SOM 4D Kohonen.
-
-    Defaults conservadores para treinamento de 100 samples em CPU.
-    """
-    # Dimensões
-    n_neurons: int = 8                 # 8 neurônios em arranjo 4D
-    # grid_shape (2,2,2,1): 8 células em espaço 4D (4ª dimensão degenerada
-    # — apenas para permitir coordenação 4D; pode ser estendida para (2,2,2,2)
-    # com 16 neurônios no futuro, quando max_modules subir de 8 para 16).
-    grid_shape: Tuple[int, int, int, int] = (2, 2, 2, 1)
-    d_cache: int = 128                  # dimensão do cache aprendível
-    d_model_in: int = 128               # dimensão da entrada (saída u8cell_T)
-    nhead: int = 4
-    d_ff: int = 256
-    dropout: float = 0.1
-
-    # Schedule do SOM
-    sigma_0: float = 1.5                # raio inicial (cobre ~3 vizinhos no grid 2x2x2x2)
-    sigma_f: float = 0.3                # raio final (apenas BMU + vizinho imediato)
-    alpha_0: float = 0.5                # taxa de aprendizado inicial
-    alpha_f: float = 0.05               # taxa de aprendizado final
-    total_steps: int = 200              # T (passos para decair sigma_0→sigma_f)
-
-    # Controle de atualização
-    update_neighbors_in_eval: bool = False  # nunca atualizar em eval
-    random_pad_scale: float = 0.02     # escala para valores aleatórios de parâmetros faltantes
-    bmu_temp: float = 1.0               # temperatura do softmax de seleção de BMU
-
-    # ───────────────────────────────────────────────────────────────────────
-    # V6.1 — PGVector-like lookup thresholds
-    # ───────────────────────────────────────────────────────────────────────
-    # Quando o usuário treina com datasets não preparados para Kohonen 4D,
-    # o mapa precisa responder a uma pergunta: "este vetor de entrada já
-    # está representado no SOM?" A resposta pode ser:
-    #   - COMPLETE   : existe um neurônio com similaridade alta (≥ pgvector_complete_thr)
-    #                  E esse neurônio está "filled" (já recebeu update SOM)
-    #   - INCOMPLETE : existe um neurônio parecido, mas NÃO está filled
-    #                  (faltam elementos — os pesos EWC ainda não foram computados)
-    #   - MISSING    : nenhum neurônio tem similaridade ≥ pgvector_missing_thr
-    #                  (o mapa não tem aquela info de forma alguma)
-    pgvector_complete_thr: float = 0.6   # sim ≥ thr → COMPLETE
-    pgvector_missing_thr: float = 0.2    # sim < thr → MISSING
-    # (entre missing_thr e complete_thr → INCOMPLETE)
-
-
-# ============================================================================
-# Coordenadas 4D do grid 2x2x2x2
-# ============================================================================
-
-def _build_grid_coords(grid_shape: Tuple[int, int, int, int]) -> torch.Tensor:
-    """Constrói as coordenadas 4D de cada neurônio no grid.
-
-    Returns: (n_neurons, 4) — coordenadas inteiras de cada neurônio.
-    Para grid (2,2,2,2): 8 coordenadas [(i,j,k,l) for i,j,k,l in {0,1}]
-    """
-    coords = []
-    for i in range(grid_shape[0]):
-        for j in range(grid_shape[1]):
-            for k in range(grid_shape[2]):
-                for l in range(grid_shape[3]):
-                    coords.append([i, j, k, l])
-    return torch.tensor(coords, dtype=torch.float32)  # (n_neurons, 4)
-
-
-def _pairwise_grid_distances(coords: torch.Tensor) -> torch.Tensor:
-    """Distância euclidiana 4D entre todos os pares de neurônios.
-
-    Args: coords (n_neurons, 4)
-    Returns: (n_neurons, n_neurons) — dist[i,j] = ||coords[i] - coords[j]||_2
-    """
-    # ||c_i - c_j||^2 = ||c_i||^2 + ||c_j||^2 - 2 c_i·c_j
-    sq = (coords ** 2).sum(dim=-1)  # (n,)
-    d2 = sq.unsqueeze(0) + sq.unsqueeze(1) - 2 * coords @ coords.T
-    return torch.sqrt(d2.clamp(min=0.0) + 1e-12)
-
-
-# ============================================================================
-# OrqCell V6 — Kohonen 4D SOM
-# ============================================================================
-
-class OrqCell(nn.Module):
-    """Orquestrador com SOM 4D Kohonen (V6).
-
-    Substitui o OrqCell V5 (que era cross-attention puro) por um mapa
-    auto-organizável onde os 8 neurônios são as 8 saídas dos BiGRU4 já
-    processados pelo u8cell_T, e o TransformerUnit atua como dimensão
-    semântica aprendível que substitui a distância euclidiana clássica.
-
-    Args:
-        config: Kohonen4DConfig
-
-    Forward (train mode):
-        H: (batch, K, d_model_in) — K saídas dos K módulos u8cell_T
-        step: int — passo global de treino (para schedule de σ e α)
-        → o: (batch, d_cache)
-        → aux_dict com sigma_t, alpha_t, bmu_id, etc.
-
-    Forward (eval mode):
-        Mesma assinatura, mas sem atualizar pesos do mapa.
-    """
-    def __init__(self, config: Optional[Kohonen4DConfig] = None, **kwargs):
-        """Inicializa o OrqCell V6 Kohonen 4D SOM.
-
-        Backward-compatible com V5 API: aceita kwargs (cache_len, d_cache,
-        d_model_in, nhead, d_ff, dropout) e converte para Kohonen4DConfig.
-        Se `config` for fornecido, tem prioridade.
-        """
-        super().__init__()
-        # Se config não fornecido, construir a partir de kwargs (V5 compat)
-        # V6 always uses 8 neurons regardless of cache_len (cache_len is ignored
-        # — V5 used 16 positions but V6 Kohonen requires exactly 8 neurons
-        # arranged in 4D grid (2,2,2,1)).
-        if config is None:
-            d_cache = kwargs.get("d_cache", 128)
-            d_model_in = kwargs.get("d_model_in", 128)
-            nhead = kwargs.get("nhead", 4)
-            d_ff = kwargs.get("d_ff", 256)
-            dropout = kwargs.get("dropout", 0.1)
-            config = Kohonen4DConfig(
-                n_neurons=8,
-                grid_shape=(2, 2, 2, 1),
-                d_cache=d_cache,
-                d_model_in=d_model_in,
-                nhead=nhead,
-                d_ff=d_ff,
-                dropout=dropout,
-            )
-        self.cfg = config
-        assert self.cfg.n_neurons == 8, "V6 Kohonen 4D requer 8 neurônios (grid 2x2x2x1)"
-        assert math.prod(self.cfg.grid_shape) == 8, "grid_shape deve ter produto 8"
-
-        # ───────────────────────────────────────────────────────────────────
-        # Pesos do mapa: cache aprendível + proj_in
-        # ───────────────────────────────────────────────────────────────────
-        # Cada neurônio tem:
-        #   - Um vetor de cache (d_cache) — seu "peso próprio" no espaço semântico
-        #   - Uma projeção de entrada (d_model_in → d_cache)
-        self.neuron_cache = nn.Parameter(
-            torch.randn(self.cfg.n_neurons, self.cfg.d_cache) * 0.02
-        )
-        self.proj_in = nn.Linear(self.cfg.d_model_in, self.cfg.d_cache)
-
-        # ───────────────────────────────────────────────────────────────────
-        # Dimensão semântica: TransformerUnit (cross-attention com 8 neurônios)
-        # ───────────────────────────────────────────────────────────────────
-        # Em vez de computar a distância euclidiana ||x - w_i|| no espaço
-        # original, usamos a atenção do TransformerUnit como uma função de
-        # vizinhança APRENDÍVEL. Cada neurônio é uma posição (1 das 8) na
-        # sequência de attention. A atenção calcula score(x, neuron_i).
-        # Como o TransformerUnit compartilha parâmetros, todos os 8 neurônios
-        # usam o mesmo Q,K,V proj — a diferenciação entre neurônios vem do
-        # neuron_cache (que é o valor aprendido individualmente por neurônio).
-        self.cross_attn = nn.MultiheadAttention(
-            self.cfg.d_cache, self.cfg.nhead, dropout=self.cfg.dropout, batch_first=True
-        )
-
-        # FFN pre-norm (consistente com TransformerUnit)
-        self.linear1 = nn.Linear(self.cfg.d_cache, self.cfg.d_ff)
-        self.linear2 = nn.Linear(self.cfg.d_ff, self.cfg.d_cache)
-        self.norm1 = nn.LayerNorm(self.cfg.d_cache)
-        self.norm2 = nn.LayerNorm(self.cfg.d_cache)
-        self.dropout1 = nn.Dropout(self.cfg.dropout)
-        self.dropout2 = nn.Dropout(self.cfg.dropout)
-        self.dropout_ff = nn.Dropout(self.cfg.dropout)
-
-        # ───────────────────────────────────────────────────────────────────
-        # Grid 4D e distâncias pré-computadas (não aprendíveis)
-        # ───────────────────────────────────────────────────────────────────
-        grid_coords = _build_grid_coords(self.cfg.grid_shape)            # (8, 4)
-        grid_dist = _pairwise_grid_distances(grid_coords)                # (8, 8)
-        self.register_buffer("grid_coords", grid_coords, persistent=False)
-        self.register_buffer("grid_dist", grid_dist, persistent=False)
-
-        # ───────────────────────────────────────────────────────────────────
-        # Estado do SOM (não-parametros, mas monitorável)
-        # ───────────────────────────────────────────────────────────────────
-        # step_t é atualizado por update_step() ou passado em forward()
-        self.register_buffer("step_t", torch.tensor(0, dtype=torch.long), persistent=False)
-        self.register_buffer("sigma_t", torch.tensor(self.cfg.sigma_0, dtype=torch.float32),
-                             persistent=False)
-        self.register_buffer("alpha_t", torch.tensor(self.cfg.alpha_0, dtype=torch.float32),
-                             persistent=False)
-
-        # ───────────────────────────────────────────────────────────────────
-        # V6.1 — neuron_filled_mask: marca quais neurônios já receberam
-        # update SOM (BMU + vizinhos). Inicialmente todos False (mapa vazio).
-        # Usado por:
-        #   - pgvector_lookup(): distingue COMPLETE (filled + sim alta)
-        #     de INCOMPLETE (sim alta mas não filled)
-        #   - EWC V6.1: só computa Fisher para neurônios filled — não
-        #     "reutiliza" pesos que o SOM já tem (evita dupla penalização)
-        # ───────────────────────────────────────────────────────────────────
-        self.register_buffer(
-            "neuron_filled_mask",
-            torch.zeros(self.cfg.n_neurons, dtype=torch.bool),
-            persistent=False,
-        )
-
-        # Histórico de BMU (para monitoramento)
-        self.bmu_history: List[int] = []
-        self.bmu_history_max = 1000  # limita crescimento
-
-    # =====================================================================
-    # Schedule do SOM (σ e α)
-    # =====================================================================
-    def _schedule_sigma(self, t: int) -> float:
-        """σ_t = σ_0 * (σ_f/σ_0)^(t/T) — decaimento exponencial."""
-        T = max(1, self.cfg.total_steps)
-        t = max(0, min(t, T))
-        ratio = self.cfg.sigma_f / max(self.cfg.sigma_0, 1e-8)
-        return self.cfg.sigma_0 * (ratio ** (t / T))
-
-    def _schedule_alpha(self, t: int) -> float:
-        """α_t = α_0 * (α_f/α_0)^(t/T) — decaimento exponencial."""
-        T = max(1, self.cfg.total_steps)
-        t = max(0, min(t, T))
-        ratio = self.cfg.alpha_f / max(self.cfg.alpha_0, 1e-8)
-        return self.cfg.alpha_0 * (ratio ** (t / T))
-
-    def update_step(self, t: int) -> None:
-        """Atualiza step_t e recalcula sigma_t e alpha_t.
-
-        Chamar a cada step de treino para manter o schedule.
-        """
-        self.step_t.fill_(t)
-        self.sigma_t.fill_(self._schedule_sigma(t))
-        self.alpha_t.fill_(self._schedule_alpha(t))
-
-    # =====================================================================
-    # Função de vizinhança Gaussiana
-    # =====================================================================
-    def _neighborhood(self, bmu_idx: int, sigma: float) -> torch.Tensor:
-        """Λ(d_bmu, σ) = exp(-d²/(2σ²)) para todos os neurônios.
-
-        Args:
-            bmu_idx: índice do BMU
-            sigma: raio de vizinhança atual
-
-        Returns: (n_neurons,) — peso de cada neurônio (1.0 no BMU, decaindo)
-        """
-        d = self.grid_dist[bmu_idx]  # (n_neurons,)
-        return torch.exp(-(d ** 2) / (2.0 * sigma * sigma + 1e-12))
-
-    # =====================================================================
-    # Compatibilidade de shape (dados não preparados para Kohonen 4D)
-    # =====================================================================
-    def _ensure_neuron_compat(self, H: torch.Tensor) -> torch.Tensor:
-        """Garante que H tem dimensão compatível com d_model_in.
-
-        Se H não tem a dimensão esperada pelo mapa (datasets não preparados
-        para Kohonen 4D), faz pad/resize com valores aleatórios controlados.
-        """
-        if H.size(-1) == self.cfg.d_model_in:
-            return H
-        # Projeta H para d_model_in via pad/truncate com ruído controlado
-        target = self.cfg.d_model_in
-        current = H.size(-1)
-        if current < target:
-            # Pad com ruído gaussiano de magnitude pequena
-            pad_shape = H.shape[:-1] + (target - current,)
-            pad = torch.randn(pad_shape, device=H.device, dtype=H.dtype) * self.cfg.random_pad_scale
-            return torch.cat([H, pad], dim=-1)
-        else:
-            # Truncate (preserva os primeiros target componentes)
-            return H[..., :target]
-
-    def map_contains_info(self, x: torch.Tensor) -> torch.Tensor:
-        """Verifica se o mapa já possui a informação de x (batch-wise).
-
-        Computa a similaridade entre x (projetado) e cada neurônio do mapa.
-        Se max sim < threshold, considera que o mapa não tem aquela info.
-
-        Args:
-            x: (batch, d_model_in) ou (batch, K, d_model_in)
-
-        Returns: (batch,) — booleno: True se o mapa tem info de x
-        """
-        x = self._ensure_neuron_compat(x)
-        if x.dim() == 3:
-            # Pool sobre K módulos
-            x = x.mean(dim=1)
-        x_proj = self.proj_in(x)  # (batch, d_cache)
-        # Similaridade = cosseno entre x_proj e cada neuron_cache
-        x_norm = F.normalize(x_proj, dim=-1)                          # (batch, d_cache)
-        c_norm = F.normalize(self.neuron_cache, dim=-1)                # (8, d_cache)
-        sims = x_norm @ c_norm.T                                       # (batch, 8)
-        max_sim, _ = sims.max(dim=-1)                                  # (batch,)
-        # Threshold conservador: 0.5 (similaridade cosseno)
-        return max_sim >= 0.5
-
-    # =====================================================================
-    # V6.1 — PGVector-like lookup
-    # =====================================================================
-    def pgvector_lookup(self, x: torch.Tensor) -> Dict[str, Any]:
-        """Busca PGVector-like: verifica se o mapa SOM tem a info de x.
-
-        Distingue três estados (conforme requisitado pelo usuário):
-
-          COMPLETE   : existe um neurônio com sim ≥ pgvector_complete_thr
-                       E esse neurônio está "filled" (neuron_filled_mask=True)
-                       → info presente e consolidada pelo SOM + EWC
-          INCOMPLETE : existe um neurônio com sim ≥ pgvector_missing_thr,
-                       mas NÃO está filled (faltam elementos — pesos EWC
-                       ainda não foram computados para aquela região do mapa)
-                       → info parcialmente presente; o EWC deve completar
-          MISSING    : nenhum neurônio tem sim ≥ pgvector_missing_thr
-                       → info ausente; usar EWC com pesos aleatórios iniciais
-
-        O usuário observou: "em uma busca PGVector posso encontrar informações
-        incompletas (faltam elementos e no caso os pesos que devem ser
-        calculados por EWC) e completas". Este método implementa essa triagem.
-
-        Args:
-            x: (batch, d_model_in) ou (batch, K, d_model_in)
-
-        Returns: dict com:
-            - status: str — "COMPLETE" | "INCOMPLETE" | "MISSING"
-                (status do BMU do batch; se batch misto, usa o pior caso)
-            - status_per_sample: List[str] — status por sample do batch
-            - best_neuron_idx: int — índice do neurônio mais similar (média)
-            - best_sim: float — similaridade máxima (média do batch)
-            - coverage_ratio: float — fração de neurônios com sim ≥ missing_thr
-            - filled_ratio: float — fração de neurônios filled no mapa
-            - missing_neurons: List[int] — neurônios ainda não filled
-            - n_filled: int — número de neurônios filled
-        """
-        x = self._ensure_neuron_compat(x)
-        if x.dim() == 3:
-            x = x.mean(dim=1)
-        x_proj = self.proj_in(x)  # (batch, d_cache)
-        x_norm = F.normalize(x_proj, dim=-1)
-        c_norm = F.normalize(self.neuron_cache, dim=-1)
-        sims = x_norm @ c_norm.T  # (batch, 8)
-        max_sim, best_idx = sims.max(dim=-1)  # (batch,), (batch,)
-
-        # Status por sample
-        statuses: List[str] = []
-        for i in range(x.size(0)):
-            sim_i = float(max_sim[i].item())
-            bmu_i = int(best_idx[i].item())
-            if sim_i < self.cfg.pgvector_missing_thr:
-                statuses.append("MISSING")
-            elif sim_i >= self.cfg.pgvector_complete_thr and bool(self.neuron_filled_mask[bmu_i]):
-                statuses.append("COMPLETE")
-            else:
-                # entre os thresholds OU sim alta mas não filled
-                statuses.append("INCOMPLETE")
-
-        # Status do batch: pior caso (MISSING > INCOMPLETE > COMPLETE)
-        if "MISSING" in statuses:
-            batch_status = "MISSING"
-        elif "INCOMPLETE" in statuses:
-            batch_status = "INCOMPLETE"
-        else:
-            batch_status = "COMPLETE"
-
-        # Coverage ratio: fração de neurônios com sim ≥ missing_thr (média do batch)
-        coverage = float((sims >= self.cfg.pgvector_missing_thr).float().mean().item())
-        filled_ratio = float(self.neuron_filled_mask.float().mean().item())
-        missing_neurons = [
-            i for i in range(self.cfg.n_neurons)
-            if not bool(self.neuron_filled_mask[i])
-        ]
-
-        return {
-            "status": batch_status,
-            "status_per_sample": statuses,
-            "best_neuron_idx": int(best_idx.float().mean().item()),
-            "best_sim": float(max_sim.mean().item()),
-            "coverage_ratio": coverage,
-            "filled_ratio": filled_ratio,
-            "missing_neurons": missing_neurons,
-            "n_filled": int(self.neuron_filled_mask.sum().item()),
-        }
-
-    def fill_missing_neurons_random(self, scale: Optional[float] = None) -> int:
-        """V6.1: preenche neurônios não-filled com valores aleatórios.
-
-        Conforme user requirement: "se não houver atribuir valores aleatórios
-        aos parâmetros faltantes". Este método é chamado pelos níveis de
-        treinamento e hipótese quando pgvector_lookup retorna INCOMPLETE/MISSING
-        e o usuário quer garantir que o mapa tenha algo (mesmo aleatório) em
-        todas as posições — para que o forward não degenere.
-
-        Args:
-            scale: magnitude dos valores aleatórios (default: cfg.random_pad_scale)
-
-        Returns:
-            nº de neurônios preenchidos
-        """
-        if scale is None:
-            scale = self.cfg.random_pad_scale
-        n_filled_now = 0
-        with torch.no_grad():
-            for i in range(self.cfg.n_neurons):
-                if not bool(self.neuron_filled_mask[i]):
-                    # Atribui valores aleatórios controlados ao cache do neurônio
-                    self.neuron_cache[i].add_(
-                        torch.randn_like(self.neuron_cache[i]) * scale
-                    )
-                    self.neuron_filled_mask[i] = True
-                    n_filled_now += 1
-        return n_filled_now
-
-    def get_filled_neurons(self) -> List[int]:
-        """V6.1: retorna lista de índices de neurônios filled."""
-        return [i for i in range(self.cfg.n_neurons)
-                if bool(self.neuron_filled_mask[i])]
-
-    # =====================================================================
-    # Best Matching Unit
-    # =====================================================================
-    def _find_bmu(self, x_proj: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Encontra o BMU para cada sample do batch.
-
-        Args:
-            x_proj: (batch, d_cache) — entrada já projetada
-
-        Returns:
-            bmu_idx: (batch,) — índice do neurônio vencedor
-            bmu_sim: (batch,) — similaridade do BMU (para monitoramento)
-        """
-        # Similaridade cosseno entre x e cada neurônio (alta = match)
-        x_norm = F.normalize(x_proj, dim=-1)
-        c_norm = F.normalize(self.neuron_cache, dim=-1)
-        sims = x_norm @ c_norm.T / self.cfg.bmu_temp  # (batch, 8)
-        bmu_sim, bmu_idx = sims.max(dim=-1)
-        return bmu_idx, bmu_sim
-
-    # =====================================================================
-    # Forward principal
-    # =====================================================================
-    def forward(
-        self,
-        H: torch.Tensor,
-        step: Optional[int] = None,
-        return_aux: bool = False,
-    ) -> torch.Tensor:
-        """H: (batch, K, d_model_in) → o: (batch, d_cache).
-
-        Args:
-            H: (batch, K, d_model_in) — K saídas dos K módulos u8cell_T
-            step: passo global de treino (None = usar step_t interno)
-            return_aux: se True, retorna (o, aux_dict) com métricas V6
-
-        Returns:
-            o: (batch, d_cache) — saída agregada do orquestrador
-            aux_dict (se return_aux): sigma_t, alpha_t, bmu_id, etc.
-        """
-        # 1. Garantir compatibilidade de shape (datasets não preparados p/ Kohonen 4D)
-        H = self._ensure_neuron_compat(H)
-        batch_size, K, _ = H.shape
-
-        # 2. Atualizar schedule do SOM
-        if step is not None:
-            self.update_step(step)
-        sigma_t = float(self.sigma_t.item())
-        alpha_t = float(self.alpha_t.item())
-
-        # 3. Projeta H para o espaço do cache: (batch, K, d_cache)
-        H_proj = self.proj_in(H)
-
-        # 4. Pool de H_proj para encontrar o BMU: (batch, d_cache)
-        H_pooled = H_proj.mean(dim=1)
-
-        # 5. Encontrar BMU por sample
-        bmu_idx, bmu_sim = self._find_bmu(H_pooled)  # (batch,), (batch,)
-
-        # 6. Calcular peso de vizinhança para cada sample e neurônio
-        # Λ(d_bmu, σ) — para cada sample, pega a linha do BMU no grid_dist
-        # bmu_idx: (batch,) → neighborhood: (batch, n_neurons)
-        bmu_indices = bmu_idx.tolist()
-        neighborhoods = torch.stack(
-            [self._neighborhood(b, sigma_t) for b in bmu_indices], dim=0
-        )  # (batch, n_neurons)
-
-        # 7. Cross-attention (dimensão semântica do TransformerUnit)
-        #    queries = neuron_cache expandido, keys/values = H_proj
-        #    Isso substitui a simples distância euclidiana por uma função
-        #    de vizinhança APRENDÍVEL via attention.
-        cache = self.neuron_cache.unsqueeze(0).expand(batch_size, -1, -1)  # (batch, 8, d_cache)
-        normed_cache = self.norm1(cache)
-        normed_H = self.norm1(H_proj)
-        attn_out, _ = self.cross_attn(normed_cache, normed_H, normed_H, need_weights=False)
-        cache = cache + self.dropout1(attn_out)
-
-        # FFN pre-norm
-        normed = self.norm2(cache)
-        ff_out = self.linear2(self.dropout_ff(F.relu(self.linear1(normed))))
-        cache = cache + self.dropout2(ff_out)  # (batch, 8, d_cache)
-
-        # 8. Atualizar pesos do mapa SOM (apenas em treino)
-        #    NOTA: O update SOM é feito APÓS backward (em update_map_after_backward)
-        #    para não quebrar o autograd graph. Aqui apenas armazenamos o estado
-        #    necessário (H_pooled, neighborhoods, bmu_idx) para o update.
-        #    Em eval mode, o mapa está congelado (sem update).
-        if self.training:
-            # Cache do estado para update posterior
-            self._pending_update = (H_pooled.detach(), neighborhoods.detach(), bmu_idx.detach())
-            update_metrics = {
-                "update_magnitude": 0.0,  # será preenchido em update_map_after_backward
-                "bmu_drift": 0.0,
-                "map_frozen": False,
-            }
-        else:
-            self._pending_update = None
-            update_metrics = {
-                "update_magnitude": 0.0,
-                "bmu_drift": 0.0,
-                "map_frozen": True,
-            }
-
-        # 9. Saída agregada: combina BMU + média ponderada por Λ + cache attention
-        # Peso final do neurônio i = Λ(d_i_bmu, σ) normalizado
-        nh_norm = neighborhoods / (neighborhoods.sum(dim=-1, keepdim=True) + 1e-8)  # (batch, 8)
-        # Saída = soma ponderada do cache attention pelos pesos de vizinhança
-        o = (cache * nh_norm.unsqueeze(-1)).sum(dim=1)  # (batch, d_cache)
-
-        # 10. Atualiza histórico de BMU (para monitoramento)
-        if len(self.bmu_history) < self.bmu_history_max:
-            self.bmu_history.extend(bmu_indices[:min(8, batch_size)])
-
-        if not return_aux:
-            return o
-
-        # ───────────────────────────────────────────────────────────────────
-        # Métricas V6 para monitoramento
-        # ───────────────────────────────────────────────────────────────────
-        # BMU diversity: nº de BMUs distintos no batch
-        unique_bmu = len(set(bmu_indices))
-        # Neighborhood spread: média de Λ sobre o batch e neurônios
-        spread = float(neighborhoods.mean().item())
-        # Map entropy: entropia da distribuição de BMU no batch
-        bmu_counts = torch.zeros(self.cfg.n_neurons, device=bmu_idx.device)
-        for b in bmu_indices:
-            bmu_counts[b] += 1
-        bmu_probs = bmu_counts / (bmu_counts.sum() + 1e-8)
-        map_entropy = float(
-            -(bmu_probs * (bmu_probs + 1e-8).log()).sum().item()
-        )
-
-        aux = {
-            "sigma_t": sigma_t,
-            "alpha_t": alpha_t,
-            "step_t": int(self.step_t.item()),
-            "bmu_id_mean": float(bmu_idx.float().mean().item()),
-            "bmu_diversity": unique_bmu,
-            "bmu_sim_mean": float(bmu_sim.mean().item()),
-            "neighborhood_spread": spread,
-            "map_entropy": map_entropy,
-            "neighborhoods_sample": neighborhoods[0].detach().cpu().tolist(),
-            "n_filled": int(self.neuron_filled_mask.sum().item()),
-            "filled_ratio": float(self.neuron_filled_mask.float().mean().item()),
-            **update_metrics,
-        }
-        return o, aux
-
-    # =====================================================================
-    # Atualização SOM (neurônio vencedor + vizinhos) — chamada APÓS backward
-    # =====================================================================
-    def update_map_after_backward(self, alpha_t: Optional[float] = None) -> Dict[str, float]:
-        """Aplica o update SOM nos pesos do mapa. Chamar APÓS loss.backward().
-
-        Implementa a regra de Kohonen:
-            w_i ← w_i + α · Λ(d_i_bmu, σ) · (x - w_i)
-
-        Esta função DEVE ser chamada depois do backward pass, com
-        torch.no_grad() implícito, para não interferir com o autograd graph.
-
-        Args:
-            alpha_t: taxa de aprendizado (default: usar alpha_t interno)
-
-        Returns:
-            Dict com métricas da atualização
-        """
-        if not hasattr(self, "_pending_update") or self._pending_update is None:
-            return {"update_magnitude": 0.0, "bmu_drift": 0.0, "map_frozen": True}
-
-        H_pooled, neighborhoods, bmu_idx = self._pending_update
-        if alpha_t is None:
-            alpha_t = float(self.alpha_t.item())
-
-        with torch.no_grad():
-            # Δw_i = α · mean_batch( Λ_i · (x - w_i) )
-            x_exp = H_pooled.unsqueeze(1)          # (batch, 1, d_cache)
-            w_exp = self.neuron_cache.unsqueeze(0)  # (1, n_neurons, d_cache)
-            diff = x_exp - w_exp                     # (batch, n_neurons, d_cache)
-            nh = neighborhoods.unsqueeze(-1)         # (batch, n_neurons, 1)
-            delta = (nh * diff).mean(dim=0)          # (n_neurons, d_cache)
-            update_magnitude = float(delta.norm().item())
-            old_cache_norm = float(self.neuron_cache.norm().item())
-            self.neuron_cache.add_(alpha_t * delta)
-            new_cache_norm = float(self.neuron_cache.norm().item())
-            bmu_drift = abs(new_cache_norm - old_cache_norm)
-
-            # V6.1 — marca neurônios "filled": BMU + vizinhos com Λ > 0.1
-            # (próximo o suficiente do BMU para terem sido atualizados de
-            # forma significativa). Isto é usado por:
-            #   - pgvector_lookup(): distingue COMPLETE de INCOMPLETE
-            #   - EWC V6.1: só computa Fisher para neurônios filled
-            # O usuário pediu: "cada neurônio da grade tem peso próprio" —
-            # então cada neurônio preenchido individualmente.
-            mean_nh = neighborhoods.mean(dim=0)  # (n_neurons,)
-            newly_filled_mask = mean_nh > 0.1
-            self.neuron_filled_mask = self.neuron_filled_mask | newly_filled_mask
-
-        # Clear pending update
-        self._pending_update = None
-        return {
-            "update_magnitude": update_magnitude,
-            "bmu_drift": bmu_drift,
-            "map_frozen": False,
-            "n_filled": int(self.neuron_filled_mask.sum().item()),
-            "newly_filled": int(newly_filled_mask.sum().item()),
-        }
-
-    # =====================================================================
-    # Utilitário para EWC (V6)
-    # =====================================================================
-    def get_neuron_params(self, neuron_idx: int) -> Dict[str, torch.Tensor]:
-        """Retorna os parâmetros associados a um neurônio específico.
-
-        Usado pelo EWC V6 para computar Fisher Information por neurônio
-        (cada neurônio tem peso próprio — o usuário observou isso).
-        """
-        if neuron_idx < 0 or neuron_idx >= self.cfg.n_neurons:
-            raise IndexError(f"neuron_idx {neuron_idx} fora do range [0, {self.cfg.n_neurons})")
-        return {
-            "cache": self.neuron_cache[neuron_idx],
-            # O proj_in e cross_attn são compartilhados entre neurônios,
-            # mas cada neurônio tem seu próprio cache. O EWC deve penalizar
-            # a mudança no cache de cada neurônio individualmente.
-        }
-
-    def get_per_neuron_fisher_targets(self) -> List[str]:
-        """Lista de nomes de parâmetros do mapa (para EWC V6)."""
-        return ["neuron_cache"]
-
-
-__all__ = ["OrqCell", "Kohonen4DConfig"]