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

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.

Auxiliary Loss Adaptation for Image Inpainting

Auxiliary losses commonly used in image inpainting lead to better reconstruction performance by incorporating prior knowledge of missing regions. However, it usually requires a lot of effort to fully exploit the potential of auxiliary losses, or otherwise, improperly weighted auxiliary losses would distract the model from the inpainting task, and the effectiveness of an auxiliary loss might vary during the training process. Hence the design of auxiliary losses takes strong domain expertise. To mitigate the problem, in this work, we introduce the Auxiliary Loss Adaptation for Image Inpainting (ALA) algorithm to dynamically adjust the parameters of the auxiliary loss. Our method is based on the principle that the best auxiliary loss is the one that helps increase the performance of the main loss most through several steps of gradient descent. We then examined two commonly used auxiliary losses in inpainting and used ALA to adapt their parameters. Experimental results show that ALA induces more competitive inpainting results than fixed auxiliary losses. In particular, simply combining auxiliary loss with ALA, existing inpainting methods can achieve increased performances without explicitly incorporating delicate network design or structure knowledge prior.

Meta Corrupted Pixels Mining for Medical Image Segmentation

Deep neural networks have achieved satisfactory performance in piles of medical image analysis tasks. However the training of deep neural network requires a large amount of samples with high-quality annotations. In medical image segmentation, it is very laborious and expensive to acquire precise pixel-level annotations. Aiming at training deep segmentation models on datasets with probably corrupted annotations, we propose a novel Meta Corrupted Pixels Mining (MCPM) method based on a simple meta mask network. Our method is targeted at automatically estimate a weighting map to evaluate the importance of every pixel in the learning of segmentation network. The meta mask network which regards the loss value map of the predicted segmentation results as input, is capable of identifying out corrupted layers and allocating small weights to them. An alternative algorithm is adopted to train the segmentation network and the meta mask network, simultaneously. Extensive experimental results on LIDC-IDRI and LiTS datasets show that our method outperforms state-of-the-art approaches which are devised for coping with corrupted annotations.