When Geometry Aligns: Dihedral Hidden-State Transformations in UNet, ViT, and DiT Architectures

Authors: Mojtaba Faramarzi, Alex Lamb, Irina Rish
Paper: https://arxiv.org/abs/2607.03580
Code: N/A
Model: N/A

TL;DR

WHAT was done? The authors introduce a theoretical and empirical framework for applying spatial reflection interventions directly to intermediate representations within non-equivariant vision backbones, including convolutional U-Nets, Vision Transformers (ViTs), and Diffusion Transformers (DiTs). They formalize geometric consistency across coupled computational pathways—such as multi-head attention projections and U-Net skip connections—as the necessary condition for maintaining internal feature stability during hidden-state transformations.

WHY it matters? As generative AI moves toward direct feature-level editing, activation steering, and representation tuning, manipulating internal activations inside standard backbones frequently introduces unseen spatial coordinate mismatches. By demonstrating that internal transformations remain stable and act as capacity-reducing regularizers when interacting pathways stay aligned in a common coordinate frame, this work provides clear structural guidelines for performing safe feature-space augmentations and modular edits without requiring full architectural redesigns or custom group-equivariant networks.

Executive summary: Modern machine learning models are increasingly edited, steered, or augmented at the level of internal activations rather than input images. However, modifying a hidden layer inside standard architectures often breaks implicit spatial assumptions. This paper reveals that internal geometric manipulations—such as flipping intermediate feature maps—do not inherently disrupt model computation. Instead, failure occurs specifically when transformations are applied partially across coupled branches, creating coordinate frame mismatches that compound across iterative sampling steps in diffusion models. By enforcing geometric consistency across all interacting pathways, researchers can safely perform feature-space interventions that improve feature stability and generalization across U-Net, ViT, and DiT architectures.

Details

The Hidden Spatial Inconsistency in Neural Representation Steering

Modern vision and generative models rely on deep feature hierarchies where intermediate hidden states encode rich spatial geometry. While prior literature has extensively explored input-level data augmentations and post-hoc model editing techniques, directly applying geometric transformations to intermediate activation tensors inside pretrained backbones remains poorly understood. When an internal representation is altered, downstream layers expect a coherent coordinate frame to process spatial relationships correctly. Standard architectures like convolutional U-Nets or Diffusion Transformers (DiTs) are not inherently group-equivariant; hence, unaligned internal manipulations risk injecting catastrophic spatial mismatch into downstream operations. The core bottleneck investigated in this paper is identifying what structural conditions allow an internal geometric intervention to preserve coherent computation across multi-head attention blocks and encoder-decoder skip connections.

Existing approaches to spatial invariance typically enforce strict architectural equivariance by redesigning layers from scratch, which prevents their application to standard pretrained weights. Conversely, post-hoc editing frameworks often modify intermediate activations without accounting for multi-branch coordinate alignment. The critical delta introduced by this work is the principle of geometric consistency: rather than modifying the underlying neural architecture, the paper demonstrates that standard backbones can accommodate direct hidden-state transformations during fine-tuning provided that all interacting computational pathways operating on the transformed feature map are transformed in unison. By contrasting geometrically consistent interventions against inconsistent partial transformations across baseline architectures such as Proszewska et al. (2025) and Yang et al. (2024), the authors establish a foundational framework for hidden-state geometric regularization.

Geometrically Inconsistent (Output Flip Only):
Head 1 (Standard) ---> [O_1] ----------------┐
Head 2 (Flipped)  ---> [O_2] ---> [Flip] ---> Concat ---> Project [W^O] ---> [Distorted Frame]

Geometrically Consistent (Coherent Transformation):
Head 1 (Standard) ---> [O_1] --┐
Head 2 (Standard) ---> [O_2] --┴-> Concat --> [Flip All] --> Project [W^O] --> [Aligned Frame]

Geometric Consistency First Principles: Coordinate Frames across Interacting Pathways

To formalize hidden-state interventions, let z^(ℓ) denote the intermediate activation tensor at layer or block ℓ. The theoretical substrate of this work considers spatial transformations drawn from the reflection subset of the dihedral group D₄​, defined as T={T_hor​,T_ver​,T_diag​,T_anti​}⊂D₄​, representing horizontal, vertical, main-diagonal, and anti-diagonal flips. For a 2D feature map or patch grid arranged on a P×P spatial lattice, a transformation τ∈T remaps token indices via a permutation matrix Π_T​. The state space is thus transformed from F^(ℓ) to F̃^(ℓ)=τ(F^(ℓ)).

The core theoretical assumption asserts that an intervention is geometrically consistent if and only if every interacting branch consuming F̃^(ℓ) operates within the identical spatial frame. In multi-head self-attention, input tokens X∈R^N×d_model​ with N=P² generate query, key, and value matrices Q_h​,K_h​,V_h​∈R^N×d_v​ for head h. A single head’s attention map A(Q,K,V)=softmax(QK^⊤/sqrt(d_k​)​)V satisfies permutation equivariance under simultaneous transformation A(Π_T​Q,Π_T​K,Π_T​V)=Π_T​A(Q,K,V). However, when multi-head outputs are combined via projection matrix W^O=[W₁^O​…W_H^O​], consistency requires that every interacting head shares the permutation Π_T​.

In convolutional U-Nets, an analogous principle governs skip-connected encoder-decoder fusion. Let F_enc^(ℓ)​∈R^{C_enc​×H×W} denote the encoder feature map and F_dec^(ℓ′)∈R^C_dec×H×W denote the corresponding decoder feature map at matched spatial resolution. If ϕ represents a skip-fusion operator, such as concatenation followed by convolution, consistency requires transforming both branches simultaneously so that G^(ℓ,ℓ′)=ϕ(τ(F_dec^(ℓ′)​),τ(F_enc^(ℓ)​))=Π_T​ϕ(F_dec^(ℓ′)​,F_enc^(ℓ)​). Applying the transformation to only one branch breaks this relation, generating spatial coordinate mismatch at fusion. From a statistical perspective, constraining hypothesis class H to a symmetry-consistent subclass H_sym​={h∈H:h(x)=h(T(x))} yields empirical Rademacher complexity Ȓ_S(H_sym)≤Ȓ_S(H), establishing geometric consistency as an idealized form of capacity-reducing regularization.

From Patch Grid to Attention Projection: Mechanistic Flow of Flipped Heads

To trace the precise mechanistic flow of a geometric intervention, consider a single input tensor passing through a Vision Transformer block, such as ViT-B/16. On a patch grid of dimensions P×P, token index t(i,j)=(i−1)P+j maps spatial coordinate (i,j) to a flat sequence index. Standard self-attention projects tokens into head outputs O_h∈R^N×d_v​. Under a targeted flipped-head intervention on head h^∗, the output matrix O_h^∗​ is reshaped into grid tensor O_h^∗​∈R^P×P×d_v​, transformed via horizontal flip (T_hor​O_h^∗​)[i,j,:]=O_h^∗​[i,P−j+1,:], and flattened back into token format Õ_h^∗​.

When this transformed head is combined with untransformed heads in an attention-inconsistent module, the final token representation at spatial location (i,j) becomes

\(y_{t(i,j)} = \sum_{h \neq h^*} W_h^O o_h[t(i, j)] + W_{h^*}^O o_{h^*}[t(i, P - j + 1)]\)

This operation injects mirrored contextual information from head h^∗ while keeping all remaining heads anchored at the original spatial location.

As visualized in Figure 5, an output-only flip creates significant coordinate distortion heatmaps ∣A_incon​−A_con​∣, whereas a fully consistent QKV transformation preserves or correctly mirrors spatial geometry across the attention field.

Input Latent (x_t) 
       │
       ▼
[Encoder Blocks] ─── F_enc^(l) ───┐ (Consistent: Apply τ to both)
       │                          ├───> Fusion φ(τ(F_dec), τ(F_enc)) ───> [Decoder]
[Decoder Blocks] ─── F_dec^(l') ──┘

When propagated through deeper layers or unrolled across T iterative denoising steps in diffusion models, these per-step attention perturbations accumulate linearly. Linearizing the denoiser update z_t−1​=Ψ_t​(z_t​,ε̑_t​) shows trajectory deviation accumulating as Δ_T​≈∑_t=1:T​D_t​J_t​E_t​, where J_t​=∂ε̑_t​/∂O_t​ is the Jacobian of the noise predictor and D_t​ measures sampler sensitivity. In single-pass architectures like ViT, spatial routing adjusts smoothly to mirror target regions without losing semantic focus, as shown in Figure 4. However, in iterative Diffusion Transformers, partial mismatches compound into severe global displacement and feature drift.

Optimization, Loss Formulation, and the 50% Stochastic Intervention Schedule

Implementing hidden-state geometric interventions requires minimal modifications to standard optimization pipelines. The training framework operates on noisy latents x_t​ at diffusion step t with injected noise ε∼N(0,I). For a chosen intermediate block ℓ sampled uniformly from available target candidates B^∗∼Uniform(C) and a transformation τ∼T, the modified activation F̃^(ℓ)(x_t​,t)=τ(F^(ℓ)(x_t​,t)) is propagated through the network to produce output ε_θ^(τ,ℓ)​(x_t​,t). The primary loss function minimizes noise-prediction error under hidden-state augmentation:

\(\mathcal{L}_{\text{aug}}(\theta) = \mathbb{E}_{x_0, \varepsilon, t, \ell, \tau} \left[ \|\varepsilon - \varepsilon_\theta^{(\tau,\ell)}(x_t, t)\|_2^2 \right]\)

In the main Stable Diffusion 2.1 U-Net testbed, fine-tuning is conducted using the AdamW optimizer with a learning rate of 10⁻⁵, batch size 4, and a standard DDPM scheduler over 5,000 steps on 4×64×64 VAE latents derived from 500 Oxford-IIIT Pet images resized to 512×512. Candidate transformation locations C encompass ResNet blocks (C_res​), attention blocks (C_attn​), and skip-fusion junctions (C_skip​).

A key practical finding is the effectiveness of a 50% stochastic intervention schedule. Applying the hidden-state transformation to every mini-batch risks over-regularizing the backbone, whereas applying it too infrequently leaves representations unconstrained. The 50% consistent schedule randomly interleaves intervention mini-batches with standard clean baseline updates, striking an optimal balance between maintaining functional denoising accuracy and enforcing geometric stability.

Quantitative Validation: Diagnostics, Ablations, and Internal Drift Probing

To measure internal representation stability independently from final image output quality, the authors introduce three activation-level diagnostic metrics computed after realigning the intervened activation back to the clean frame via A_align=T⁻¹(Ã). Self-Consistency Shift (SCS) measures feature-level L1 distance

\(\text{SCS} = \frac{1}{BCHW} \|A - A^{\text{align}}\|_1\)

Activation Mass Scatter (AMS) measures spatial dispersion around the normalized spatial center of mass μ_b​ via

\(\text{AMS} = \frac{1}{B} \sum_{b=1}^B \sum_{h,w} \bar{m}_b(h, w) \|(h, w) - \mu_b\|_2^2\)

Drift captures global spatial displacement between clean and aligned mass centers

\(\text{Drift} = \frac{1}{B} \sum_{b=1}^B \|\mu_b - \mu_b^{\text{align}}\|_2^2\)

Denoising quality is tracked simultaneously via Noise-Prediction MSE

\(\text{NP-MSE} = \mathbb{E}_{x_t, t} [\|\hat{\varepsilon}_\theta(x_t, t) - \varepsilon\|_2^2]\)

As illustrated across the time-series evaluations in Figure 1, inconsistent intervention modes (attention-inconsistent and inconsistent-skip) suffer severe stability degradation: SCS spikes above 50 and AMS exceeds 35, accompanied by increased NP-MSE noise-prediction error. In contrast, 50% consistent variants maintain lower feature shift, reduced spatial scatter, and minimal spatial drift while achieving superior noise prediction.

Image-space evaluations across 105,000 generated samples presented in Figure 2 demonstrate that 50% consistent modes maintain competitive FID, KID, CLIP text-alignment scores, and LPIPS diversity compared to the baseline, confirming that internal stability gains do not sacrifice visual quality.

Qualitative uncurated samples across all seven modes are illustrated in Figure 6.

To explain why attention fields exhibit variable stability under reflection, the authors conduct probing analyses on synthetic 4×4 and 8×8 grid patterns using probe consistency score

\(S^{(\ell)} = \frac{1}{M} \sum_{i=1}^M \text{cosine}(\mathcal{A}^{(\ell)}(x_i), \mathcal{A}^{(\ell)}(T(x_i)))\)

As shown in Table 3 and Figure 3, ViT Block 5 attention fields yield moderate consistency on directional patterns like horizontal stripes (0.523–0.581) and diagonal gradients (0.624–0.650), but achieve perfect consistency (1.000) on isotropic Gaussian blobs, revealing inherent geometric anisotropy in attention layers.

Complementary proof-of-concept FID experiments summarized in Table 1 highlight performance gains when applying 50% consistent regularization across standard benchmarks. On CIFAR-10 (32×32 DDPM baseline from Proszewska et al. (2025)), regularized training reduces FID from 4.46 to 4.30±0.05. On CelebA-64 DDPM, FID reaches 4.53±0.33 compared to the 4.51 baseline. On MNIST using the spiking transformer backbone SDiT (Yang et al. (2024)), FID improves from 5.54 to 5.29±0.08.

Furthermore, as detailed in Table 2, fine-tuning a pretrained ViT-B/16 on CIFAR-100 with input-or-hidden (I/H) reflection augmentation improves top-1 classification accuracy from 92.30% (standard 224×224 input augmentation) to 93.19%.

Contextualizing Geometric Interventions in Modern Vision Backbones

The findings in this paper build upon three major lines of deep learning research: spatial data augmentation, group equivariant network design, and post-hoc hidden-state editing. Traditional data augmentation techniques, such as those surveyed by Shorten & Khoshgoftaar (2019) or feature-space block masking like PatchUp (Faramarzi et al., 2022), operate primarily at input layers or apply unconstrained noise to hidden states. While effective for classification, unconstrained feature noise lacks the structural alignment required to preserve fine spatial details in generative models.

In terms of architecture, group equivariant neural networks (Cohen & Welling, 2016; Finzi et al., 2021) enforce exact symmetry constraints by replacing standard convolutions and attention layers with group-steerable operators. While mathematically elegant, steerable backbones cannot leverage massive pretrained weights from standard non-equivariant architectures like Stable Diffusion or DiT. On the other hand, hidden-state editing techniques (Meng et al., 2022; Hertz et al., 2022) alter intermediate activations to steer text-to-image synthesis or factual knowledge, but rarely consider whether multi-head attention projections or skip connections remain aligned in a shared spatial frame. This paper directly bridges these domains by showing that standard, non-equivariant models can accommodate internal geometric change during fine-tuning without architectural modifications, provided that coupled computational paths remain geometrically consistent.

Scope Constraints: Reflection Subsets and Latent Inconsistencies

Despite its strong theoretical and empirical insights, several limitations should be noted. First, the implemented transformation set is strictly restricted to the reflection subset T={T_hor​,T_ver​,T_diag​,T_anti​} of the dihedral group D₄​. Rotations were explicitly excluded from the main diffusion experiments because VAE latent representations in Stable Diffusion are not rotation-consistent; rotating latent grids introduces severe reconstruction artifacts when passed through pretrained VAE decoders. Extending geometric consistency to full continuous rotation groups or affine transformations will require specialized alignment mechanisms.

Second, the main quantitative U-Net evaluation relies on a lightweight subset of 500 images from the Oxford-IIIT Pet dataset. While this setup allowed rigorous evaluation across seven intervention modes, three random seeds, and 105,000 generated images, testing on larger-scale datasets like ImageNet-1k or LAION remains an important next step. Third, the Rademacher complexity analysis provides an idealized, capacity-reducing regularization argument for function classes rather than a strict quantitative generalization bound for non-linear, multi-step diffusion sampling pipelines.

Strategic Assessment and Architectural Implications

This work establishes geometric consistency as a foundational design principle for feature-space interventions in spatially structured neural networks. The central insight—that internal manipulations fail due to multi-branch coordinate frame mismatch rather than the transformation itself—has significant implications for parameter-efficient fine-tuning (PEFT) methods like LoRA, activation steering, and modular editing. When designing internal adapters or editing modules for multi-head attention and U-Net backbones, enforcing coherent spatial coordinate alignment across all interacting pathways prevents catastrophic feature drift and trajectory error accumulation.

For researchers working on vision and generative backbones, incorporating stochastic 50% geometrically consistent hidden-state augmentations during fine-tuning offers a lightweight, parameter-free regularizer that improves feature stability and generalization without requiring costly architectural redesigns. Future work should expand this framework to continuous spatial transformations, explore rotation-consistent latent spaces, and integrate geometric alignment constraints into real-time activation steering protocols.