Co-domain symmetry and U(1)-equivariant networks

Complex-valued data often has a meaningful co-domain symmetry: the overall phase of a signal is arbitrary (the absolute time reference of a radar pulse, the global phase of an MR coil, the carrier phase of a comms symbol), so a sensible model should either ignore that overall phase (U(1)-invariance) or rotate its features along with the input (U(1)-equivariance) — and never silently break it.

complextorch.nn provides a coherent set of building blocks for both regimes. They originate from two papers:

• SurReal (Chakraborty, Xing, Yu — arxiv:1910.11334) — treats \(\mathbb{C} \setminus \{0\}\) as the Riemannian manifold \(\mathbb{R}^+ \times SO(2)\) and defines a weighted-Fréchet-mean (wFM) convolution and a tangent-space ReLU on it.

• Co-Domain Symmetry (CDS) (Singhal, Xing, Yu, CVPR 2022) — builds Invariant (“I-type”) and Equivariant (“E-type”) networks via phase-modulation layers, phase-thresholding activations, and a magnitude-only batch norm.

What’s “equivariant” and “invariant”?

For a global phase rotation by \(\psi\) — i.e. \(x \mapsto e^{j\psi}\,x\) — an operator \(M(\cdot)\) is

• U(1)-equivariant if \(M(e^{j\psi}\,x) = e^{j\psi}\,M(x)\) — output rotates the same way the input does.

• U(1)-invariant if \(M(e^{j\psi}\,x) = M(x)\) — output is unchanged.

A network is equivariant up to its classifier if every layer is equivariant; a fully invariant network ends with an invariant head (typically a distance-to-prototypes or a magnitude-only readout).

The headline correctness tests for each module live in tests/invariants/test_equivariance.py.

Module map

Strictly U(1)-equivariant

Module	Operation
complextorch.nn.PhaseShift	\(y = e^{j\phi}\,z\)
complextorch.nn.ComplexScaling	\(y = (\alpha + j\beta)\,z\)
complextorch.nn.MagBatchNorm1d / 2d / 3d	$y = z \cdot \mathrm{BN}(
complextorch.nn.EquivariantPhaseReLU	channel-mean-relative phase mask
complextorch.nn.wFMConvStrict2d	SurReal Eq. 14-16, convex weights (\(\sum w = 1, w \geq 0\))

Strictly U(1)-invariant

Module	Operation
complextorch.nn.PhaseDivConv1d / 2d / 3d	$y = x \cdot \overline{g(x)} / (
complextorch.nn.PhaseConjConv1d / 2d / 3d	\(y = x \cdot \overline{g(x)}\) (invariant when \(g\) is C-linear)
complextorch.nn.PrototypeDistance (E-type call)	logits unchanged when both \(z\) and the reference rotate by \(e^{j\psi}\)

Tangent-space / manifold operators

Module	Operation	Property
complextorch.nn.tReLU	\(r \mapsto \max(r, 1)\), \(\arg z \mapsto \max(\arg z, 0)\)	SurReal Eq. 21-22 — not equivariant by design, analogous to standard ReLU’s lack of translation-equivariance
complextorch.nn.wFMConv1d / wFMConv2d	port of RotLieNet ComplexConv2Deffgroup	approximate, with non-paper pre-modulation; prefer wFMConvStrict2d for paper-faithful math
complextorch.nn.wFMReLU	learned affine on $\log	z
complextorch.nn.wFMDistanceLinear	distance-to-Fréchet-mean head	output is real (invariants for classification)

Composing an equivariant block

Any composition of strictly U(1)-equivariant ops is itself U(1)-equivariant. For example,

import complextorch.nn as cnn

equivariant_block = nn.Sequential(
    cnn.wFMConvStrict2d(in_channels, out_channels, kernel_size=3, padding=0),
    cnn.ComplexScaling(out_channels),
    cnn.EquivariantPhaseReLU(out_channels),
    cnn.MagBatchNorm2d(out_channels),
)

To make the whole network invariant for classification, follow such a stack with a complextorch.nn.PrototypeDistance head and pass a reference vector (typically a sum-pool of the features) so that prototypes co-rotate with the input:

features = backbone(x)                           # equivariant
z = features.flatten(2).mean(dim=2)              # [B, C], still equivariant
ref = features.sum(dim=1, keepdim=False).mean()  # [B, 1], rotates with x
logits = head(z, reference=ref)                  # invariant

This is exactly the pattern used by complextorch.models.CDSEquivariant.

Where invariance comes from

• PhaseDivConv uses the fact that a global phase rotation cancels in numerator and denominator: \((e^{j\psi}\,x) \cdot \overline{e^{j\psi}\,g(x)} / |e^{j\psi}\,g(x)|^2 = x \cdot \overline{g(x)} / |g(x)|^2\).

• PhaseConjConv is invariant for the same reason when \(g\) is complex-linear (the cfloat-native default). The CDS paper described it as “phase-mixing” because the reference code decomposed \(g\) as two real convs (which still happens to be C-linear, so the same cancellation applies; the “phase-mixing” framing was loose).

• PrototypeDistance with reference=: when both \(z\) and the reference rotate by \(e^{j\psi}\), the per-prototype distances \(|z - y \cdot p_k|\) are unchanged, so logits are invariant.

Caveats

• Padding breaks equivariance. nn.Unfold(padding > 0) zero-pads reals/phases independently. A zero magnitude cannot be rotated meaningfully — the \((\log|0|, \arg 0)\) representation is ambiguous. wFMConvStrict2d is strictly equivariant only with padding=0; with positive padding, boundary positions degrade. The docstring carries the full note.

• tReLU is not equivariant. This mirrors standard ReLU not being translation-equivariant in real-valued nets — the tangent-space ReLU is a principled lift, not a free lunch.

• wFMConv2d (existing) ≠ wFMConvStrict2d (new). The former is the port of RotLieNet’s ComplexConv2Deffgroup, with a non-paper pre-modulation and fold(unfold(·)) smear; the latter is the paper-faithful Eq. 14-16 implementation. Prefer the strict variant when you need provable equivariance.