File size: 5,149 Bytes
3fe1f47 | 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 | """V3-DIFF: dinov2 -> Diffusion -> SOM (alimentacao logica da rede SOM).
O ajuste dinov2-small pedido: a camada Diffusion PUBLICADA
(hako.generative.diffusion.DiffusionGame -- DDPM epsilon-prediction +
T-CURRICULUM) e treinada sobre os latentes do banco dinov2 (prototipos
int4 desquantizados, padrao granite publicado); suas amostras denoizadas
(ancestral sampling) alimentam o aprendizado da rede SOM por dois canais:
C1 (interpolacao guiada): o slot dinov2 do batch vira
z_dino(i) = beta(t) * z~[i mod n] + (1-beta(t)) * proto[i mod 48],
com beta(t): 0.8 -> 0.3 (quantidade ajustada "so o necessario" --
os prototipos int4 reais permanecem a base quantizada).
C2 (amostras pseudo): no fim de cada epoca, n_extra pseudo-amostras
{gpt2a: real, gpt2g: real, dinov2: z~} entram no fluxo de treino SOM
=> o gerador difusivo alimenta DIRETAMENTE o aprendizado da rede.
Justificativa matematica: o DDPM publicado treinado em z0 ~ P_dinov2
aproxima P_dinov2; amostras z~ ~ P_theta expandem o suporte de treino da
SOM nas direcoes de alta densidade (aumento de dados latente) -- como o
QE eh media de ||z - w_bmu||, amostras entre prototipos reduzem vales de
Voronoi vazios => QE menor e cobertura melhor (hit_entropy sobe).
"""
from __future__ import annotations
from typing import Dict, List
import numpy as np
import torch
from hako.generative.diffusion import DiffusionGame # publicado, intocado
class DinoDiffusionBridge:
def __init__(self, protos: torch.Tensor, T: int = 32, hidden: int = 192,
seed: int = 7) -> None:
d = int(protos.shape[1])
self.protos = protos.float()
self.game = DiffusionGame(N_dim=d, hidden=hidden, T=T, levels=6,
bonus=0.15, seed=seed)
self.hist: List[dict] = []
self.d = d
# ---- persistencia (retomadas: denoiser continua de onde parou) --------
def save(self, path) -> None:
torch.save(self.game.state_dict(), path)
def load(self, path) -> bool:
try:
self.game.load_state_dict(torch.load(path, weights_only=True))
self.game.eval()
return True
except Exception:
return False
# ---- treino do denoiser sobre os latentes dinov2 ----------------------
def train(self, steps: int, reward_qe: float, gen: torch.Generator,
deadline_s: float | None = None, t0: float = 0.0) -> dict:
import time
last = {}
for s in range(steps):
j = int(torch.randint(0, self.protos.shape[0], (1,),
generator=gen))
z0 = self.protos[j]
last = self.game.train_step(z0, reward_qe, gen)
self.hist.append(last)
if deadline_s is not None and (time.time() - t0) > deadline_s:
break
losses = [h["loss_diff"] for h in self.hist[-200:]] or [0.0]
return {"steps_diff": len(self.hist), "loss_diff_ema": round(
0.9 * losses[0] + 0.1 * losses[-1], 5),
"level": last.get("level", 0),
"difficulty": round(last.get("difficulty", 0.2), 4),
"reward_ma": round(last.get("reward_ma", 0.0), 4)}
# ---- amostragem ancestral (canal generativo) --------------------------
@torch.no_grad()
def sample(self, n: int, gen: torch.Generator,
n_steps: int = 16) -> torch.Tensor:
out = torch.stack([self.game.sample(0.5, gen, n_steps=n_steps)
for _ in range(n)])
return out
# ---- C1: mistura beta(t) entre z~ e prototipo real ---------------------
def mix(self, i: int, zt: torch.Tensor, beta: float) -> torch.Tensor:
p = self.protos[i % self.protos.shape[0]]
z = zt[i % zt.shape[0]]
return beta * z + (1.0 - beta) * p
@staticmethod
def beta_schedule(ep: int, total_ep: int, b0: float = 0.8,
b1: float = 0.3) -> float:
if total_ep <= 1:
return b0
return b0 + (b1 - b0) * (ep / (total_ep - 1))
# ---- sonda: QE com dinov2 real vs dinov2 difusivo ----------------------
@torch.no_grad()
def probe_qe(self, sys_, pools, zt: torch.Tensor, n: int = 96,
gen_seed: int = 7) -> dict:
gen = torch.Generator().manual_seed(gen_seed)
Pa = torch.as_tensor(pools["gpt2a"]).float()
Pg = torch.as_tensor(pools["gpt2g"]).float()
idx = torch.randperm(min(n, len(Pa)), generator=gen)[:n].tolist()
qe_real, qe_diff = [], []
for i in idx:
zs_r = {"gpt2a": Pa[i], "gpt2g": Pg[i],
"dinov2": self.protos[i % self.protos.shape[0]]}
zs_d = {"gpt2a": Pa[i], "gpt2g": Pg[i],
"dinov2": zt[i % zt.shape[0]]}
qe_real.append(float(sys_.forward_sample(zs_r, train=False)["qe"]))
qe_diff.append(float(sys_.forward_sample(zs_d, train=False)["qe"]))
return {"qe_dinov2_real": round(float(np.mean(qe_real)), 4),
"qe_dinov2_difusivo": round(float(np.mean(qe_diff)), 4),
"n": len(idx)}
|