Rep3D: Re-parameterize Large 3D Kernels with Low-Rank Receptive Modeling for Medical Imaging

TL;DR

Rep3D addresses the challenge of efficiently training large-kernel 3D convolutions for medical image segmentation by embedding a learnable spatial prior that mirrors ERF-driven diffusion. It provides a theoretical link between parallel large+small kernel branches and element-wise learning rates, then implements a low-rank receptive bias modeling (LRBM) via a lightweight generator to modulate kernel updates during training. The resulting plain-encoder architecture uses a 21^3 depthwise kernel with a two-stage modulation network to achieve local-to-global convergence, yielding state-of-the-art Dice scores across five 3D segmentation benchmarks. This approach improves optimization stability and performance while maintaining inference efficiency, offering a scalable and interpretable alternative to transformer-based methods in 3D medical imaging.

Abstract

In contrast to vision transformers, which model long-range dependencies through global self-attention, large kernel convolutions provide a more efficient and scalable alternative, particularly in high-resolution 3D volumetric settings. However, naively increasing kernel size often leads to optimization instability and degradation in performance. Motivated by the spatial bias observed in effective receptive fields (ERFs), we hypothesize that different kernel elements converge at variable rates during training. To support this, we derive a theoretical connection between element-wise gradients and first-order optimization, showing that structurally re-parameterized convolution blocks inherently induce spatially varying learning rates. Building on this insight, we introduce Rep3D, a 3D convolutional framework that incorporates a learnable spatial prior into large kernel training. A lightweight two-stage modulation network generates a receptive-biased scaling mask, adaptively re-weighting kernel updates and enabling local-to-global convergence behavior. Rep3D adopts a plain encoder design with large depthwise convolutions, avoiding the architectural complexity of multi-branch compositions. We evaluate Rep3D on five challenging 3D segmentation benchmarks and demonstrate consistent improvements over state-of-the-art baselines, including transformer-based and fixed-prior re-parameterization methods. By unifying spatial inductive bias with optimization-aware learning, Rep3D offers an interpretable, and scalable solution for 3D medical image analysis. The source code is publicly available at https://github.com/leeh43/Rep3D.

Figures (3)

• Figure 1: (a) Traditional structural re-parameterization methods (e.g., CSLA blocks) re-parameterize small and large kernel convolutions to improve representational capacity, but apply linear optimization with same learning rate across the kernels, demonstrating a faster convergence in local then global. (b) In contrast, Rep3D introduces a learnable spatial bias via a generator network $f_\theta$, which modulates each element in the large kernel using a prior based on distance decay. This adaptive modulation enables local-to-global update dynamics aligned with ERF behavior, enhancing both training stability and model performance for 3D volumetric tasks.
• Figure 2: In contrast to (a) structural or (b) gradient-based re-parameterization, Rep3D introduces a novel re-parameterization strategy that injects a learnable spatial bias into large kernel convolutions for optimization. During training, a lightweight generator network produces a modulation mask conditioned on a distance-based prior, which adaptively scales gradient updates across the kernel. This enables spatially-aware learning dynamics that reflect local-to-global variations in the effective receptive field (ERF).
• Figure 3: As kernel size increases, depthwise convolutions in 3D UX-Net exhibit increasingly diffuse ERFs, gradually expanding the gradient dynamics from local to broader spatial regions. Incorporating LRBM further enhances weighting toward global areas by modulating the spatial contribution of distant elements. In contrast, Rep3D produces a well-distributed ERF that preserves strong central activation while extending contextual influence across the full kernel.