PowerMachine commited on
Commit
70ae8f3
·
verified ·
1 Parent(s): 7e0bd86

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
Files changed (1) hide show
  1. 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 V6: Orquestrador com Mapa Auto-Organizável de Kohonen 4D.
2
-
3
- ═══════════════════════════════════════════════════════════════════════════════
4
- V6 — MELHORIA MATEMÁTICA FUNDAMENTAL (vs V5)
5
- ═══════════════════════════════════════════════════════════════════════════════
6
-
7
- Em V5, o OrqCell tinha um cache aprendível (L, d_cache) genérico, sem estrutura
8
- topológica. Em V6, o cache é um **Mapa Auto-Organizável de Kohonen (SOM) 4D**:
9
-
10
- ARRANJO QUADRIDIMENSIONAL:
11
- Eixo 1 — módulo BiGRU4 (8 nós, um por BiGRU4 paralelo dentro de cada u8cell_T)
12
- Eixo 2 — dimensão semântica do TransformerUnit (S_sem nós, default 4)
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
- """OrqCell V6: Orquestrador com SOM 4D de Kohonen + cross-attention.
294
 
295
  Args:
296
- cache_len: nº de posições no cache aprendível (L) — mantido para
297
- compatibilidade, mas em V6 o cache é o SOM 4D.
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
- self.use_som = use_som
 
 
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
- self,
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
- # V6: Computa cache via SOM 4D
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