---
title: Neural Latent-Space Integrator
url: https://www.emergentmind.com/topics/neural-latent-space-integrator
type: topic
---

# Neural Latent-Space Integrator

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)  
    ```python
    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())
    ```  

4.  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):  
    ```python
    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).  

5.  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.  

6.  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.

Source: https://www.emergentmind.com/topics/neural-latent-space-integrator