Sheaf-PC

JAX implementation accompanying the NeurReps 2025 paper Sheaf Cohomology of Linear Predictive Coding Networks.

Sheaf-PC represents a linear predictive-coding network as a cellular sheaf. The coboundary maps node activations to edge-wise prediction errors, inference follows diffusion under the sheaf Laplacian, and a Hodge decomposition separates diffusive and harmonic error. The predictive-coding energy is edge-factorized, with one prediction-error term per connection.

Installation

Sheaf-PC supports Python 3.10-3.12. Install uv, then run:

uv sync

The experiments generate a noisy Gaussian identity task on demand; no dataset download is required.

Quick start

Run the smoke test:

uv run python train.py --config configs/smoke.yaml

Run a 10-hidden-layer knotted network with theta = 0:

uv run python train.py \
  --config configs/knotted.yaml \
  --theta 0.0 \
  --output-dir results/knotted_theta0

Each run writes config.json, metrics.csv, and summary.json to its output directory.

Reproducing the experiments

# Figures 2-4: knotted-network feedback-angle sweep
uv run python scripts/run_theta_sweep.py

# Appendix C: all-to-all network-size sweep
uv run python scripts/run_all_to_all.py

# Render the figures
uv run python scripts/plot_figures.py

Experiment	Hidden state	Batch	Target noise	Learning rate	Steps
Knotted theta sweep	10 x 2D	128	0.001	0.1	1000
All-to-all size sweep	2-15 x 4D	64	0.01	1.0	1000

Sweep results are written under results/theta_sweep/ and results/all_to_all/; rendered figures are written to results/figures/.

The sweep scripts accept comma-separated seeds:

uv run python scripts/run_theta_sweep.py --seeds 0,1,2,3,4

Python API

import jax

from sheaf_pc import boundary_ids, knotted_chain, relative_coboundary

sheaf = knotted_chain(jax.random.PRNGKey(0), theta=0.33)
x_id, y_id = boundary_ids(sheaf)
D, C, free, clamped = relative_coboundary(sheaf, (x_id, y_id))

Importing sheaf_pc enables JAX 64-bit mode and highest matmul precision for the Hodge and spectral calculations.

Checks

uv run ruff check .
uv run python -m pytest -q
uv build

Citation

@inproceedings{seely2025sheaf,
  title         = {Sheaf Cohomology of Linear Predictive Coding Networks},
  author        = {Seely, Jeffrey},
  booktitle     = {NeurIPS 2025 Workshop on Symmetry and Geometry in Neural Representations (NeurReps)},
  year          = {2025},
  eprint        = {2511.11092},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2511.11092}
}

Machine-readable metadata is available in CITATION.cff.

License

Sheaf-PC is released under the MIT License.