gpu/gpu_monte_carlo.py

#!/usr/bin/env python3
"""GPU-scale independent audit of finite-temperature attention metastability.

This uses the stochastic self-attention process SA_P from Eq. (5.3) of
arXiv:2508.09628.  It checks the constant-step first phase on an equilateral
configuration satisfying the paper's geometric hypotheses, then measures
first-exit times from the singleton-cluster epsilon neighborhoods used by
Theorem 5.4.  All particles update simultaneously and each query samples an
independent categorical key from its finite-beta attention weights.
"""

from __future__ import annotations

import json
import math
import platform
import time

import torch


SEED = 809704


def regular_polygon(count: int, radius: float, device: torch.device) -> torch.Tensor:
    angles = torch.arange(count, device=device, dtype=torch.float32) * (2.0 * math.pi / count)
    return radius * torch.stack((torch.cos(angles), torch.sin(angles)), dim=1)


def stochastic_step(
    state: torch.Tensor,
    beta: float,
    gamma: float,
    generator: torch.Generator,
) -> torch.Tensor:
    """One vectorized SA_P step for state shaped (replicas, particles, 2)."""

    scores = torch.einsum("rid,rjd->rij", state, state)
    probabilities = torch.softmax(float(beta) * scores, dim=2)
    uniforms = torch.rand(
        (*state.shape[:2], 1),
        dtype=state.dtype,
        device=state.device,
        generator=generator,
    )
    indices = (uniforms > torch.cumsum(probabilities, dim=2)).sum(dim=2)
    indices.clamp_(max=state.shape[1] - 1)
    selected = torch.gather(state, 1, indices[..., None].expand(-1, -1, state.shape[2]))
    return (1.0 - gamma) * state + gamma * selected


def first_phase(device: torch.device, generator: torch.Generator) -> dict[str, object]:
    vertices = regular_polygon(3, 1.0, device)
    interiors = 0.25 * vertices
    initial = torch.cat((vertices, interiors), dim=0)
    beta = 1_000_000.0
    gamma = 0.1
    c0 = 1.5
    diameter = math.sqrt(3.0)
    initial_distance = 0.75
    tau = min(c0 / (2.0 * diameter), math.sqrt(2.0 * c0) / 2.0, (1.0 - gamma) * initial_distance)
    t1 = math.floor(math.log(tau / initial_distance) / math.log(1.0 - gamma))
    replicas = 262_144
    state = initial.expand(replicas, -1, -1).clone()
    for _ in range(t1):
        state = stochastic_step(state, beta, gamma, generator)
    vertex_ok = torch.all(
        torch.linalg.vector_norm(state[:, :3] - vertices[None, :, :], dim=2) <= beta ** (-0.25) + 1e-6,
        dim=1,
    )
    interior_ok = torch.all(
        torch.linalg.vector_norm(state[:, 3:] - vertices[None, :, :], dim=2) <= 1.1 * tau + 1e-6,
        dim=1,
    )
    success = vertex_ok & interior_ok

    sweep: list[dict[str, float]] = []
    sweep_replicas = 65_536
    for local_beta in (2.0, 4.0, 8.0, 16.0, 32.0, 64.0):
        local = initial.expand(sweep_replicas, -1, -1).clone()
        for _ in range(t1):
            local = stochastic_step(local, local_beta, gamma, generator)
        labels = torch.argmax(torch.einsum("rid,kd->rik", local[:, 3:], vertices), dim=2)
        target = torch.arange(3, device=device)[None, :]
        correct = torch.mean((labels == target).to(torch.float64)).item()
        distance = torch.mean(torch.linalg.vector_norm(local[:, 3:] - vertices[None, :, :], dim=2)).item()
        sweep.append(
            {
                "beta": local_beta,
                "correct_cell_fraction": correct,
                "mean_distance_to_assigned_vertex": distance,
            }
        )

    return {
        "beta": beta,
        "gamma": gamma,
        "T1": t1,
        "tau": tau,
        "tested_C": 1.1,
        "replicas": replicas,
        "successes": int(success.sum().item()),
        "empirical_success_probability": float(success.to(torch.float64).mean().item()),
        "paper_lower_bound": 1.0 - beta ** (-0.125),
        "moderate_beta_sweep_replicas": sweep_replicas,
        "moderate_beta_sweep": sweep,
    }


def first_exit_sweep(device: torch.device, generator: torch.Generator) -> dict[str, object]:
    vertices = regular_polygon(3, 0.5, device)
    gamma = 0.02
    diameter = math.sqrt(3.0) * 0.5
    epsilon = 2.0 * diameter * gamma
    c0 = 0.375
    replicas = 16_384
    max_steps = 25_000
    beta_values = (4.0, 8.0, 12.0, 16.0, 20.0, 24.0)
    rows: list[dict[str, float | int]] = []
    total_particle_updates = 0

    for beta in beta_values:
        state = vertices.expand(replicas, -1, -1).clone()
        alive = torch.ones(replicas, dtype=torch.bool, device=device)
        exits = torch.full((replicas,), max_steps + 1, dtype=torch.int64, device=device)
        steps_executed = 0
        for step in range(1, max_steps + 1):
            state = stochastic_step(state, beta, gamma, generator)
            distance = torch.linalg.vector_norm(state - vertices[None, :, :], dim=2)
            newly_exited = alive & torch.any(distance > epsilon + 1e-6, dim=1)
            exits[newly_exited] = step
            alive &= ~newly_exited
            steps_executed = step
            total_particle_updates += replicas * 3
            if not torch.any(alive):
                break

        uncensored = exits[exits <= max_steps].to(torch.float64)
        if uncensored.numel() == 0:
            median = float(max_steps + 1)
            mean = float(max_steps + 1)
            q10 = float(max_steps + 1)
            q90 = float(max_steps + 1)
        else:
            median = float(torch.quantile(uncensored, 0.5).item())
            mean = float(uncensored.mean().item())
            q10 = float(torch.quantile(uncensored, 0.1).item())
            q90 = float(torch.quantile(uncensored, 0.9).item())
        row = {
            "beta": beta,
            "replicas": replicas,
            "steps_executed": steps_executed,
            "median_exit_time_uncensored": median,
            "mean_exit_time_uncensored": mean,
            "q10_exit_time_uncensored": q10,
            "q90_exit_time_uncensored": q90,
            "censored_fraction": float((exits > max_steps).to(torch.float64).mean().item()),
            "empirical_survival_at_t5": float((exits >= 5).to(torch.float64).mean().item()),
        }
        rows.append(row)
        print("METASTABILITY_ROW=" + json.dumps(row, sort_keys=True), flush=True)

    betas = torch.tensor([float(row["beta"]) for row in rows], dtype=torch.float64)
    log_medians = torch.log(
        torch.tensor([float(row["median_exit_time_uncensored"]) for row in rows], dtype=torch.float64)
    )
    design = torch.stack((betas, torch.ones_like(betas)), dim=1)
    coeff = torch.linalg.lstsq(design, log_medians).solution
    predicted = design @ coeff
    residual = torch.sum((log_medians - predicted) ** 2)
    total = torch.sum((log_medians - torch.mean(log_medians)) ** 2)
    r_squared = float((1.0 - residual / total).item())

    return {
        "geometry": "three singleton clusters at an equilateral triangle of radius 0.5",
        "gamma": gamma,
        "epsilon": epsilon,
        "epsilon_over_gamma": epsilon / gamma,
        "required_minimum_2diameter": 2.0 * diameter,
        "c0": c0,
        "replicas_per_beta": replicas,
        "max_steps": max_steps,
        "total_particle_updates": total_particle_updates,
        "log_median_exit_vs_beta_slope": float(coeff[0].item()),
        "log_linear_fit_r_squared": r_squared,
        "rows": rows,
    }


def main() -> None:
    if not torch.cuda.is_available():
        raise RuntimeError("This audit requires a CUDA GPU; refusing a silent CPU fallback.")
    device = torch.device("cuda")
    # Monte Carlo category sampling is insensitive to fp64 roundoff, while
    # fp32 uses the accelerator's native high-throughput path.
    torch.set_default_dtype(torch.float32)
    generator = torch.Generator(device=device)
    generator.manual_seed(SEED)
    start = time.perf_counter()
    gpu = torch.cuda.get_device_properties(device)
    result = {
        "schema_version": 1,
        "paper": "Attention's forward pass and Frank-Wolfe",
        "openreview_id": "zrn7rRuvhW",
        "arxiv_id": "2508.09628",
        "seed": SEED,
        "implementation": "independent PyTorch vectorization of stochastic SA_P Eq. (5.3)",
        "simulation_dtype": "float32",
        "software": {
            "python": platform.python_version(),
            "torch": torch.__version__,
            "cuda_runtime": torch.version.cuda,
        },
        "hardware": {
            "gpu_name": gpu.name,
            "gpu_memory_bytes": gpu.total_memory,
            "device_count": torch.cuda.device_count(),
        },
    }
    result["theorem_5_2_first_phase"] = first_phase(device, generator)
    result["theorem_5_4_metastability"] = first_exit_sweep(device, generator)
    torch.cuda.synchronize()
    result["wall_time_seconds"] = time.perf_counter() - start
    first = result["theorem_5_2_first_phase"]
    meta = result["theorem_5_4_metastability"]
    result["checks"] = {
        "first_phase_exceeds_paper_lower_bound": (
            first["empirical_success_probability"] >= first["paper_lower_bound"]
        ),
        "exit_time_log_linear_r_squared_above_0_99": meta["log_linear_fit_r_squared"] > 0.99,
        "positive_exponential_slope": meta["log_median_exit_vs_beta_slope"] > 0.0,
    }
    print("RESULT_JSON=" + json.dumps(result, sort_keys=True), flush=True)


if __name__ == "__main__":
    main()