V5: W8A8 aprimorado (α aprendível via STE) + bug fix apply_w8a8_to_model + 4 hiperparâmetros auto-ajustáveis (T, τ, λ_ent, init_gate) + OOM-Killer fixes + num_layers_hyp=8 fixo
Browse files- src/bigru_t/model/orq_cell.py +27 -408
src/bigru_t/model/orq_cell.py
CHANGED
|
@@ -1,320 +1,37 @@
|
|
| 1 |
-
"""orq_cell.py — OrqCell
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
Eixo 3 — grid Kohonen X (S_x nós, default 4)
|
| 14 |
-
Eixo 4 — grid Kohonen Y (S_y nós, default 4)
|
| 15 |
-
|
| 16 |
-
Total de protótipos = 8 × S_sem × S_x × S_y (default: 8 × 4 × 4 × 4 = 512)
|
| 17 |
-
|
| 18 |
-
Cada protótipo é um vetor ∈ R^{d_cache}. A topologia 4D define a vizinhança:
|
| 19 |
-
protótipos "próximos" no grid representam regiões semânticas próximas no
|
| 20 |
-
espaço latente das 8 BiGRU4 + TransformerUnit.
|
| 21 |
-
|
| 22 |
-
ALGORITMO DE TREINO (Kohonen, diferenciável via STE):
|
| 23 |
-
Para cada entrada h ∈ R^{d_cache} (saída projetada de H):
|
| 24 |
-
1. Computa distâncias: d_i = ||h - w_i||² para todos os protótipos i
|
| 25 |
-
2. BMU (Best Matching Unit): i* = argmin_i d_i
|
| 26 |
-
3. Vizinhança gaussiana: h_i = exp(-dist_grid(i, i*)² / (2σ²))
|
| 27 |
-
4. Atualização (STE para argmin): w_i ← w_i + α · h_i · (h - w_i)
|
| 28 |
-
onde α é a taxa de aprendizado (decrescente), σ é a largura de vizinhança
|
| 29 |
-
|
| 30 |
-
O argmin não é diferenciável — usamos STE (gradiente passa direto) e uma
|
| 31 |
-
loss auxiliar de quantização: L_q = Σ_i h_i · ||h - w_i||² (ponderada pela
|
| 32 |
-
vizinhança, soft winner-takes-all).
|
| 33 |
-
|
| 34 |
-
ÊNFASE NA VIZINHANÇA:
|
| 35 |
-
Em SOMs clássicos, σ decresce monotonicamente. Em V6, σ é **adaptativo**:
|
| 36 |
-
σ(t) = σ_min + (σ_max - σ_min) · exp(-t/τ_σ)
|
| 37 |
-
onde τ_σ é um hiperparâmetro aprendível (via MetaConfigurator V7).
|
| 38 |
-
Isso permite ao modelo decidir quão "loose" a vizinhança deve ser — em
|
| 39 |
-
fases de exploração (loss alta), σ grande (mais protótipos atualizados);
|
| 40 |
-
em fases de especialização (loss baixa), σ pequeno (apenas BMU).
|
| 41 |
-
|
| 42 |
-
INTEGRAÇÃO COM CROSS-ATTENTION:
|
| 43 |
-
O SOM atua como "roteador topológico": para cada posição l do cache,
|
| 44 |
-
computa-se o BMU e a vizinhança, e pondera-se a contribuição dos K módulos
|
| 45 |
-
via cross-attention. A saída agregada combina:
|
| 46 |
-
- Informação topológica (BMU + vizinhança gaussiana)
|
| 47 |
-
- Informação de atenção (cross-attention cache↔módulos)
|
| 48 |
-
|
| 49 |
-
REFERÊNCIAS:
|
| 50 |
-
- Kohonen, T. (1982). "Self-organized formation of topologically correct
|
| 51 |
-
feature maps". Biological Cybernetics, 43(1), 59-69.
|
| 52 |
-
- Kohonen, T. (2001). "Self-Organizing Maps" 3rd ed. Springer.
|
| 53 |
-
- Rojas, R. (2013). "Neural Networks" (cap. 17, SOM diferenciável).
|
| 54 |
-
- Van Hulle, M. M. (2012). "Self-organizing maps" (Handbook of Nat. Computing).
|
| 55 |
"""
|
| 56 |
from __future__ import annotations
|
| 57 |
|
| 58 |
-
import math
|
| 59 |
import torch
|
| 60 |
import torch.nn as nn
|
| 61 |
import torch.nn.functional as F
|
| 62 |
-
from typing import Optional
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
# ═══════════════════════════════════════════════════════════════════════════
|
| 66 |
-
# STE para argmin (Soft Winner-Takes-All com backward bem definido)
|
| 67 |
-
# ═══════════════════════════════════════════════════════════════════════════
|
| 68 |
-
class STEArgmin(torch.autograd.Function):
|
| 69 |
-
"""Straight-Through Estimator para argmin.
|
| 70 |
-
|
| 71 |
-
Forward: retorna one-hot do argmin (winner-takes-all)
|
| 72 |
-
Backward: gradiente passa direto para todos os tensores de entrada
|
| 73 |
-
(distribuído uniformemente — aproximação soft)
|
| 74 |
-
|
| 75 |
-
Isso permite que o gradiente flua para os protótipos via a loss de
|
| 76 |
-
quantização L_q = Σ_i h_i · ||h - w_i||², mesmo que o BMU seja argmin.
|
| 77 |
-
"""
|
| 78 |
-
|
| 79 |
-
@staticmethod
|
| 80 |
-
def forward(ctx, distances: torch.Tensor) -> torch.Tensor:
|
| 81 |
-
"""distances: (..., N_prototypes) → one_hot: (..., N_prototypes)."""
|
| 82 |
-
bmu_idx = distances.argmin(dim=-1) # (...)
|
| 83 |
-
one_hot = F.one_hot(bmu_idx, num_classes=distances.shape[-1]).to(
|
| 84 |
-
distances.dtype
|
| 85 |
-
) # (..., N_prototypes)
|
| 86 |
-
ctx.save_for_backward(distances, one_hot)
|
| 87 |
-
return one_hot
|
| 88 |
-
|
| 89 |
-
@staticmethod
|
| 90 |
-
def backward(ctx, grad_output):
|
| 91 |
-
distances, one_hot = ctx.saved_tensors
|
| 92 |
-
# STE: gradiente passa direto (distribuído proporcional ao 1/distance)
|
| 93 |
-
# Aproximação soft: gradientes maiores para protótipos mais próximos
|
| 94 |
-
# d/dw_i ||h - w_i||² = -2(h - w_i) → já flui naturalmente
|
| 95 |
-
return grad_output
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
def ste_argmin(distances: torch.Tensor) -> torch.Tensor:
|
| 99 |
-
"""Argmin com STE. Retorna one-hot do BMU."""
|
| 100 |
-
return STEArgmin.apply(distances)
|
| 101 |
-
|
| 102 |
-
|
| 103 |
-
# ═══════════════════════════════════════════════════════════════════════════
|
| 104 |
-
# KohonenSOM4D — Mapa Auto-Organizável 4D
|
| 105 |
-
# ═══════════════════════════════════════════════════════════════════════════
|
| 106 |
-
class KohonenSOM4D(nn.Module):
|
| 107 |
-
"""Mapa Auto-Organizável de Kohonen 4D para o OrqCell V6.
|
| 108 |
-
|
| 109 |
-
Arranjo quadridimensional:
|
| 110 |
-
Eixo 0: módulo BiGRU4 (n_bigru=8)
|
| 111 |
-
Eixo 1: dimensão semântica do TransformerUnit (S_sem)
|
| 112 |
-
Eixo 2: grid Kohonen X (S_x)
|
| 113 |
-
Eixo 3: grid Kohonen Y (S_y)
|
| 114 |
-
|
| 115 |
-
Protótipos: tensor (n_bigru, S_sem, S_x, S_y, d_cache) — flatten para
|
| 116 |
-
N_prototypes = n_bigru × S_sem × S_x × S_y.
|
| 117 |
-
|
| 118 |
-
Vizinhança: gaussiana 4D no grid (distância Euclidiana entre índices 4D).
|
| 119 |
-
|
| 120 |
-
Args:
|
| 121 |
-
n_bigru: nº de BiGRU4 (default 8 — fixo na arquitetura u8cell_T)
|
| 122 |
-
S_sem: dimensão semântica do TransformerUnit (default 4)
|
| 123 |
-
S_x: tamanho do grid Kohonen X (default 4)
|
| 124 |
-
S_y: tamanho do grid Kohonen Y (default 4)
|
| 125 |
-
d_cache: dimensão do cache/protótipos
|
| 126 |
-
sigma_min: σ mínimo (vizinhança mínima — default 0.5)
|
| 127 |
-
sigma_max: σ máximo (vizinhança inicial — default 2.0)
|
| 128 |
-
alpha_init: taxa de aprendizado inicial (default 0.1)
|
| 129 |
-
|
| 130 |
-
Forward:
|
| 131 |
-
h: (batch, d_cache) →
|
| 132 |
-
bmu_onehot: (batch, N_prototypes) — one-hot do BMU
|
| 133 |
-
neighborhood: (batch, N_prototypes) — pesos gaussianos da vizinhança
|
| 134 |
-
quant_loss: escalar — Σ_i h_i · ||h - w_i||² (loss auxiliar)
|
| 135 |
-
bmu_4d: (batch, 4) — índice 4D do BMU (para logging)
|
| 136 |
-
"""
|
| 137 |
-
|
| 138 |
-
def __init__(
|
| 139 |
-
self,
|
| 140 |
-
n_bigru: int = 8,
|
| 141 |
-
S_sem: int = 4,
|
| 142 |
-
S_x: int = 4,
|
| 143 |
-
S_y: int = 4,
|
| 144 |
-
d_cache: int = 128,
|
| 145 |
-
sigma_min: float = 0.5,
|
| 146 |
-
sigma_max: float = 2.0,
|
| 147 |
-
alpha_init: float = 0.1,
|
| 148 |
-
):
|
| 149 |
-
super().__init__()
|
| 150 |
-
self.n_bigru = n_bigru
|
| 151 |
-
self.S_sem = S_sem
|
| 152 |
-
self.S_x = S_x
|
| 153 |
-
self.S_y = S_y
|
| 154 |
-
self.d_cache = d_cache
|
| 155 |
-
self.N_prototypes = n_bigru * S_sem * S_x * S_y
|
| 156 |
-
|
| 157 |
-
# Protótipos: (n_bigru, S_sem, S_x, S_y, d_cache) — inicializados pequenos
|
| 158 |
-
self.prototypes = nn.Parameter(
|
| 159 |
-
torch.randn(n_bigru, S_sem, S_x, S_y, d_cache) * 0.02
|
| 160 |
-
)
|
| 161 |
|
| 162 |
-
# Coordenadas 4D do grid (para computar vizinhança)
|
| 163 |
-
# Pré-computado: shape (N_prototypes, 4)
|
| 164 |
-
coords = []
|
| 165 |
-
for b in range(n_bigru):
|
| 166 |
-
for s in range(S_sem):
|
| 167 |
-
for x in range(S_x):
|
| 168 |
-
for y in range(S_y):
|
| 169 |
-
coords.append([float(b), float(s), float(x), float(y)])
|
| 170 |
-
self.register_buffer(
|
| 171 |
-
"grid_coords", torch.tensor(coords, dtype=torch.float32)
|
| 172 |
-
) # (N_prototypes, 4)
|
| 173 |
|
| 174 |
-
# Hiperparâmetros de vizinhança (registrados como buffers, atualizados
|
| 175 |
-
# externamente pelo scheduler ou pelo MetaConfigurator)
|
| 176 |
-
self.register_buffer("sigma_min", torch.tensor(float(sigma_min)))
|
| 177 |
-
self.register_buffer("sigma_max", torch.tensor(float(sigma_max)))
|
| 178 |
-
self.register_buffer("alpha", torch.tensor(float(alpha_init)))
|
| 179 |
-
self.register_buffer("step_count", torch.tensor(0, dtype=torch.long))
|
| 180 |
-
# τ_σ: constante de tempo do decaimento de σ (aprendível no futuro)
|
| 181 |
-
self.register_buffer("tau_sigma", torch.tensor(100.0))
|
| 182 |
-
|
| 183 |
-
@property
|
| 184 |
-
def sigma_current(self) -> torch.Tensor:
|
| 185 |
-
"""σ atual = σ_min + (σ_max - σ_min) · exp(-t/τ_σ)."""
|
| 186 |
-
decay = torch.exp(-self.step_count.float() / self.tau_sigma)
|
| 187 |
-
return self.sigma_min + (self.sigma_max - self.sigma_min) * decay
|
| 188 |
-
|
| 189 |
-
@property
|
| 190 |
-
def alpha_current(self) -> torch.Tensor:
|
| 191 |
-
"""α atual = α_init · exp(-t/τ_α) com τ_α = 2·τ_σ."""
|
| 192 |
-
decay = torch.exp(-self.step_count.float() / (2.0 * self.tau_sigma))
|
| 193 |
-
return self.alpha * decay
|
| 194 |
-
|
| 195 |
-
def forward(
|
| 196 |
-
self,
|
| 197 |
-
h: torch.Tensor,
|
| 198 |
-
return_aux: bool = False,
|
| 199 |
-
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 200 |
-
"""Computa BMU, vizinhança e loss de quantização.
|
| 201 |
-
|
| 202 |
-
Args:
|
| 203 |
-
h: (batch, d_cache) — entrada (saída projetada de H)
|
| 204 |
-
return_aux: se True, retorna também bmu_4d e quant_loss detalhado
|
| 205 |
-
|
| 206 |
-
Returns:
|
| 207 |
-
bmu_onehot: (batch, N_prototypes) — one-hot do BMU (com STE)
|
| 208 |
-
neighborhood: (batch, N_prototypes) — pesos gaussianos
|
| 209 |
-
quant_loss: escalar — Σ_i h_i · ||h - w_i||² (auxiliary loss)
|
| 210 |
-
"""
|
| 211 |
-
batch_size = h.shape[0]
|
| 212 |
-
|
| 213 |
-
# Flatten protótipos: (N_prototypes, d_cache)
|
| 214 |
-
W = self.prototypes.reshape(self.N_prototypes, self.d_cache)
|
| 215 |
-
|
| 216 |
-
# Distâncias Euidianas ao quadrado: (batch, N_prototypes)
|
| 217 |
-
# ||h - w_i||² = ||h||² - 2·h·w_i + ||w_i||²
|
| 218 |
-
h_sq = (h ** 2).sum(dim=-1, keepdim=True) # (batch, 1)
|
| 219 |
-
w_sq = (W ** 2).sum(dim=-1, keepdim=True).t() # (1, N_prototypes)
|
| 220 |
-
cross = h @ W.t() # (batch, N_prototypes)
|
| 221 |
-
distances = h_sq - 2.0 * cross + w_sq # (batch, N_prototypes)
|
| 222 |
-
distances = distances.clamp(min=0.0) # estabilidade numérica
|
| 223 |
-
|
| 224 |
-
# BMU via argmin com STE
|
| 225 |
-
bmu_onehot = ste_argmin(distances) # (batch, N_prototypes)
|
| 226 |
-
|
| 227 |
-
# Vizinhança gaussiana 4D
|
| 228 |
-
# Para cada batch, pegar o BMU e computar pesos para todos os protótipos
|
| 229 |
-
bmu_idx = distances.argmin(dim=-1) # (batch,)
|
| 230 |
-
bmu_coords = self.grid_coords[bmu_idx] # (batch, 4)
|
| 231 |
-
|
| 232 |
-
# Distância 4D entre cada protótipo e o BMU
|
| 233 |
-
# grid_coords: (N_prototypes, 4), bmu_coords: (batch, 4)
|
| 234 |
-
# → (batch, N_prototypes, 4) via broadcasting
|
| 235 |
-
diff = self.grid_coords.unsqueeze(0) - bmu_coords.unsqueeze(1) # (batch, N, 4)
|
| 236 |
-
grid_dist_sq = (diff ** 2).sum(dim=-1) # (batch, N_prototypes)
|
| 237 |
-
|
| 238 |
-
sigma = self.sigma_current.clamp(min=1e-3)
|
| 239 |
-
neighborhood = torch.exp(-grid_dist_sq / (2.0 * sigma ** 2)) # (batch, N)
|
| 240 |
-
|
| 241 |
-
# Loss de quantização (soft winner-takes-all ponderado pela vizinhança)
|
| 242 |
-
# L_q = Σ_i h_i · ||h - w_i||² (média sobre batch)
|
| 243 |
-
# Já temos distances = ||h - w_i||²
|
| 244 |
-
quant_loss = (neighborhood * distances).sum(dim=-1).mean() # escalar
|
| 245 |
-
|
| 246 |
-
# Atualiza step_count (apenas em modo treino)
|
| 247 |
-
if self.training:
|
| 248 |
-
with torch.no_grad():
|
| 249 |
-
self.step_count += 1
|
| 250 |
-
|
| 251 |
-
if not return_aux:
|
| 252 |
-
return bmu_onehot, neighborhood, quant_loss
|
| 253 |
-
|
| 254 |
-
# Auxiliar para logging
|
| 255 |
-
bmu_4d = bmu_coords.long() # (batch, 4)
|
| 256 |
-
return bmu_onehot, neighborhood, quant_loss, bmu_4d
|
| 257 |
-
|
| 258 |
-
@torch.no_grad()
|
| 259 |
-
def update_prototypes_online(self, h: torch.Tensor, bmu_onehot, neighborhood):
|
| 260 |
-
"""Atualização online dos protótipos (estilo SOM clássico).
|
| 261 |
-
|
| 262 |
-
w_i ← w_i + α · h_i · (h - w_i)
|
| 263 |
-
|
| 264 |
-
Chamado externamente quando se quer fazer a atualização SOM explícita
|
| 265 |
-
(em adição à atualização por gradiente via quant_loss).
|
| 266 |
-
"""
|
| 267 |
-
alpha = self.alpha_current
|
| 268 |
-
W = self.prototypes.reshape(self.N_prototypes, self.d_cache)
|
| 269 |
-
# h: (batch, d_cache), neighborhood: (batch, N_prototypes)
|
| 270 |
-
# Δw_i = α · Σ_batch h_i,batch · (h_batch - w_i) / batch_size
|
| 271 |
-
delta = neighborhood.t() @ h # (N_prototypes, d_cache)
|
| 272 |
-
delta = delta - neighborhood.sum(dim=0, keepdim=True).t() * W # (N, d_cache)
|
| 273 |
-
W.add_(alpha * delta / max(h.shape[0], 1))
|
| 274 |
-
self.prototypes.data = W.reshape(
|
| 275 |
-
self.n_bigru, self.S_sem, self.S_x, self.S_y, self.d_cache
|
| 276 |
-
)
|
| 277 |
-
|
| 278 |
-
def extra_repr(self) -> str:
|
| 279 |
-
return (
|
| 280 |
-
f"n_bigru={self.n_bigru}, S_sem={self.S_sem}, S_x={self.S_x}, "
|
| 281 |
-
f"S_y={self.S_y}, d_cache={self.d_cache}, "
|
| 282 |
-
f"N_prototypes={self.N_prototypes}, "
|
| 283 |
-
f"sigma={float(self.sigma_current.item()):.3f}, "
|
| 284 |
-
f"alpha={float(self.alpha_current.item()):.4f}, "
|
| 285 |
-
f"step={int(self.step_count.item())}"
|
| 286 |
-
)
|
| 287 |
-
|
| 288 |
-
|
| 289 |
-
# ═══════════════════════════════════════════════════════════════════════════
|
| 290 |
-
# OrqCell V6 — com KohonenSOM4D + Cross-Attention
|
| 291 |
-
# ═══════════════════════════════════════════════════════════════════════════
|
| 292 |
class OrqCell(nn.Module):
|
| 293 |
-
"""
|
| 294 |
|
| 295 |
Args:
|
| 296 |
-
cache_len: nº de posições no cache aprendível (L)
|
| 297 |
-
|
| 298 |
-
d_cache: dimensão do cache/protótipos
|
| 299 |
d_model_in: dimensão de entrada (saída dos u8cell_T)
|
| 300 |
nhead: nº de cabeças da cross-attention
|
| 301 |
d_ff: dimensão interna do FFN
|
| 302 |
dropout: probabilidade de dropout
|
| 303 |
-
n_bigru: nº de BiGRU4 no SOM (default 8)
|
| 304 |
-
S_sem, S_x, S_y: dimensões do SOM 4D
|
| 305 |
-
use_som: se True (default V6), usa SOM 4D; se False, cache V5 (compat)
|
| 306 |
|
| 307 |
Forward:
|
| 308 |
H: (batch, K, d_model_in) – saídas ponderadas dos K módulos
|
| 309 |
→ o: (batch, d_cache)
|
| 310 |
-
→ (se return_aux=True) aux: dict com bmu_4d, quant_loss, neighborhood_entropy
|
| 311 |
-
|
| 312 |
-
V6 vs V5:
|
| 313 |
-
* Cache V5 = (L, d_cache) parâmetros livres, sem topologia
|
| 314 |
-
* Cache V6 = SOM 4D (8, S_sem, S_x, S_y, d_cache) com topologia Kohonen
|
| 315 |
-
* V6 adiciona loss de quantização L_q ao aux_dict
|
| 316 |
-
* V6 mantém cross-attention (cache=queries, módulos=keys/values)
|
| 317 |
-
* V6 computa BMU + vizinhança para cada batch
|
| 318 |
"""
|
| 319 |
|
| 320 |
def __init__(
|
|
@@ -325,46 +42,23 @@ class OrqCell(nn.Module):
|
|
| 325 |
nhead: int,
|
| 326 |
d_ff: int,
|
| 327 |
dropout: float = 0.1,
|
| 328 |
-
n_bigru: int = 8,
|
| 329 |
-
S_sem: int = 4,
|
| 330 |
-
S_x: int = 4,
|
| 331 |
-
S_y: int = 4,
|
| 332 |
-
use_som: bool = True,
|
| 333 |
):
|
| 334 |
super().__init__()
|
| 335 |
self.cache_len = cache_len
|
| 336 |
self.d_cache = d_cache
|
| 337 |
-
|
|
|
|
|
|
|
| 338 |
|
| 339 |
# Projeção da entrada para o espaço do cache
|
| 340 |
self.proj_in = nn.Linear(d_model_in, d_cache)
|
| 341 |
|
| 342 |
-
if use_som:
|
| 343 |
-
# V6: SOM 4D de Kohonen como cache topológico
|
| 344 |
-
self.som = KohonenSOM4D(
|
| 345 |
-
n_bigru=n_bigru,
|
| 346 |
-
S_sem=S_sem,
|
| 347 |
-
S_x=S_x,
|
| 348 |
-
S_y=S_y,
|
| 349 |
-
d_cache=d_cache,
|
| 350 |
-
)
|
| 351 |
-
# Projeção do SOM (N_prototypes, d_cache) → (L, d_cache) para
|
| 352 |
-
# compatibilidade com a cross-attention (queries = L posições)
|
| 353 |
-
self.som_to_cache = nn.Linear(self.som.N_prototypes, cache_len, bias=False)
|
| 354 |
-
# Inicialização identidade-like (primeiras L linhas de N_prototypes)
|
| 355 |
-
with torch.no_grad():
|
| 356 |
-
eye = torch.eye(cache_len, self.som.N_prototypes)
|
| 357 |
-
self.som_to_cache.weight.copy_(eye * 0.1)
|
| 358 |
-
else:
|
| 359 |
-
# V5: cache aprendível simples (compat)
|
| 360 |
-
self.cache = nn.Parameter(torch.randn(cache_len, d_cache) * 0.02)
|
| 361 |
-
|
| 362 |
# Cross-attention: queries = cache, keys/values = H_proj
|
| 363 |
self.cross_attn = nn.MultiheadAttention(
|
| 364 |
d_cache, nhead, dropout=dropout, batch_first=True
|
| 365 |
)
|
| 366 |
|
| 367 |
-
# FFN (pre-norm style)
|
| 368 |
self.linear1 = nn.Linear(d_cache, d_ff)
|
| 369 |
self.linear2 = nn.Linear(d_ff, d_cache)
|
| 370 |
self.norm1 = nn.LayerNorm(d_cache)
|
|
@@ -373,59 +67,21 @@ class OrqCell(nn.Module):
|
|
| 373 |
self.dropout2 = nn.Dropout(dropout)
|
| 374 |
self.dropout_ff = nn.Dropout(dropout)
|
| 375 |
|
| 376 |
-
def forward(
|
| 377 |
-
|
| 378 |
-
H: torch.Tensor,
|
| 379 |
-
return_aux: bool = False,
|
| 380 |
-
) -> torch.Tensor:
|
| 381 |
-
"""H: (batch, K, d_model_in) → o: (batch, d_cache).
|
| 382 |
-
|
| 383 |
-
Se return_aux=True, retorna (o, aux_dict) com métricas do SOM V6.
|
| 384 |
-
"""
|
| 385 |
batch_size = H.size(0)
|
| 386 |
|
| 387 |
# Projeta H para o espaço do cache
|
| 388 |
H_proj = self.proj_in(H) # (batch, K, d_cache)
|
| 389 |
|
| 390 |
-
#
|
| 391 |
-
|
| 392 |
-
# =====================================================================
|
| 393 |
-
# h_mean: representação agregada de H para alimentar o SOM
|
| 394 |
-
h_mean = H_proj.mean(dim=1) # (batch, d_cache)
|
| 395 |
-
|
| 396 |
-
if self.use_som:
|
| 397 |
-
# SOM forward: retorna bmu_onehot, neighborhood, quant_loss, [bmu_4d]
|
| 398 |
-
som_out = self.som(h_mean, return_aux=True)
|
| 399 |
-
bmu_onehot, neighborhood, quant_loss, bmu_4d = som_out
|
| 400 |
-
|
| 401 |
-
# Reconstrói cache topológico:
|
| 402 |
-
# Para cada batch, o cache é uma combinação linear dos protótipos,
|
| 403 |
-
# ponderada pela vizinhança gaussiana do BMU.
|
| 404 |
-
# cache_b = Σ_i neighborhood[b,i] · W[i] → (batch, L, d_cache)
|
| 405 |
-
W_flat = self.som.prototypes.reshape(
|
| 406 |
-
self.som.N_prototypes, self.d_cache
|
| 407 |
-
) # (N_prototypes, d_cache)
|
| 408 |
-
# cache_b: (batch, N_prototypes, d_cache) via broadcasting
|
| 409 |
-
cache_topo = neighborhood.unsqueeze(-1) * W_flat.unsqueeze(0)
|
| 410 |
-
# Reduz para (batch, L, d_cache) via projeção linear
|
| 411 |
-
# cache_topo: (batch, N_prototypes, d_cache) → transp → (batch, d_cache, N_prototypes)
|
| 412 |
-
# → som_to_cache → (batch, d_cache, L) → transp → (batch, L, d_cache)
|
| 413 |
-
cache = self.som_to_cache(cache_topo.transpose(1, 2)).transpose(1, 2)
|
| 414 |
-
else:
|
| 415 |
-
# V5 compat: cache é parâmetro (L, d_cache), expande para batch
|
| 416 |
-
cache = self.cache.unsqueeze(0).expand(batch_size, -1, -1)
|
| 417 |
-
quant_loss = torch.tensor(0.0, device=H.device)
|
| 418 |
-
bmu_4d = None
|
| 419 |
-
neighborhood = None
|
| 420 |
|
| 421 |
-
# =====================================================================
|
| 422 |
# Cross-attention: queries = cache, keys/values = H_proj
|
| 423 |
-
#
|
| 424 |
normed_cache = self.norm1(cache)
|
| 425 |
-
normed_H = self.norm1(H_proj)
|
| 426 |
-
attn_out, _ = self.cross_attn(
|
| 427 |
-
normed_cache, normed_H, normed_H, need_weights=False
|
| 428 |
-
)
|
| 429 |
cache = cache + self.dropout1(attn_out)
|
| 430 |
|
| 431 |
# FFN pre-norm
|
|
@@ -435,41 +91,4 @@ class OrqCell(nn.Module):
|
|
| 435 |
|
| 436 |
# Saída = média sobre as L posições do cache
|
| 437 |
o = cache.mean(dim=1) # (batch, d_cache)
|
| 438 |
-
|
| 439 |
-
if not return_aux:
|
| 440 |
-
return o
|
| 441 |
-
|
| 442 |
-
# =====================================================================
|
| 443 |
-
# V6: Métricas auxiliares do SOM
|
| 444 |
-
# =====================================================================
|
| 445 |
-
aux = {
|
| 446 |
-
"som_quant_loss": float(quant_loss.item())
|
| 447 |
-
if isinstance(quant_loss, torch.Tensor)
|
| 448 |
-
else float(quant_loss),
|
| 449 |
-
}
|
| 450 |
-
if self.use_som:
|
| 451 |
-
aux["som_sigma"] = float(self.som.sigma_current.item())
|
| 452 |
-
aux["som_alpha"] = float(self.som.alpha_current.item())
|
| 453 |
-
aux["som_step"] = int(self.som.step_count.item())
|
| 454 |
-
aux["som_n_prototypes"] = self.som.N_prototypes
|
| 455 |
-
# Entropia da vizinhança (mede quão espalhada está a ativação)
|
| 456 |
-
if neighborhood is not None:
|
| 457 |
-
# Normaliza neighborhood para distribuição de probabilidade
|
| 458 |
-
neigh_norm = neighborhood / (neighborhood.sum(dim=-1, keepdim=True) + 1e-8)
|
| 459 |
-
entropy = -(neigh_norm * (neigh_norm + 1e-8).log()).sum(dim=-1).mean()
|
| 460 |
-
aux["som_neighborhood_entropy"] = float(entropy.item())
|
| 461 |
-
# Concentração do BMU (quanto menor, mais concentrado)
|
| 462 |
-
aux["som_bmu_concentration"] = float(
|
| 463 |
-
bmu_onehot.sum(dim=-1).mean().item()
|
| 464 |
-
)
|
| 465 |
-
if bmu_4d is not None:
|
| 466 |
-
# Distribuição dos BMUs por eixo (para logging)
|
| 467 |
-
aux["som_bmu_axis0_mean"] = float(bmu_4d[:, 0].float().mean().item())
|
| 468 |
-
aux["som_bmu_axis1_mean"] = float(bmu_4d[:, 1].float().mean().item())
|
| 469 |
-
aux["som_bmu_axis2_mean"] = float(bmu_4d[:, 2].float().mean().item())
|
| 470 |
-
aux["som_bmu_axis3_mean"] = float(bmu_4d[:, 3].float().mean().item())
|
| 471 |
-
|
| 472 |
-
return o, aux
|
| 473 |
-
|
| 474 |
-
|
| 475 |
-
__all__ = ["OrqCell", "KohonenSOM4D", "ste_argmin", "STEArgmin"]
|
|
|
|
| 1 |
+
"""orq_cell.py — OrqCell: orquestrador com cache aprendível + cross-attention.
|
| 2 |
+
|
| 3 |
+
Reaproveita a filosofia do orquestrador cooperativo do v13.9.2, mas:
|
| 4 |
+
- Cache aprendível (L, d_cache) substitui o "estado de trust" do v13.9.2
|
| 5 |
+
- Cross-attention (cache = queries, módulos = keys/values) substitui o
|
| 6 |
+
gating por Confidence Gate
|
| 7 |
+
- Média sobre as L posições do cache como saída agregada
|
| 8 |
+
|
| 9 |
+
O cache é um "roteador aprendível" que decide quais módulos ativos devem
|
| 10 |
+
contribuir para a representação final. Diferente do TrustNetwork do v13.9.2
|
| 11 |
+
(TD-learning com softmax norm), o OrqCell usa attention pura — a seleção
|
| 12 |
+
final dos módulos ativos é feita pelo Lema 1 no UnifiedModel.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 13 |
"""
|
| 14 |
from __future__ import annotations
|
| 15 |
|
|
|
|
| 16 |
import torch
|
| 17 |
import torch.nn as nn
|
| 18 |
import torch.nn.functional as F
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 20 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 21 |
class OrqCell(nn.Module):
|
| 22 |
+
"""Orquestrador com cache aprendível e cross-attention.
|
| 23 |
|
| 24 |
Args:
|
| 25 |
+
cache_len: nº de posições no cache aprendível (L)
|
| 26 |
+
d_cache: dimensão do cache
|
|
|
|
| 27 |
d_model_in: dimensão de entrada (saída dos u8cell_T)
|
| 28 |
nhead: nº de cabeças da cross-attention
|
| 29 |
d_ff: dimensão interna do FFN
|
| 30 |
dropout: probabilidade de dropout
|
|
|
|
|
|
|
|
|
|
| 31 |
|
| 32 |
Forward:
|
| 33 |
H: (batch, K, d_model_in) – saídas ponderadas dos K módulos
|
| 34 |
→ o: (batch, d_cache)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 35 |
"""
|
| 36 |
|
| 37 |
def __init__(
|
|
|
|
| 42 |
nhead: int,
|
| 43 |
d_ff: int,
|
| 44 |
dropout: float = 0.1,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 45 |
):
|
| 46 |
super().__init__()
|
| 47 |
self.cache_len = cache_len
|
| 48 |
self.d_cache = d_cache
|
| 49 |
+
|
| 50 |
+
# Cache aprendível (L, d_cache) — inicializado pequeno para estabilidade
|
| 51 |
+
self.cache = nn.Parameter(torch.randn(cache_len, d_cache) * 0.02)
|
| 52 |
|
| 53 |
# Projeção da entrada para o espaço do cache
|
| 54 |
self.proj_in = nn.Linear(d_model_in, d_cache)
|
| 55 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
# Cross-attention: queries = cache, keys/values = H_proj
|
| 57 |
self.cross_attn = nn.MultiheadAttention(
|
| 58 |
d_cache, nhead, dropout=dropout, batch_first=True
|
| 59 |
)
|
| 60 |
|
| 61 |
+
# FFN (pre-norm style, consistente com TransformerUnit)
|
| 62 |
self.linear1 = nn.Linear(d_cache, d_ff)
|
| 63 |
self.linear2 = nn.Linear(d_ff, d_cache)
|
| 64 |
self.norm1 = nn.LayerNorm(d_cache)
|
|
|
|
| 67 |
self.dropout2 = nn.Dropout(dropout)
|
| 68 |
self.dropout_ff = nn.Dropout(dropout)
|
| 69 |
|
| 70 |
+
def forward(self, H: torch.Tensor) -> torch.Tensor:
|
| 71 |
+
"""H: (batch, K, d_model_in) → o: (batch, d_cache)"""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 72 |
batch_size = H.size(0)
|
| 73 |
|
| 74 |
# Projeta H para o espaço do cache
|
| 75 |
H_proj = self.proj_in(H) # (batch, K, d_cache)
|
| 76 |
|
| 77 |
+
# Expande cache para o batch
|
| 78 |
+
cache = self.cache.unsqueeze(0).expand(batch_size, -1, -1) # (batch, L, d_cache)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 79 |
|
|
|
|
| 80 |
# Cross-attention: queries = cache, keys/values = H_proj
|
| 81 |
+
# Pre-norm
|
| 82 |
normed_cache = self.norm1(cache)
|
| 83 |
+
normed_H = self.norm1(H_proj) # compartilha norm1 (mesma dimensão)
|
| 84 |
+
attn_out, _ = self.cross_attn(normed_cache, normed_H, normed_H, need_weights=False)
|
|
|
|
|
|
|
| 85 |
cache = cache + self.dropout1(attn_out)
|
| 86 |
|
| 87 |
# FFN pre-norm
|
|
|
|
| 91 |
|
| 92 |
# Saída = média sobre as L posições do cache
|
| 93 |
o = cache.mean(dim=1) # (batch, d_cache)
|
| 94 |
+
return o
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|