Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform

Zhongchen Zhao; Jixin Wang; Qi Xie; Hui Lin; Lei Zhang; Deyu Meng; Zongben Xu

cs.CVJul 23, 2026

Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform

Authors: Zhongchen Zhao, Jixin Wang, Qi Xie, Hui Lin, Lei Zhang, Deyu Meng, Zongben Xu

Organizations: 1Xi’an Jiaotong University · 2The Hong Kong Polytechnic University · 3OPPO Research

Abstract

Equivariant networks embed geometric symmetries as structural priors through weight sharing, achieving remarkable parameter efficiency across vision tasks. However, this parameter efficiency does not translate into compute efficiency: existing implementations unroll the structured weights into dense matrices and dispatch them to generic dense kernels, so the FLOPs of an equivariant layer are no smaller than those of a non-equivariant counterpart. In this paper, we observe that the equivariant linear (EQ-Linear) layer---the most fundamental and frequently used module in modern equivariant architectures---is essentially a circular convolution along the group dimension composed with a linear transform along the channel dimension. Building on this observation, we propose Flash EQ-Linear, an exact acceleration algorithm that reduces the complexity from $\mathcal{O}(NDC)$ to $\mathcal{O}(NDC/T)$ by combining the Fourier convolution theorem along the group dimension with the conjugate symmetry of the real DFT. We further provide dedicated CUDA kernels for Flash EQ-Linear, covering both forward and backward passes and both FP32 and FP16 precision. At the operator level, Flash EQ-Linear achieves up to ${2\times}$ forward speedup over PyTorch's F.linear; at the network level, Flash EQ-ViT and Flash EQ-Swin achieve up to ${1.7\times}$ end-to-end speedup over both equivariant and non-equivariant baselines. To our knowledge, this is the first time equivariant networks strictly dominate their non-equivariant counterparts along all three axes simultaneously: accuracy, parameter efficiency, and inference speed.Code is available at https://github.com/zhongchenzhao/FlashEQLinear.

Explore similar work

May 21, 2025cs.CV

Quick ViTs: Speeding up Vision Transformers through Equivariance

Natural images exhibit strong geometric regularities: local structures, such as edges, corners, and textures, appear in many orientations and mirror configurations. Since Vision Transformers (ViTs) operate on square image patches, these transformations naturally correspond to the dihedral symmetry group

\mathrm{D}_8

, also known as the octic group. Recent work has shown that ViTs can be made reflection equivariant and more efficient than standard ViTs simultaneously by implementing the linear layers in the Fourier domain of the reflection group. In this work, we extend the equivariance to reflections and rotations and analyze the scalability of the resulting networks. Our Quick ViTs, based on octic equivariant linear layers, achieve 5.33x reductions in FLOPs and up to 8x reductions in memory compared to ordinary linear layers. By analyzing the arithmetic intensity of these layers, we identify theoretical limits on how much the FLOP savings translate into throughput improvements on modern GPUs. However, these limitations disappear as the embedding dimensions increase. Enabled by their computational efficiency, we conduct a broader empirical evaluation of equivariant ViTs than in previous work. Upon training supervised (DeiT-III) and self-supervised (DINOv2) on ImageNet-1K, we find that our Quick ViTs match or exceed baseline accuracy while at the same time providing substantial efficiency gains.

David Nordström, Johan Edstedt, Fredrik Kahl +1

Self-Supervised Vision Transformers Vision Transformer Equivariant Neural Networks Imagenet Equivariance Real-Time Inference+1

Ask PDF

May 5, 2026cs.CV

Normalization Equivariance for Arbitrary Backbones, with Application to Image Denoising

Normalization Equivariance (NE) is a structural prior that improves robustness to distribution shift in image-to-image tasks. A function

f

is normalization equivariant iff

f(a y + b\mathbf{1}) = a f(y) + b\mathbf{1}

for all

a>0

and

b\in\mathbb{R}

. Existing NE methods constrain every internal layer to NE-compatible operations. These constraints add runtime cost and exclude standard transformer components such as softmax attention and LayerNorm. We introduce Wrapped Normalization Equivariance (WNE), a parameter-free wrapper that normalizes the input, applies any backbone, and denormalizes the output. We prove every NE function admits this factorization, so the wrapper exactly parameterizes the class of NE functions. On blind denoising, wrapping CNN and transformer architectures improves robustness under noise-level mismatch with no measurable GPU overhead, while architectural NE baselines are up to

1.6\times

slower.

Youssef Saied, François Fleuret

Exact Equivariance Batch Normalization Image Denoising Normalization Softmax Attention

Ask PDF

Sep 8, 2026cs.LG

Equivariance Breaks the Learning Rate

Equivariant networks are commonly trained with Adam, yet recent work reports that matrix-structured optimizers such as Muon can perform better on these architectures without explaining why. We identify one source of this difference inside equivariant linear layers. Each irrep block learns a channel-mixing matrix

W_l

shared across its

2l+1

components, giving the expanded map

W_l \otimes I_{2l+1}

. For a single application of the layer, the gradient of

W_l

sums

2l+1

outer product contributions and has rank at most

2l+1

. Adam rescales stored weights individually without using the irrep boundaries, so one learning rate can produce different spectral step sizes across blocks within a layer. We address this mismatch by normalizing each block update separately, without introducing a new hyperparameter. This changes only the scale of the update, leaving Adam's moment estimates and its direction within each block unchanged. We evaluate the mechanism in a controlled

\mathrm{SO}(3)

-equivariant model with a matched dense control and in an e3nn interatomic potential model trained on rMD17 and MD22. The toy setup isolates a mismatch that grows with width while the dense control shows no corresponding growth. In the interatomic potential model, block normalization and tuning Adam's momentum coefficients independently improve performance, but neither alone matches Muon. Combined, they make Adam competitive with Muon on all datasets, indicating that blockwise step control and momentum accumulation account for much of Muon's advantage.

Andrei Manolache, Mathias Niepert

Adam Equivariant Neural Networks Equivariance Learning Rates Muon Model Weights+1

Ask PDF