repro-bundle/v1/scripts/ste_sim_batch.py

"""Batched multi-setting STE simulation.

Run multiple (quantizer, hyperparameter) configurations simultaneously as a
single (B, d) tensor where B = n_seeds * n_settings. Each setting's quantizer
is applied to its slice of rows. This amortises the GPU per-step kernel-launch
overhead (~350 us regardless of B) across all settings, so adding more settings
is essentially free in wall-clock.

Used for Claim 2 (Figures 2, 3) and Claim 3 (Figure 6) sweeps.
"""
from __future__ import annotations

import math
import time
from dataclasses import dataclass, field
from typing import List

import numpy as np
import torch


@dataclass
class Setting:
    """One configuration (quantizers + hyperparams)."""
    name: str
    qw_levels: list
    qw_theta: list
    qw_omega: float
    qw_Delta: float
    qx_levels: list | None = None     # None = identity input quantizer
    qx_theta: list | None = None
    qx_omega: float = 0.0
    qx_Delta: float = 0.0
    eta: float = 0.04
    lam: float = 1.0
    sigma2_x: float = 1.0
    kappa_x: float = 1.0
    # the following are usually shared across settings:
    d: int = 900
    rho: float = 1.0
    sigma2: float = 0.0
    w_init_std: float = 1.0
    w_star_value: float = 1.0
    n_seeds: int = 5


def run_ste_batch(settings: List[Setting], n_steps: int, log_period: int,
                  device: str = "cuda", dtype: torch.dtype = torch.float32,
                  seed_offset: int = 0):
    """Run a batched STE simulation across all settings.

    Returns:
      results: list (one per setting) of dicts with arrays
               (taus, m, q, m_psi, q_psi, eps_g)   shape (T, n_seeds)
      wall_seconds: float
    """
    use_cuda = device == "cuda" and torch.cuda.is_available()
    dev = torch.device("cuda" if use_cuda else "cpu")
    assert len(settings) > 0
    d = settings[0].d
    n_settings = len(settings)
    n_seeds = settings[0].n_seeds
    B = n_settings * n_seeds

    # -- pre-build per-setting tensors --
    qw_levels_t = []
    qw_theta_t = []
    qx_levels_t = []
    qx_theta_t = []
    use_id_x_list = []
    eta_list = []
    lam_list = []
    sigma2_x_list = []
    kappa_x_list = []
    for s in settings:
        qw_levels_t.append(torch.tensor(s.qw_levels, dtype=dtype, device=dev))
        qw_theta_t.append(torch.tensor(s.qw_theta, dtype=dtype, device=dev))
        if s.qx_levels is None:
            qx_levels_t.append(None)
            qx_theta_t.append(None)
            use_id_x_list.append(True)
        else:
            qx_levels_t.append(torch.tensor(s.qx_levels, dtype=dtype, device=dev))
            qx_theta_t.append(torch.tensor(s.qx_theta, dtype=dtype, device=dev))
            use_id_x_list.append(False)
        eta_list.append(s.eta)
        lam_list.append(s.lam)
        sigma2_x_list.append(s.sigma2_x)
        kappa_x_list.append(s.kappa_x)

    eta_arr = torch.tensor(eta_list, dtype=dtype, device=dev)         # (n_settings,)
    lam_arr = torch.tensor(lam_list, dtype=dtype, device=dev)
    sigma2_x_arr = torch.tensor(sigma2_x_list, dtype=dtype, device=dev)
    kappa_x_arr = torch.tensor(kappa_x_list, dtype=dtype, device=dev)
    # broadcast to (B,) by repeating each setting's value n_seeds times
    eta_b = eta_arr.repeat_interleave(n_seeds).view(B, 1)              # (B, 1)
    lam_b = lam_arr.repeat_interleave(n_seeds).view(B, 1)
    sigma2_x_b = sigma2_x_arr.repeat_interleave(n_seeds)              # (B,)
    kappa_x_b = kappa_x_arr.repeat_interleave(n_seeds)                # (B,)

    # initial state
    gen = torch.Generator(device=dev).manual_seed(1234 + seed_offset)
    w = torch.randn(B, d, dtype=dtype, device=dev, generator=gen) * settings[0].w_init_std
    # (assume all settings share w_init_std and w_star_value)
    w_star = torch.full((d,), settings[0].w_star_value, dtype=dtype, device=dev)
    rho = settings[0].rho
    sigma2 = settings[0].sigma2
    noise_std = math.sqrt(sigma2) if sigma2 > 0 else 0.0

    inv_sqrt_d = 1.0 / math.sqrt(d)
    inv_d = 1.0 / d

    # log buffers
    n_logs = n_steps // log_period + 1
    m_log = np.empty((n_logs, B), dtype=np.float64)
    q_log = np.empty((n_logs, B), dtype=np.float64)
    m_psi_log = np.empty((n_logs, B), dtype=np.float64)
    q_psi_log = np.empty((n_logs, B), dtype=np.float64)
    eps_log = np.empty((n_logs, B), dtype=np.float64)

    def quantize_w_batch(w):
        """Apply per-setting weight quantizer to its slice of rows."""
        out = torch.empty_like(w)
        for si in range(n_settings):
            sl = slice(si * n_seeds, (si + 1) * n_seeds)
            out[sl] = qw_levels_t[si][torch.bucketize(w[sl], qw_theta_t[si])]
        return out

    def quantize_x_batch(x):
        out = torch.empty_like(x)
        for si in range(n_settings):
            sl = slice(si * n_seeds, (si + 1) * n_seeds)
            if use_id_x_list[si]:
                out[sl] = x[sl]
            else:
                out[sl] = qx_levels_t[si][torch.bucketize(x[sl], qx_theta_t[si])]
        return out

    def log_state(idx):
        with torch.no_grad():
            psi_w = quantize_w_batch(w)
            m = (w * w_star).sum(dim=1) * inv_d
            q = (w * w).sum(dim=1) * inv_d
            m_psi = (psi_w * w_star).sum(dim=1) * inv_d
            q_psi = (psi_w * psi_w).sum(dim=1) * inv_d
            eps = sigma2 + rho + sigma2_x_b * q_psi - 2.0 * kappa_x_b * m_psi
        m_log[idx] = m.cpu().numpy()
        q_log[idx] = q.cpu().numpy()
        m_psi_log[idx] = m_psi.cpu().numpy()
        q_psi_log[idx] = q_psi.cpu().numpy()
        eps_log[idx] = eps.cpu().numpy()

    idx = 0
    log_state(idx); idx += 1
    if use_cuda:
        torch.cuda.synchronize()
    t0 = time.time()

    for outer in range(n_steps // log_period):
        for _ in range(log_period):
            x = torch.randn(B, d, dtype=dtype, device=dev, generator=gen)
            psi_x = quantize_x_batch(x)
            psi_w = quantize_w_batch(w)
            y = (x * w_star).sum(dim=1) * inv_sqrt_d
            if noise_std > 0:
                y = y + torch.randn(B, dtype=dtype, device=dev, generator=gen) * noise_std
            yhat = (psi_w * psi_x).sum(dim=1) * inv_sqrt_d
            r = yhat - y                                              # (B,)
            # per-row eta, lam: shape (B, 1)
            w = w - eta_b * (r.unsqueeze(1) * psi_x * inv_sqrt_d + lam_b * inv_d * psi_w)
        log_state(idx); idx += 1

    if use_cuda:
        torch.cuda.synchronize()
    wall = time.time() - t0

    taus = np.arange(n_logs) * (log_period / d)

    results = []
    for si, s in enumerate(settings):
        sl = slice(si * n_seeds, (si + 1) * n_seeds)
        results.append({
            "name": s.name,
            "taus": taus,
            "m": m_log[:, sl],
            "q": q_log[:, sl],
            "m_psi": m_psi_log[:, sl],
            "q_psi": q_psi_log[:, sl],
            "eps_g": eps_log[:, sl],
        })
    return results, wall