EGN — Equivariant Geodesic Networks

A universal classifier on the SPD manifold, packaged the way a convolutional network is.

import numpy as np
from egn import EGNClassifier

X = np.random.randn(512, 22, 256)      # 512 trials, 22 channels, 256 samples
y = np.random.randint(0, 4, 512)

clf = EGNClassifier(epochs=30).fit(X, y)
print(clf.score(X, y))

That is the whole API for the common case. The matrix size, the number of manifold channels, the class count and the label vocabulary are all inferred from the data, so the same object trains on EEG epochs, skeleton covariances, radar returns and image descriptors without a per-dataset subclass.

pip install egn        # torch and numpy are the only dependencies

---

What it does

Every intermediate representation is an exact symmetric positive definite matrix under the affine-invariant Riemannian metric. There is no projection step in the forward pass, and constrained parameters are updated by retraction, so a Stiefel frame stays orthonormal and an SPD prototype stays positive definite regardless of the step size.

x ──ToSPD──► (B, K, d, d) ──EGNBlock ×L──► (B, K′, m, m) ──pool──► (B, m, m) ──head──► logits

The correspondence with a CNN is deliberate, and the intuitions transfer:

CNN	EGN
Conv2d(c_in, c_out, k)	BiMap(d_in, d_out, c_in, c_out)
channels	manifold channels (branches)
stride / downsampling	matrix size reduction d_in → d_out
BatchNorm2d	SPDBatchNorm (whitens by a Fréchet mean)
ReLU	SpectralActivation (acts on eigenvalues)
Dropout	GeodesicDropout (interpolates towards the identity along a geodesic)
bias	GeometricBias (a metric isometry, not an addition)
global average pool	RiemannianPool (Fréchet or log-Euclidean barycentre)
linear classifier	TangentHead or GeodesicPrototypeHead

---

Input conventions

ToSPD accepts, and infers when input_kind="auto":

input	meaning	output
(B, n, n)	covariance / descriptor already on the manifold	(B, 1, n, n)
(B, K, n, n)	multi-branch SPD	(B, K, n, n)
(B, C, T)	multichannel signal	(B, K, C, C)
(B, T, D)	sequence of features (input_kind="sequence")	(B, K, D, D)
(B, C, H, W)	feature map / image	(B, K, C, C)

branches=K splits the sample axis into K windows and forms one covariance per window — the manifold equivalent of a multi-channel stem.

Rank-2 input is rejected with an explanatory error rather than being turned into a singular rank-one outer product, which would only fail later inside a logarithm.

---

Why the GPU is fast now

The earlier research code ran the whole network in float64 and was slower on GPU than on CPU. Six things caused that, and all six are addressed:

1. float64 everywhere. Consumer and inference-class GPUs execute double precision at 1/32 of their float32 rate; a T4 in fp64 is genuinely slower than a decent CPU. The dtype policy is now config.spectral_dtype = "auto" — float32 on CUDA, float64 on CPU. Set it to "float64" only to reproduce a theory check.
2. Host synchronisation inside the forward pass. The old Fréchet mean called .item() on a residual every iteration and the spectral activation called .item() on its threshold, draining the CUDA queue several times per layer. Nothing in the forward pass calls .item() any more; the mean runs a fixed iteration budget.
3. Repeated eigendecompositions. sqrtm_pair returns S^{1/2} and S^{-1/2} from one decomposition; the prototype head whitens by each prototype once per forward instead of once per (sample, prototype) pair; distances use eigvalsh, which never forms eigenvectors.
4. Batching. Every operator takes an arbitrary leading shape and issues exactly one eigh per call. config.eig_chunk caps the batch when memory, not throughput, is the constraint.
5. Per-step metric readback. Training statistics accumulate on the device and are read once per epoch.
6. Pooling cost. The default pool is the closed-form log-Euclidean barycentre; the iterative Fréchet mean is one flag away (pool="frechet") when the channel spread makes it worth the iterations.

Measure your own machine rather than trusting any of this:

python -m egn.benchmark --sizes 16 32 64 --batch 256 1024 4096

The output separates eigh throughput per dtype from end-to-end model throughput, which is what tells you whether you are precision bound or launch-latency bound. If throughput is flat in the batch size, the kernels are too small to saturate the device — raise the batch before anything else.

The honest caveat. These are small-matrix, decomposition-heavy workloads. A GPU wins decisively at large batch sizes and float32; at batch 32 with 8×8 matrices it may still lose to a CPU, because the kernels never fill the device. That is a property of the operation, not of this implementation.

---

Scaling out

Distributed data parallel is a launch flag, not a rewrite:

torchrun --nproc_per_node=4 examples/train_ddp.py

EGNClassifier detects the process group, wraps the model in DistributedDataParallel, installs a DistributedSampler and reduces metrics across ranks. The model is materialised before wrapping, so the replicas have parameters to broadcast.

DDP works with the Riemannian optimiser without special handling: DDP all-reduces the Euclidean gradients in its backward hook, and every rank then applies the same deterministic retraction to the same parameters, so replicas stay identical.

data_parallel=True uses nn.DataParallel for a quick single-process multi-GPU run. It is not recommended — it re-scatters the model every step, which on a network of short kernel launches costs more than it saves.

---

Building your own architecture

import torch.nn as nn
from egn.nn import EGNBlock, GeodesicPrototypeHead, RiemannianPool, ToSPD

model = nn.Sequential(
    ToSPD(kind="signal", branches=4),
    EGNBlock(22, 16, in_channels=4, out_channels=8, mix=True),
    EGNBlock(16, 8,  in_channels=8, out_channels=8),
    RiemannianPool("frechet", iters=5),
    GeodesicPrototypeHead(8, num_classes=4),
)

Or use the factories, which follow the torchvision convention:

from egn import egn_tiny, egn_small, egn_base
model = egn_base(num_classes=4)      # deeper trunk, mixed channels, geodesic head

Choosing a head

TangentHead (default) takes one logarithm per sample and applies a linear classifier in the tangent space at a learnable reference point. Its cost is independent of the class count.

GeodesicPrototypeHead scores by squared geodesic distance to trainable SPD prototypes, p(c | Σ) = softmax(−d²(Σ, P_c)/τ). Fully geometric, and its cost grows with the number of prototypes. Use it when you want prototypes you can inspect, or when classes are naturally described as regions of the manifold.

Both are invariant in the joint sense: congruencing the input and the reference by the same matrix leaves the logits unchanged. Invariance while holding the prototypes fixed is false — congruencing only the input changes every distance.

---

Geometry as a public API

egn.geometry and egn.functional are usable on their own, with no model involved:

from egn.geometry import distance, frechet_mean, geodesic, riemannian_log

d = distance(A, B)                      # affine-invariant geodesic distance
M = frechet_mean(batch, dim=1)          # Riemannian barycentre
mid = geodesic(A, B, 0.5)               # midpoint on the manifold

Everything is differentiable, including through repeated eigenvalues: the backward pass uses the Loewner divided-difference matrix, and clamped eigenvalues receive a zero subgradient instead of an unbounded one.

---

Tests

pytest -q

The suite asserts the geometry, not just the shapes: exp/log inverse to machine precision, affine invariance of the distance, geodesic constant speed, the Fréchet mean as a fixed point and its equivariance, gradcheck on every spectral operator including an exact eigenvalue tie, and the closed-form prototype gradient against autograd. It also asserts the counter-examples — that a convex combination is not a geodesic, and that a rectangular BiMap is not output-congruent — because those are the claims that are easy to overstate.

---

Citation

@software{khan2026egn,
  author = {Khan, Md Raihan},
  title  = {EGN: Equivariant Geodesic Networks on the SPD manifold},
  year   = {2026},
  url    = {https://github.com/kraihan/egn}
}

License

MIT. This package contains no third-party research code; see MIGRATION.md if you are coming from the original research repository.