File size: 6,439 Bytes
908fd1a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """VXOneNet (spec F2 §2): proyección + embeddings de texto + tronco + cabezas.
Estructura por bloque: Linear → LayerNorm → GELU → Dropout(0.1); bloque 1
mapea R^{d_in} → R^w, bloques 2..L son w→w con residual h_b = h_{b-1} +
bloque_b(h_{b-1}). LayerNorm (no BatchNorm): streaming por shards, batch-1 en
inferencia y drift temporal (F2 §2.1).
Cabezas por dataset (F2 §3.1):
- UNSW: noul(1) + choice_family(10) + corn(3 condicionales)
- CIC: noul(1) + choice_family(8|6) + corn + choice_ddos(4) [solo strat]
- EMBER: noul(1) + choice_family(11, unknown entrenado)
Dims congeladas (§2.2) y conteo de parámetros verificado (§2.3): los tests
comparan contra la tabla exacta de la spec.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
from torch import nn
# ── Dims congeladas por escala (F2 §2.2) ────────────────────────────────
SCALES: dict[str, dict[str, int]] = {
"nano": {"L": 2, "w_unsw": 224, "w_cic": 224, "w_ember": 192,
"d0_unsw": 112, "d0_cic": 112, "d0_ember": 20, "e": 16},
"small": {"L": 3, "w_unsw": 576, "w_cic": 576, "w_ember": 576,
"d0_unsw": 288, "d0_cic": 288, "d0_ember": 96, "e": 32},
"base": {"L": 4, "w_unsw": 1280, "w_cic": 1280, "w_ember": 1280,
"d0_unsw": 512, "d0_cic": 512, "d0_ember": 384, "e": 64},
}
DROPOUT = 0.1
@dataclass(frozen=True)
class HeadSpec:
"""Cabezas activas de una corrida (nombres congelados)."""
n_candidates: int # choice_family (K)
ddos_candidates: int = 0 # 0 = head ddos ausente
has_score: bool = True # False en EMBER (S6)
@property
def names(self) -> tuple[str, ...]:
ns = ["noul", "choice"]
if self.has_score:
ns.append("score")
if self.ddos_candidates:
ns.append("ddos")
return tuple(ns)
def head_spec_for(source: str, scale: str, split_scheme: str,
n_candidates: int) -> HeadSpec:
"""Candidate set por corrida (D12): familias con ≥1 fila en train.
La verificación de n_candidates ocurre en data.py; aquí solo se tipa.
"""
if source == "unsw":
return HeadSpec(n_candidates=n_candidates)
if source == "cic":
return HeadSpec(n_candidates=n_candidates,
ddos_candidates=4 if split_scheme == "strat" else 0)
if source == "ember":
return HeadSpec(n_candidates=n_candidates, has_score=False)
raise ValueError(f"source desconocido: {source!r}")
class _Block(nn.Module):
def __init__(self, d_in: int, w: int) -> None:
super().__init__()
self.lin = nn.Linear(d_in, w)
self.ln = nn.LayerNorm(w)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return nn.functional.dropout(
nn.functional.gelu(self.ln(self.lin(x))), DROPOUT,
training=self.training,
)
class VXOneNet(nn.Module):
"""Motor VX-One v0.1. d_num incluye el flag de fila (d+1 ya aplicado
por el feature stage: la entrada de P0 es [u ; r] ∈ R^{d+1})."""
def __init__(
self,
source: str,
scale: str,
heads: HeadSpec,
d_num: int, # d canónico (sin flag); P0 entra con d+1
vocab_sizes: dict[str, int] | None = None, # {campo: V+2} (solo UNSW)
) -> None:
super().__init__()
cfg = SCALES[scale]
self.source = source
self.scale = scale
self.heads = heads
if source == "ember":
w, d0, L = cfg["w_ember"], cfg["d0_ember"], cfg["L"]
else:
w = cfg["w_unsw"] if source == "unsw" else cfg["w_cic"]
d0 = cfg["d0_unsw"] if source == "unsw" else cfg["d0_cic"]
L = cfg["L"]
e = cfg["e"]
self.vocab_sizes = vocab_sizes or {}
m = len(self.vocab_sizes)
# P0: Linear((d+1) → d0) — el +1 es el flag de fila (§2.3)
self.proj = nn.Linear(d_num + 1, d0)
# Embeddings lookup: una tabla por campo, (V+2) × e
self.emb = nn.ModuleDict(
{f: nn.Embedding(vs, e) for f, vs in self.vocab_sizes.items()}
)
d_in = d0 + m * e
# Tronco: bloque 1 d_in→w, bloques 2..L w→w con residual
self.blocks = nn.ModuleList([_Block(d_in, w)] + [_Block(w, w) for _ in range(L - 1)])
# Cabezas (§2.3): Noul w+1; Choice wK+K; CORN w+3; ddos 4w+4
self.noul = nn.Linear(w, 1)
self.choice = nn.Linear(w, heads.n_candidates)
self.has_score = heads.has_score
if heads.has_score:
self.corn_u = nn.Linear(w, 1, bias=False) # u^T h
self.corn_b = nn.Parameter(torch.tensor([1.0, 0.0, -1.0])) # init §2.1
self.has_ddos = heads.ddos_candidates > 0
if self.has_ddos:
self.ddos = nn.Linear(w, heads.ddos_candidates)
# ── forward ──────────────────────────────────────────────────────
def trunk(self, u_r: torch.Tensor, txt_slots: dict[str, torch.Tensor] | None = None) -> torch.Tensor:
h = self.proj(u_r)
if self.emb:
parts = [h]
for f, tab in self.emb.items():
parts.append(tab(txt_slots[f]))
h = torch.cat(parts, dim=1)
for i, block in enumerate(self.blocks):
h = h + block(h) if i > 0 else block(h)
return h
def forward(self, u_r: torch.Tensor, txt_slots: dict[str, torch.Tensor] | None = None) -> dict[str, torch.Tensor]:
"""u_r float32 [B, d+1]; txt_slots {campo: int64 [B]} (solo UNSW)."""
h = self.trunk(u_r, txt_slots)
out: dict[str, torch.Tensor] = {
"noul": self.noul(h).squeeze(-1), # logits [B]
"choice": self.choice(h), # logits [B, K]
}
if self.has_score:
out["corn_logits"] = (self.corn_u(h).squeeze(-1).unsqueeze(1)
+ self.corn_b.unsqueeze(0)) # [B, 3]
if self.has_ddos:
out["ddos"] = self.ddos(h) # logits [B, 4]
return out
# ── conteo de parámetros (contrato §2.3; Apéndice B) ─────────────
def n_parameters(self) -> int:
return sum(p.numel() for p in self.parameters())
|