File size: 6,081 Bytes
463d94c | 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 | #!/usr/bin/env python3
"""CPU audit using actually trained score networks.
The earlier audits were analytic path families. This producer adds a small,
fully deterministic denoising-score model: fixed Gaussian-mixture data,
fixed-noise DSM targets, three dimensions, three noise levels, and three
independent seeds. It reports the learned-vs-analytic score energy and feeds
the measured per-generation errors into the exact fresh-data recurrence.
It is deliberately a scope audit, not a claim that a tiny network replaces
the paper's asymptotic stochastic proof.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import numpy as np
import torch
from torch import nn
ROOT = Path(__file__).resolve().parents[1]
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
class ScoreNet(nn.Module):
def __init__(self, dimension: int) -> None:
super().__init__()
self.net = nn.Sequential(
nn.Linear(dimension + 1, 48), nn.Tanh(),
nn.Linear(48, 48), nn.Tanh(),
nn.Linear(48, dimension),
)
def forward(self, x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
return self.net(torch.cat([x, torch.log(sigma)], dim=1))
def mixture_sample(rng: np.random.Generator, n: int, d: int) -> np.ndarray:
labels = rng.integers(0, 2, size=n)
means = np.where(labels[:, None] == 0, -1.0, 1.0)
return means + 0.35 * rng.normal(size=(n, d))
def exact_mixture_score(y: np.ndarray, sigma: float) -> np.ndarray:
"""Score of 0.5*N(-1, .35^2+sigma^2)+0.5*N(1, .35^2+sigma^2)."""
variance = 0.35**2 + sigma**2
left = np.exp(-np.sum((y + 1.0) ** 2, axis=1) / (2.0 * variance))
right = np.exp(-np.sum((y - 1.0) ** 2, axis=1) / (2.0 * variance))
weight_right = right / np.maximum(left + right, np.finfo(float).tiny)
posterior_mean = -1.0 + 2.0 * weight_right
return (posterior_mean[:, None] - y) / variance
def train_one(d: int, seed: int, steps: int = 700) -> dict[str, float | int]:
np_rng = np.random.default_rng(10000 + 97 * d + seed)
torch.manual_seed(20000 + 97 * d + seed)
clean = mixture_sample(np_rng, 1536, d).astype(np.float32)
clean_t = torch.from_numpy(clean)
sigmas = (0.08, 0.16, 0.32)
noisy_parts, target_parts, sigma_parts = [], [], []
for sigma in sigmas:
noise = np_rng.normal(size=clean.shape).astype(np.float32)
noisy = clean + sigma * noise
noisy_parts.append(torch.from_numpy(noisy))
target_parts.append(torch.from_numpy(-noise / sigma))
sigma_parts.append(torch.full((len(clean), 1), sigma, dtype=torch.float32))
x_train = torch.cat(noisy_parts)
target = torch.cat(target_parts)
sigma_train = torch.cat(sigma_parts)
model = ScoreNet(d)
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
for _ in range(steps):
opt.zero_grad(set_to_none=True)
prediction = model(x_train, sigma_train)
loss = torch.mean((prediction - target) ** 2)
loss.backward()
opt.step()
eval_clean = mixture_sample(np_rng, 2048, d).astype(np.float32)
energies = []
for sigma in sigmas:
noise = np_rng.normal(size=eval_clean.shape).astype(np.float32)
noisy = eval_clean + sigma * noise
x = torch.from_numpy(noisy)
s = torch.full((len(x), 1), sigma, dtype=torch.float32)
with torch.no_grad():
learned = model(x, s).numpy()
exact = exact_mixture_score(noisy.astype(np.float64), sigma)
energies.append(float(np.mean((learned - exact) ** 2)))
return {
"dimension": d,
"seed": seed,
"training_examples": int(len(x_train)),
"training_steps": steps,
"noise_levels": len(sigmas),
"score_error_energy_mean": float(np.mean(energies)),
"score_error_energy_max": float(np.max(energies)),
"score_error_energy_by_sigma": [float(x) for x in energies],
}
def recurrence(errors: np.ndarray, alpha: float, initial: float) -> tuple[np.ndarray, np.ndarray]:
beta = (1.0 - alpha) ** 2
values = np.empty(len(errors), dtype=float)
current = float(initial)
for i, error in enumerate(errors):
current = beta * current + float(error)
values[i] = current
direct = np.array([
sum(beta ** (n - j) * float(errors[j]) for j in range(n + 1)) + beta ** (n + 1) * initial
for n in range(len(errors))
])
return values, direct
def main() -> None:
rows = [train_one(d, seed) for d in (1, 2, 4) for seed in (0, 1, 2)]
recurrence_rows = []
for row in rows:
base = max(1e-8, min(0.5, row["score_error_energy_mean"]))
errors = base / np.square(np.arange(1, 65, dtype=float))
for alpha in (0.1, 0.5, 0.9):
actual, direct = recurrence(errors, alpha, initial=0.07)
recurrence_rows.append({
"dimension": row["dimension"], "seed": row["seed"], "alpha": alpha,
"error_energy_first": float(errors[0]), "generations": len(errors),
"beta": (1.0 - alpha) ** 2,
"max_identity_residual": float(np.max(np.abs(actual - direct))),
"final_divergence": float(actual[-1]),
})
result = {
"schema": "trained-score-network-scope-audit-v1",
"cpu_only": True,
"model": "48-48 tanh score network trained by fixed-noise denoising score matching",
"dimensions": [1, 2, 4], "seeds": [0, 1, 2],
"training_rows": rows, "training_row_count": len(rows),
"max_score_error_energy": max(row["score_error_energy_max"] for row in rows),
"recurrence_rows": recurrence_rows,
"max_recurrence_identity_residual": max(row["max_identity_residual"] for row in recurrence_rows),
"all_cpu": True,
}
output = ROOT / "outputs" / "learned_score_scope_audit.json"
output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
print(json.dumps(result, indent=2, sort_keys=True))
if __name__ == "__main__":
main()
|