FFT-parameterized circulant layers, explicit Langevin dynamics, reference-matrix checks, and reproducible benchmark artifacts.
A dense hidden layer stores a full D D matrix. A circulant layer stores one generator vector c R^D, applies the corresponding circular convolution with FFTs, and exposes an exact matrix convention that can be checked against a materialized reference.
This repository packages that structure into scalar energy models and tests the entire path: layer algebra, parameter counts, input gradients, Langevin updates, persistent chains, serialization, and small distribution-learning experiments.
Scope: this is a reproducible proof of concept and engineering baseline. It does not claim universal runtime superiority, equal expressivity to dense networks, or a new invention in circulant matrices or EBMs.
| Dimension | Dense parameters | Spectral parameters | Reduction |
|---|---|---|---|
| 128 | 49,664 | 896 | 55.4 |
| 512 | 788,480 | 3,584 | 219.8 |
| 1,024 | 3,149,824 | 7,168 | 439.6 |
| 2,048 | 12,591,104 | 14,336 | 878.8 |
At D = 2,048, the spectral model uses nearly three orders of magnitude fewer trainable parameters. The trade-off is structural restriction and, in the current CUDA snapshot, a slower end-to-end ULA step: 1.363 ms spectral versus 0.729 ms dense on an RTX 4060 Laptop GPU. The benchmark makes that trade-off visible instead of hiding it behind asymptotic notation.
The included score-matching smoke tests are intentionally small and diagnostic, not claims of state-of-the-art modeling quality.
| Experiment | Dense score MSE | Spectral score MSE | Better snapshot |
|---|---|---|---|
| Standard Gaussian DSM | 0.603 | 0.553 | Spectral |
| Four-mode Gaussian mixture | 13.777 | 10.519 | Spectral |
| Noi 8000 sy ring score matching | 19.239 | 19.785 | Dense |
All values above are read directly from the committed JSON artifacts in benchmark_results/. See the benchmark protocol for device, batch size, repetitions, and timing conventions.
The release now includes three structured extensions derived from the architectural audit:
BlockCirculantLinear/BlockSpectralEBM: dense channel mixing with a circulant FFT block for every input/output channel pair. The implementation exposes a slow dense reference matrix and an exact Fourier-symbol operator norm.PermutedSpectralEBM: deterministic, serialized coordinate permutations between spectral layers. The fixed shuffles add no trainable parameters and intentionally break the shared cyclic coordinate symmetry; they are an expressivity heuristic, not a universal theorem.vectorized_langevin_chain: persistent-state ULA execution that reuses a detached state buffer. Fixed-noise tests prove step-for-step agreement with repeatedula_stepcalls; it is an allocation optimization, not graph caching or a new sampler.
The extension benchmark is intentionally modest and descriptive:
| D | Dense channel-map parameters | Block-circulant parameters | Reduction |
|---|---|---|---|
| 32 | 16,512 | 640 | 25.8 |
| 64 | 65,792 | 1,280 | 51.4 |
Raw measurements and the exact command are in benchmark_results/2026-07-14-extensions-cpu.json and the benchmark protocol.
| Goal | Path |
|---|---|
| Understand the math | Proof and conventions |
| Run the tests | python -m pytest |
| Try a minimal model | Install and example |
| Reproduce timings | Benchmark protocol |
| Reproduce toy experiments | scripts/ and benchmark_results/ |
| Evaluate novelty claims | Prior art and novelty boundary |
| Formal-state integration | HRR proof-search adapter |
| Contribute | CONTRIBUTING.md |
Python 3.10+ and PyTorch 2.1+ are supported.
python -m pip install -e .
python -m pip install -r requirements-dev.txt
$env:PYTEST_DISABLE_PLUGIN_AUTOLOAD = "1"
python -m pytest -q
ruff check .import torch
from spectral_ebm import SpectralEBM, langevin_sample
model = SpectralEBM(dim=32, hidden_layers=3)
initial = torch.randn(16, 32)
samples = langevin_sample(model, initial, steps=100, step_size=0.01)
print(samples.shape) # torch.Size([16, 32])For persistent chains and training losses, see spectral_ebm/chains.py and spectral_ebm/training.py.
The implementation uses one explicit convention throughout:
W[i, j] = c[(i - j) mod D]
W x = irfft(rfft(x) * rfft(c), n=D)
The first column of W is c. The spectral norm is computed exactly from the maximum magnitude of the discrete Fourier spectrum of c. For energy E(x) at temperature T, the implemented ULA update is:
x_next = x - h/(2T) * grad E(x) + sqrt(h) * epsilon
~ Normal(0, I)
The reference construction, normalization details, invariants, and limitations are written out in docs/proof.md. Projected bounds are exposed as an explicit approximation; they are not presented as exact unconstrained sampling.
# Layer and ULA benchmark
python -m benchmarks.benchmark_layers --device cuda --dims 128 256 512 --batch-size 64 --repeats 10 --warmup 5 --output benchmark_results/local-cuda.json
# Larger end-to-end measurements
python -m benchmarks.full_benchmark --device cuda --dims 1024 2048 --batch-size 64 --repeats 10 --output benchmark_results/local-cuda-large.json
# Toy distribution checks
python scripts/train_gaussian_dsm.py --output benchmark_results/local-gaussian.json
python scripts/train_mixture_dsm.py --output benchmark_results/local-mixture.json
python scripts/train_ring_score.py --output benchmark_results/local-ring.json
# Rebuild the committed README figures
python scripts/make_plots.pyThe committed artifacts include CPU and CUDA timings, large-dimension measurements, parameter counts, score-matching results, and every timing repetition. Results are hardware-specific. The current data supports a strong parameter-efficiency claim, not a universal speed claim.
BlockCirculantLinear(..., backend="triton") fuses the complex frequency-bin channel contraction in Triton while leaving the explicit torch.fft.rfft/irfft transforms on cuFFT. This removes the intermediate broadcast/einsum tensor; it does not claim that cuFFT's internal FFT stages are fused or that every GPU receives a 2x speedup.
python -m pip install -e .[triton]Use backend="torch" for the portable default. Triton requires CUDA, a compatible Triton build, and a working C/CUDA compiler toolchain; the test suite skips the accelerator test when those prerequisites are absent.
The capability probe records unavailable environments, while the measured RTX 4060 Laptop GPU artifact records an actual comparison. These artifacts do not claim a universal speedup.
DifferentiablePermutation learns a doubly-stochastic Sinkhorn matrix and supports straight-through hard assignments. PermutedSpectralEBM(permutation_mode="differentiable") inserts these layers between spectral blocks, so the coordinate shuffle receives gradients and is serialized with the model.
from spectral_ebm import PermutedSpectralEBM
model = PermutedSpectralEBM(
dim=128,
hidden_layers=3,
permutation_mode="differentiable",
permutation_hard=False,
)HRREncoder binds token vectors to positional role vectors with circular convolution, bundles them into continuous state vectors, and FormalProofSearchAdapter refines those vectors with the persistent Langevin chain.
from spectral_ebm import FormalProofSearchAdapter, HRREncoder, SpectralEBM
encoder = HRREncoder(["by", "intro", "exact", "h"], dim=128)
adapter = FormalProofSearchAdapter(encoder, SpectralEBM(128))
result = adapter.refine(
[["by", "intro", "h"], ["exact", "h"]],
steps=8,
step_size=0.01,
noise_scale=0.0,
)This is an integration interface for Lean 4 tactic tokens or parser-produced AST nodes. It does not include a Lean parser, tactic decoder, proof checker, or verified theorem result; a production deployment must connect those trusted components around the continuous candidate loop.
Run the end-to-end smoke adapter with:
python scripts/formal_search_demo.py --dim 64 --steps 4The v0.4.0 hardening surface adds three opt-in production controls:
- AmortizedHouseholderPermutation composes K trainable Householder reflections with K*D parameters and preserves the Euclidean norm up to floating-point error. It is an orthogonal mixer, not a strict one-hot permutation matrix.
- FormalProofSearchAdapter now defaults to tangent-projected spherical Langevin refinement, keeping HRR states on a chosen radius. Set spherical=False to use the existing Euclidean persistent ULA path.
- The tiled Triton frequency mixer aggregates row, output-channel, and frequency tiles and processes input channels in bounded chunks. Explicit FFT transforms remain cuFFT operations.
Run the scale and stability audit with:
python -m benchmarks.hardening_audit --device cuda --dim 4096 --batch-size 64 --channels 8 --output benchmark_results/hardening-cuda.json
The audit reports Sinkhorn versus Householder parameter memory, CUDA peak allocation when available, and per-step sphere-norm envelopes. The committed CPU smoke artifact uses a smaller dimension so it remains reproducible on ordinary development machines.
spectral_ebm/ Core layers, models, chains, samplers, and training losses
benchmarks/ Timing harnesses with synchronized CUDA measurements
scripts/ Toy experiments and reproducible figure generation
tests/ Algebra, gradients, sampling, serialization, and training tests
docs/ Proof, training, benchmark, formal-search, and novelty documentation
benchmark_results/ Raw JSON artifacts used for the published result plots
Circulant and FFT-structured projections are established prior art, as are EBMs trained with Langevin-based methods. The project documents that boundary explicitly and avoids describing the basic combination as a new invention. Any future research claim should be narrower than the baseline, compared against the closest references, and supported by a theorem or reproducible experiment. See docs/novelty.md.
The current public release is v0.4.0. Source code is licensed under the Apache License 2.0. Citation metadata is provided in CITATION.cff.
@software{arndt_spectral_ebm_2026,
author = {Arndt, Justin},
title = {spectral-ebm},
year = {2026},
url = {https://github.com/j-arndt/spectral-ebm},
license = {Apache-2.0}
}v0.4.0 is a polished enterprise-integration proof-of-concept release: the local test suite passes, the GitHub Actions matrix passes on Python 3.10 and 3.12, and the public repository contains the raw evidence needed to reproduce the claims.