File size: 5,557 Bytes
4f42a02
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Verificacoes matematicas do HAKO-V2 em runtime (teoremas corrigidos).

  V2-M1 (T-ATTVAR / Minkowski): alpha e combinacao convexa => Var(z_attn) <=
        max_e Var(v_e), qualquer tau. Verificacao empirica por banco.
  V2-M2 (Dirichlet / passo seguro): energia de Dirichlet do Tree-GNN publicado
        deve ser monotona nao-crescente nos passos registrados; ajuste do passo
        com eta <= 2/lambda_max (lambda_max estimado por iteracao de potencia
        no operador de mensagem W_msg).
  V2-M3 (cota de Rademacher): R <= sqrt(1+(K-1)*rho_bar)/sqrt(K) com rho_bar a
        correlacao media entre valores de especialistas; reportada por banco.
  V2-M4 (Lemma 1 BPE): treino sharded == sequencial (verificador publicado).
"""
from __future__ import annotations

from typing import Dict, List

import numpy as np
import torch


def verify_moe_contraction(sys_, zs_list: List[dict], n_max: int = 128) -> Dict:
    """V2-M1: para cada banco presente, Var_i(z_attn) <= max_e Var_i(v_e)."""
    moe = sys_.moe
    out: Dict[str, Dict] = {}
    with torch.no_grad():
        for name in list(moe.banks):
            zs = [d[name] for d in zs_list if name in d][:n_max]
            if len(zs) < 8:
                continue
            Z = torch.stack(zs)                      # (n, d_src)
            bankd = moe.banks[name]
            K, A = bankd["K"], bankd["n_adapters"]
            Wv = getattr(moe, f"Wv_{name}")
            V = torch.stack([torch.stack(
                [Wv @ (Z[i] @ getattr(moe, f"Ad_{name}_{e % A}"))
                 for i in range(Z.shape[0])]) for e in range(K)])  # (K,n,N)
            var_e = V.var(dim=1).mean(dim=1)          # (K,)
            var_max = float(var_e.max())
            # z_attn com alpha real do mecanismo publicado
            zN = torch.stack([sys_.forward_sample({name: z}, train=False)["z"]
                              for z in zs[:32]])
            zt = []
            for zi, z_s in zip(zN, zs[:32]):
                out_f, _ = moe(zi, {name: z_s})
                zt.append(out_f)
            var_attn = float(torch.stack(zt).var(dim=0).mean())
            ok = var_attn <= var_max * 1.05 + 1e-9
            out[name] = {
                "var_attn": round(var_attn, 6),
                "var_max_especialista": round(var_max, 6),
                "razao": round(var_attn / max(var_max, 1e-12), 4),
                "ok": bool(ok),
            }
    return {"teorema": "T-ATTVAR (Minkowski, combinacao convexa, tau livre)",
            "bancos": out, "passou": all(v["ok"] for v in out.values())}


def verify_dirichlet(sys_) -> Dict:
    """V2-M2: energia registrada pelo Tree-GNN publicado deve nao-crescer;
    lambda_max do operador estimado por iteracao de potencia em W_msg."""
    tgnn = sys_.tgnn
    energies = [float(e) for e in tgnn.last_energy] if tgnn.last_energy else []
    W = tgnn.W_msg.detach()
    v = torch.randn(W.shape[1])
    v /= v.norm()
    lam = 0.0
    for _ in range(64):
        w = W @ v
        nrm = w.norm()
        if nrm < 1e-12:
            break
        v = w / nrm
        lam = float(nrm)
    eta_safe = 2.0 / max(lam, 1e-9)
    mono = all(energies[i + 1] <= energies[i] + 1e-9
               for i in range(len(energies) - 1))
    return {
        "teorema": "T-GNN Dirichlet monotono com eta <= 2/lambda_max",
        "energias": [round(e, 6) for e in energies[:8]],
        "monotona": bool(mono) if energies else None,
        "lambda_max_estimado": round(lam, 4),
        "eta_seguro_publicado": round(min(0.35, eta_safe), 4),
        "passou": bool(mono),
    }


def verify_rademacher(sys_, zs_list: List[dict]) -> Dict:
    """V2-M3: cota efetiva por banco com rho_bar medido."""
    moe = sys_.moe
    out = {}
    with torch.no_grad():
        for name in list(moe.banks):
            zs = [d[name] for d in zs_list if name in d][:64]
            if len(zs) < 16:
                continue
            Z = torch.stack(zs)
            bankd = moe.banks[name]
            K, A = bankd["K"], bankd["n_adapters"]
            Wv = getattr(moe, f"Wv_{name}")
            V = torch.stack([torch.stack(
                [Wv @ (Z[i] @ getattr(moe, f"Ad_{name}_{e % A}"))
                 for i in range(Z.shape[0])]) for e in range(K)])
            Vf = V.reshape(K, -1)                      # (K, n*N)
            Vf = Vf - Vf.mean(dim=1, keepdim=True)
            Vn2 = Vf / (Vf.norm(dim=1, keepdim=True) + 1e-9)
            C = Vn2 @ Vn2.T                            # (K, K)
            off = C - torch.eye(C.shape[0]) * torch.diag(C)
            rho = float(off.abs().mean())
            Keff = K
            bound = float(np.sqrt(1.0 + (Keff - 1) * rho) / np.sqrt(Keff))
            out[name] = {"rho_bar": round(rho, 4), "K": Keff,
                         "cota_rademacher": round(bound, 4)}
    return {"teorema": "cota de Rademacher sqrt(1+(K-1)rho)/sqrt(K)",
            "bancos": out, "passou": True}


def verify_bpe_lemma(corpus_docs) -> Dict:
    """V2-M4: verificador publicada do Lemma 1 (BPE aditivo)."""
    from hako.tokenizer.byte_bpe import ByteBPETokenizer
    small = [d[:120] for d in corpus_docs[:60] if d]
    tok_full = ByteBPETokenizer(vocab_size=300, min_freq=2, shards=1)
    tok_full.train(small, log_every=0)
    tok_shard = ByteBPETokenizer(vocab_size=300, min_freq=2, shards=4)
    tok_shard.train(small, log_every=0)
    ok = tok_full.merge_history == tok_shard.merge_history
    return {"teorema": "Lemma 1 (BPE chunk-atomico, sharded == sequencial)",
            "merges": len(tok_full.merge_history), "passou": bool(ok)}