"""Shared networks, data, and optimization for CIS 6270 Lecture 6. The small MLPs keep time derivatives and all training steps visible. Sequence networks receive the whole padded sequence, so predictions can depend on other positions. No pretrained checkpoint or dataset download is required. """ import copy import json import math import random from pathlib import Path import numpy as np import torch from torch import nn from torch.nn import functional as F def seed_all(seed, threads=1): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.set_num_threads(threads) def mlp(inputs, outputs, width): return nn.Sequential(nn.Linear(inputs, width), nn.SiLU(), nn.Linear(width, width), nn.SiLU(), nn.Linear(width, outputs)) def time_like(t, x): """One scalar time per batch member, shape [B,1].""" t = torch.as_tensor(t, dtype=x.dtype, device=x.device) if t.numel() == 1: return t.expand(len(x), 1) return t.reshape(len(x), 1) class MapNet(nn.Module): def __init__(self, dim=2, width=64, context_dim=0): super().__init__() self.net = mlp(dim + 2 + context_dim, dim, width) def forward(self, x, s, t, context=None): inputs = [x, time_like(s, x), time_like(t, x)] if context is not None: inputs.append(context) return self.net(torch.cat(inputs, -1)) class SequenceNet(nn.Module): def __init__(self, length, vocab, width=64): super().__init__() self.length, self.vocab = length, vocab self.net = mlp(length * vocab + 2, length * vocab, width) def forward(self, x, s, t): inputs = torch.cat([x.flatten(1), time_like(s, x), time_like(t, x)], -1) return self.net(inputs).reshape(-1, self.length, self.vocab) def finite_map(model, x, s, t, context=None): """F(s,t,x) = x + (t-s) v(s,t,x), including exact F(s,s,x)=x.""" s, t = time_like(s, x), time_like(t, x) return x + (t - s) * model(x, s, t, context) def ordered_times(x, ceiling=.98): times = ceiling * torch.rand(len(x), 2, device=x.device, dtype=x.dtype) s, t = times.sort(-1).values.split(1, -1) return s, t, (s + t) / 2 def draw_batch(data, count): return data[torch.randint(len(data), (count,), device=data.device)] def interpolate(data, time, noise=None): noise = torch.randn_like(data) if noise is None else noise t = time.reshape(len(data), *([1] * (data.ndim - 1))) return (1 - t) * noise + t * data, data - noise def ema_copy(model): result = copy.deepcopy(model).eval() result.requires_grad_(False) return result @torch.no_grad() def update_ema(ema, model, decay=.99): for p, q in zip(ema.parameters(), model.parameters()): p.lerp_(q, 1 - decay) def optimize(model, objective, steps, lr=1e-3, ema=None): optimizer = torch.optim.Adam(model.parameters(), lr=lr) logs = [] for step in range(steps): loss, details = objective(step) if not torch.isfinite(loss): raise FloatingPointError(f'Nonfinite loss at step {step}: {details}') optimizer.zero_grad(set_to_none=True) loss.backward() norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 10.) if not torch.isfinite(norm): raise FloatingPointError(f'Nonfinite gradient at step {step}') optimizer.step() if ema is not None: update_ema(ema, model) logs.append({'step': step, 'loss': float(loss.detach()), **{k: float(torch.as_tensor(v).detach()) for k, v in details.items()}}) return logs CENTERS = torch.tensor([[-1.5, -1.5], [-1.5, 1.5], [1.5, -1.5], [1.5, 1.5]]) DATA_STD = .22 def mixture_data(n, generator=None): ids = torch.randint(4, (n,), generator=generator) return CENTERS[ids] + DATA_STD * torch.randn(n, 2, generator=generator) def mixture_posterior(x, t): """Exact p(X1 | (1-t)X0+tX1=x) for the four Gaussian mixture. Return component probabilities, conditional means, scalar variances. This analytic oracle is used for diagnostics and the GLASS example. """ t = time_like(t, x) a = 1 - t variance = a.square() + (t * DATA_STD).square() centers = CENTERS.to(x) delta = x[:, None, :] - t[:, None, :] * centers logits = -delta.square().sum(-1) / (2 * variance) prob = logits.softmax(-1) gain = t * DATA_STD**2 / variance means = centers + gain[:, None, :] * delta posterior_variance = DATA_STD**2 * a.square() / variance return prob, means, posterior_variance def exact_denoiser(x, t): p, means, _ = mixture_posterior(x, t) return (p[..., None] * means).sum(1) def exact_velocity(x, t): t = time_like(t, x) # Direct conditional velocity avoids cancellation at t=1. a = 1 - t variance = a.square() + (t * DATA_STD).square() centers = CENTERS.to(x) delta = x[:, None] - t[:, None] * centers weights = (-delta.square().sum(-1) / (2 * variance)).softmax(-1) component_velocity = centers + ((t * DATA_STD**2 - a) / variance)[:, None] * delta return (weights[..., None] * component_velocity).sum(1) def mixture_metrics(samples): x = samples.detach().cpu() distances = (x[:, None] - CENTERS).square().sum(-1) counts = torch.bincount(distances.argmin(-1), minlength=4).float() p = counts / len(x) logp = torch.logsumexp(-distances / (2 * DATA_STD**2), -1) logp -= math.log(4 * 2 * math.pi * DATA_STD**2) return {'finite_samples': bool(torch.isfinite(x).all()), 'mean_distance_to_center': float(distances.min(-1).values.sqrt().mean()), 'mode_fractions': p.tolist(), 'mode_entropy': float(-(p * p.clamp_min(1e-12).log()).sum()), 'mean_target_log_density': float(logp.mean())} def load_text(path, variable=False): lines = [x.strip().split() for x in Path(path).read_text().splitlines() if x.strip()] if len(lines) < 2: raise ValueError('Text data need at least two nonempty lines.') vocab = sorted(set(word for line in lines for word in line)) lengths = torch.tensor([len(x) for x in lines]) if not variable and len(set(lengths.tolist())) != 1: raise ValueError('Fixed-length methods require equal tokens per line; use expanding otherwise.') if lengths.max() > 32: raise ValueError('Teaching MLP supports up to 32 positions.') ids = {word: i for i, word in enumerate(vocab)} encoded = torch.zeros(len(lines), int(lengths.max()), dtype=torch.long) for i, line in enumerate(lines): encoded[i, :len(line)] = torch.tensor([ids[word] for word in line]) return encoded, lengths, vocab def text_metrics(ids, vocab, lengths=None, reference=None): lengths = [ids.shape[1]] * len(ids) if lengths is None else lengths.tolist() texts = [' '.join(vocab[j] for j in row[:length]) for row, length in zip(ids.tolist(), lengths)] flat = [j for row, length in zip(ids.tolist(), lengths) for j in row[:length]] counts = torch.bincount(torch.tensor(flat, dtype=torch.long), minlength=len(vocab)).float() p = counts / counts.sum().clamp_min(1) metrics = {'unique_fraction': len(set(texts)) / len(texts), 'token_entropy': float(-(p * p.clamp_min(1e-12).log()).sum()), 'mean_length': sum(lengths) / len(lengths), 'empty_fraction': lengths.count(0) / len(lengths)} if reference is not None: metrics['training_support_fraction'] = sum(x in reference for x in texts) / len(texts) return texts, metrics def write_json(path, value): Path(path).write_text(json.dumps(value, indent=2, allow_nan=False) + '\n')