Neural Latent-Space Integrator

• Neural Latent-Space Integrator is a method employing adaptive orthogonal convolution layers to ensure norm-preservation in deep networks.
• It combines BCOP and RKO techniques to construct and fuse orthogonal kernels that support strides, dilations, and grouped convolutions efficiently.
• The approach enhances adversarial robustness and stabilizes training in applications like GANs, invertible flows, and robust classifiers with minimal computational overhead.

Below is a self-contained overview of Orthogonium (the open-source Python package) and its core Adaptive Orthogonal Convolution (AOC) layer. We assume familiarity with basic PyTorch/TensorFlow APIs but no prior knowledge of orthogonal convolutions.

1. Orthogonality in Convolutional Kernels

   • Let W be the 4-D convolutional weight tensor of shape W∈ℝ[c_out, c_in, k_h, k_w]. In “Toeplitz” (flatten-and-stride) form, the conv operation y=conv(W,x; stride=s) is equivalent to y=TK·x where TK∈ℝ[(c_out·H_out·W_out)×(c_in·H_in·W_in)] is a (usually rectangular) Toeplitz matrix.
   • We say W is row-orthogonal if (S_s W)(S_s W)ᵀ=I, and column-orthogonal if (S_s W)ᵀ(S_s W)=I, where S_s selects strided rows of the full convolution matrix. When c_in·s²=c_out or s=1, the two definitions coincide and imply that the layer is 1-Lipschitz (norm preserving): ∥y∥₂=∥x∥₂ and gradients do not vanish or explode.
   • Why it matters
     • Norm preservation stabilizes very deep nets and RNNs (avoiding vanishing/exploding gradients).
     • Orthogonal layers have Lipschitz constant 1 ⇒ provable adversarial robustness (certified bounds).
     • Invertible and constant‐Jacobian convolutions in normalizing flows and invertible ResNets.
     • Stable GAN training (WGAN critic must be 1-Lipschitz).

2. Mathematical Formulation & Parameterizations

   Suppose W∈ℝ[c_out, c_in, k_h, k_w]. We wish to enforce (S_s W)ᵀ(S_s W)=I (or the row version). Orthogonium offers two main tools:

   (a) BCOP-style explicit construction * Build W by fusing elementary orthogonal blocks: * 1×1 orthogonal conv from an orthogonal matrix M∈ℝ[c_out×c_in]. (M Mᵀ=I or MᵀM=I via any standard method: QR, SVD, Cayley, Björck.) * 1×2 and 2×1 orthogonal convs via “symmetric projector” N=M Mᵀ: stack([N, I−N], dim=−1) yields P∈ℝ[c×c×1×2] with P Pᵀ=I. * Compose these by block-convolution (⊞) to obtain any k_h×k_w kernel. * If A and B are orthogonal convs, their fused kernel K=B⊞A is orthogonal too (prop. 2.3 in the paper).

   (b) RKO-style reshape orthonormalization * Reshape W→W′∈ℝ[c_out,(c_in·k_h·k_w)] and run a matrix‐orthonormalizer on W′ (e.g. Björck iteration: W′←3/2 W′−1/2 W′(W′)ᵀW′ ). * Reshape back to [c_out,c_in,k_h,k_w]. If k_h=k_w=s=stride, this layer is strictly orthogonal (prop. 3.4).

   (c) (Optional) Cayley‐transform or SVD‐based tangent‐space updates * Parameterize W′ as Cayley(A)= (I−A)(I+A)⁻¹ for skew A, or do periodic SVD on the flattened weight matrix. These yield exact WᵀW=I but are more costly.

3. The AOC Method: Fusion + Extensions

   AOC combines BCOP and RKO to get native support for stride, dilation, groups and transposed convolutions, all in the spatial domain.

   Algorithm AOCConv2d build_kernel(in_c, out_c, k, s, groups, dilation):

   1. Let c_int ← max(in_c, ⌊out_c/s²⌋).
   2. Build K_bcop∈ℝ[c_int,in_c,k−s+1,k−s+1] via BCOP blocks:
      • Generate P₁,…,P_{k−s} (2×1 blocks), M (1×1), Q₁,…,Q_{k−s} (1×2 blocks).
      • Fuse by parallel associative scan (reducing O(k−s) to O(log(k−s))).
   3. Build K_rko′∈ℝ[out_c,(c_int·s·s)] by orthonormalizing an unconstrained matrix (Björck or QR). Reshape to K_rko∈ℝ[out_c,c_int,s,s].
   4. Fuse final kernel: K ← block_convolution(K_rko, K_bcop) (size=out_c×in_c×k×k, ready for stride=s, dilation, groups).
   5. Use K in a standard low-overhead conv2d call.

* Complexity: * BCOP block fusion: O(c_int·in_c·(k−s+1)²·log(k−s)) via parallel scan. * RKO reshape+orthonormalize: O(out_c²·c_int·s²) for Björck/QR. * Final block_convolution: implemented as a single 2D conv per output channel, overall O(c_out·c_int·in_c·k²). * All overheads are independent of the input H×W and scale sub-linearly in k.

* Transposed conv: simply flip channel dims and spatial dims of K and call ConvTranspose2d. Orthogonality is preserved. * Grouped conv: build each of g sub-kernels separately; Toeplitz matrix is block-diagonal, so global orthogonality holds iff each subkernel is orthogonal. * Dilation: orthogonality is invariant under dilation (with circular padding).

Pseudocode (PyTorch style)

class AOCConv2d(nn.Module): def __init__(self, in_c, out_c, kernel_size, stride=1, dilation=1, groups=1, bias=False, bj_iters=12): super().__init__() self.in_c, self.out_c = in_c, out_c self.k, self.s, self.d, self.g = kernel_size, stride, dilation, groups # intermediate channels self.c_int = max(in_c, out_c // (stride*stride)) # BCOP parameters: stored as small 1×1,1×2,2×1 blocks self.bcop_params = self._init_bcop_blocks() # RKO parameters: an unconstrained matrix for s×s kernel self.rko_mat = nn.Parameter(torch.randn(out_c, self.c_int * stride * stride)) self.bj_iters = bj_iters if bias: self.bias = nn.Parameter(torch.zeros(out_c)) def forward(self, x): # 1) orthonormalize rko_mat→K_rko Kp = self._rko_orthonormalize(self.rko_mat) K_rko = Kp.view(self.out_c, self.c_int, self.s, self.s) # 2) build K_bcop via parallel scan of 1×2,2×1 stacks K_bcop = self._build_bcop_kernel() # 3) fuse to full K K = block_convolution(K_rko, K_bcop) # uses our fast implementation return F.conv2d(x, K, self.bias, stride=self.s, dilation=self.d, groups=self.g, padding=self._pad())

1. Orthogonium Package & API

   Installation: pip install orthogonium

   Key modules: from orthogonium.conv import AOCConv2d, AOCConvTranspose2d from orthogonium.linear import OrthogonalLinear

   Example (replace nn.Conv2d in a ResNet block):

   import torch, torch.nn as nn from orthogonium.conv import AOCConv2d class Bottleneck(nn.Module): def __init__(self, in_c, mid_c, out_c, stride): super().__init__() self.conv1 = AOCConv2d(in_c, mid_c, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm2d(mid_c) self.conv2 = AOCConv2d(mid_c, mid_c, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(mid_c) self.conv3 = AOCConv2d(mid_c, out_c, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(out_c) # skip, activation, etc.

   Constructor arguments:

   • in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, bj_iters (default 12), tol (for Björck)
   • AOCConvTranspose2d has the same signature plus output_padding.

   Hyperparameters & trade-offs: * bj_iters∈[8..20]: more iters ⇒ stricter orthogonality but slower. * groups>1 yields lighter param count / faster conv. * Memory overhead: ~+5–15% versus nn.Conv2d. * Throughput slowdown: ~1.75× at small batch, ~1.13× at large (see Table 2).

2. Experimental Results

   A) Scalability on ImageNet 1K / ResNet-34 (input 224²) Batch size | Conv2D ref train ms | AOC train ms (×overhead) ––128–––––| 137 (1.00×) | 239 (1.75×) ––256–––––| 284 (1.00×) | 354 (1.25×) ––512–––––| 550 (1.00×) | 622 (1.13×)

   Memory: AOC adds +4–7% over Conv2D and supports larger batches before OOM.

   B) Certified Robustness / Expressivity on CIFAR-10 Model | Clean Acc | Provable Acc @ϵ=36/255 | #Params ––BCOP (orig paper) | 72.2% | 58.3% | 2.6 M ––SOC (20 its) | 78.0% | 62.7% | 27 M ––AOC (clean-opt) | 80.0% | 60.1% | 41.3 M ––AOC (robust-opt) | 74.0% | 64.3% | 41.3 M

   C) Large-scale training on ImageNet 1K

   • Cosine-similarity head (maximize clean accuracy): 68.2% top-1
   • Robust margin m=1.5√2: 42.1% clean / 26.3% certifiable

   These results show that AOC layers are competitive on small datasets, and—thanks to low input-size-independent overhead—enable practical large-scale 1-Lipschitz training.

3. Practical Integration, Limitations & Use Cases

   Integration

   • In PyTorch/Torchvision swap nn.Conv2d→AOCConv2d, nn.Linear→OrthogonalLinear.
   • Ensure circular padding or set padding=kernel//2 for tight spectral bound.
   • Use bj_iters≈12 and tune up to 20 if you require ‑99.99% orthonormality.

   Limitations * Small-kernel BCOP parametrization is incomplete: some 2×2 orthogonal kernels are not covered. In practice, higher-capacity nets and use of AOC’s fused branch remedy this. * At very small batch sizes (<32), overhead can reach ∼2×. Use larger batches for best throughput. * Slightly higher memory footprint (+5–15%).

   Primary Use Cases * Adversarially robust / 1-Lipschitz classifiers (certifiably robust). * Normalizing flows & invertible ResNets (constant-det Jacobian). * GAN critics (Wasserstein GANs with strict Lipschitz). * Spectrally-regularized RNNs (stable long-term gradients). * Any architecture needing strict norm preservation or stable training.

All of the above is implemented in the open-source Orthogonium library (https://github.com/thib-s/orthogonium). It provides drop-in, highly optimized, and thoroughly unit-tested modules for orthogonal conv/linear layers, and we invite you to incorporate AOCConv2d into your next project.