1import json
  2import math
  3import random
  4from pathlib import Path
  5
  6import numpy as np
  7import torch
  8from torch import nn
  9
 10SEED = 587
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12
 13
 14def seed_all(seed=SEED):
 15    random.seed(seed)
 16    np.random.seed(seed)
 17    torch.manual_seed(seed)
 18    if torch.cuda.is_available():
 19        torch.cuda.manual_seed_all(seed)
 20
 21
 22class GRUSequence(nn.Module):
 23    def __init__(self, hidden=24):
 24        super().__init__()
 25        self.gru = nn.GRU(1, hidden, batch_first=True)
 26        self.readout = nn.Linear(hidden, 1)
 27
 28    def forward(self, u):
 29        # h_all includes h_1,...,h_T; h0 is explicitly zero.
 30        out, _ = self.gru(u)
 31        h0 = torch.zeros(u.shape[0], 1, out.shape[-1], device=u.device)
 32        h_all = torch.cat([h0, out], dim=1)
 33        y = self.readout(out).squeeze(-1)
 34        return y, h_all
 35
 36
 37def passivity_loss(u, y, h_all, gamma=0.10, incremental=None, eta=0.0):
 38    # u: B,T,1; y: B,T; h_all: B,T+1,H
 39    r = (h_all[:, 1:].pow(2).sum(-1) + y.pow(2)
 40         - h_all[:, :-1].pow(2).sum(-1) - u.squeeze(-1).pow(2))
 41    lp = torch.relu(r).mean() + gamma * torch.relu(r.sum(dim=1)).mean()
 42    inc_lp = torch.zeros((), device=u.device)
 43    if incremental is not None and eta:
 44        u2, y2, h2 = incremental
 45        du = u2 - u
 46        dy = y2 - y
 47        dh = h2 - h_all
 48        ri = (dh[:, 1:].pow(2).sum(-1) + dy.pow(2)
 49              - dh[:, :-1].pow(2).sum(-1) - du.squeeze(-1).pow(2))
 50        inc_lp = torch.relu(ri).mean()
 51        lp = lp + eta * inc_lp
 52    stats = {
 53        "max_positive_residual": float(torch.relu(r).max().detach().cpu()),
 54        "mean_cumulative_residual": float(r.sum(1).mean().detach().cpu()),
 55        "mean_hidden_norm": float(h_all[:, 1:].norm(dim=-1).mean().detach().cpu()),
 56        "penalty": float(lp.detach().cpu()),
 57        "incremental_penalty": float(inc_lp.detach().cpu()),
 58    }
 59    return lp, stats
 60
 61
 62def math_check():
 63    # Directly check sum_t r_t = ||h_T||^2 + sum ||y||^2 - ||h_0||^2 - sum ||u||^2.
 64    g = torch.Generator().manual_seed(SEED)
 65    B, T, H = 7, 13, 5
 66    u = torch.randn(B, T, 1, generator=g)
 67    h = torch.randn(B, T + 1, H, generator=g)
 68    y = torch.randn(B, T, generator=g)
 69    r = h[:, 1:].pow(2).sum(-1) + y.pow(2) - h[:, :-1].pow(2).sum(-1) - u.squeeze(-1).pow(2)
 70    lhs = r.sum(1)
 71    rhs = h[:, -1].pow(2).sum(-1) + y.pow(2).sum(1) - h[:, 0].pow(2).sum(-1) - u.squeeze(-1).pow(2).sum(1)
 72    err = float((lhs - rhs).abs().max())
 73    # A deliberately amplifying scalar trajectory has positive residuals.
 74    hs = [torch.tensor([[1.0]])]
 75    us, ys = [], []
 76    for _ in range(8):
 77        us.append(torch.zeros(1, 1))
 78        ys.append(torch.zeros(1))
 79        hs.append(1.15 * hs[-1])
 80    rr = []
 81    for t in range(8):
 82        rr.append(hs[t+1].pow(2).sum() - hs[t].pow(2).sum())
 83    positive_fraction = float((torch.stack(rr) > 0).float().mean())
 84    return {"telescoping_max_abs_error": err, "amplifying_positive_residual_fraction": positive_fraction}
 85
 86
 87def batch_data(n, T, seed):
 88    # Long-context signal: classify whether the sum of random +/-1 inputs is positive.
 89    gen = torch.Generator().manual_seed(seed)
 90    x = torch.randint(0, 2, (n, T, 1), generator=gen).float() * 2 - 1
 91    # Add a weak final tie breaker so labels are balanced and deterministic.
 92    sums = x.sum(dim=1).squeeze(-1)
 93    labels = (sums > 0).float()
 94    ties = sums == 0
 95    labels[ties] = (x[ties, -1, 0] > 0).float()
 96    return x, labels
 97
 98
 99def run_model(reg_lambda, paired, seed=SEED):
100    seed_all(seed)
101    model = GRUSequence().to(DEVICE)
102    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
103    train_u, train_y = batch_data(96, 32, 1001)
104    train_u, train_y = train_u.to(DEVICE), train_y.to(DEVICE)
105    last = {}
106    for step in range(260):
107        # Fresh small batches are unnecessary here; fixed data makes the comparison cheap/reproducible.
108        if paired:
109            noise = 0.05 * torch.randn_like(train_u)
110            u2 = train_u + noise
111            y2, h2 = model(u2)
112        y, h = model(train_u)
113        task = nn.functional.binary_cross_entropy_with_logits(y[:, -1], train_y)
114        inc = (u2, y2, h2) if paired else None
115        pl, st = passivity_loss(train_u, y, h, gamma=0.10, incremental=inc, eta=0.10 if paired else 0.0)
116        loss = task + reg_lambda * pl
117        opt.zero_grad()
118        loss.backward()
119        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
120        opt.step()
121        last = {"task": float(task.detach().cpu()), **st}
122
123    def evaluate(T, seed_offset):
124        u, target = batch_data(512, T, 4000 + seed_offset)
125        u, target = u.to(DEVICE), target.to(DEVICE)
126        with torch.no_grad():
127            y, h = model(u)
128            acc = ((torch.sigmoid(y[:, -1]) > .5) == target.bool()).float().mean()
129            _, st = passivity_loss(u, y, h, gamma=0.10)
130            # A finite input perturbation amplification at the final state.
131            u2 = u + 0.02 * torch.randn_like(u)
132            _, h2 = model(u2)
133            amp = ((h2[:, -1] - h[:, -1]).norm(dim=-1) /
134                   (u2 - u).norm(dim=(1, 2)).clamp_min(1e-8)).mean()
135        return {"accuracy": float(acc.cpu()), "amplification": float(amp.cpu()), **st}
136
137    return {"train": last, "test_T32": evaluate(32, 32), "test_T160": evaluate(160, 160)}
138
139
140def main():
141    seed_all()
142    math_result = math_check()
143    results = {
144        "device": DEVICE,
145        "math_check": math_result,
146        "baseline": run_model(0.0, False, SEED),
147        "passivity_lambda_0.01": run_model(0.01, False, SEED),
148        "passivity_lambda_0.01_incremental": run_model(0.01, True, SEED),
149    }
150    Path("results.json").write_text(json.dumps(results, indent=2))
151    print(json.dumps(results, indent=2))
152
153
154if __name__ == "__main__":
155    try:
156        main()
157    except Exception as e:
158        # CUDA can fail in a shared environment; retry entirely on CPU.
159        if DEVICE == "cuda":
160            print("CUDA failed, retrying on CPU:", repr(e))
161            torch.cuda.empty_cache()
162            DEVICE = "cpu"
163            main()
164        else:
165            raise