Skip to content

perf: Rung5(n_amp_qubits=2) is ~60× faster than mathematically-equivalent Rung4 cancellation path #86

Description

@jascal

TL;DR

Rung5(n_amp_qubits=2) and Rung4 are mathematically identical encodings —
per Rung5's own docstring: "Generalises Rung4 (the fixed n_amp_qubits=2
case) to an arbitrary amp-register width fixed at construction time."
But
the Rung5 cancellation code path runs roughly 60× faster than Rung4
on the same input at the same max_steps=10 scipy budget.

Observed on v0.8.1 (commit 7ad0ea6) against the
sm-sae benchmark fixture.

Evidence

Same 4 cancellation pairs (e⁻/e⁺, u_r/ū_r, e⁻/u_r, photon/W⁺), same
6-knob amplitude search, same optimize={"method": "scipy", "max_steps": 10, "seed": 0}, same preserve_tiers=True. Every pair drives Δ to
exactly −|before| in both configs — the optima are identical. The only
thing that differs is wall time:

pair Rung4 wall time Rung5(n_amp_qubits=2) wall time
e_e 273.1 s 5.0 s
u_pair 229.5 s 3.8 s
lepton_quark 223.9 s 3.9 s
boson_pair 241.3 s 4.1 s
total ~16 min ~17 s

(For reference, default-budget Rung4 — max_steps=50 — takes ~80 minutes
per pair. The benchmark stopped that variant after one pair.)

Where the gap likely lives

Without profiling I can't pinpoint, but the dispatch surface suggests the
Rung4 cancellation path (_run_rung4, around cancellation.py:901)
re-builds the dictionary / gram via a more expensive route than the Rung5
path (_run_rung5, around cancellation.py:1086+). Both call
scipy.optimize.differential_evolution with the same call shape, so the
cost is per-objective-evaluation — probably in gram construction.

Recommendation

Three options in roughly increasing intrusiveness:

  1. Re-implement _run_rung4 as a thin wrapper that delegates to the
    Rung5 path with n_amp_qubits=2 hardcoded.
    Guarantees mathematical
    equivalence (it's already true by the docstring) and inherits the
    speed. Probably the cheapest path to the win.
  2. Profile and bring the Rung4 gram-construction path up to parity.
    More work; preserves a separate Rung4 implementation.
  3. Deprecate Rung4 in favour of Rung5(n_amp_qubits=2) with a
    migration note. Rung5 generalises it cleanly; the API surface is
    strictly larger.

Reproduction

The sm-sae sweep harness that surfaced this lives at
scripts/polygram_sweep.py
specifically the Rung4_amp_budget and Rung5_amp_budget config
entries. Both run against polygram.from_sae_lens Features built from
the 8-particle SM slice with assign_amp_knobs=True,
assign_phase_knobs=True. Wall times are captured per-pair in the
emitted runs/polygram/sweep/sweep_results.json.

Why this matters

For benchmarks and downstream consumers (e.g.
sae-forge), choosing between
Rung4 and Rung5(n_amp_qubits=2) is currently a ~60× perf decision
without any semantic difference. That ambiguity belongs in polygram, not
in callers' configs.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions