1 Introduction↩︎

Convolutional Neural Networks (CNNs) have undergone significant architectural evolution. While early designs relied on stacking small kernels (typically \(3 \times 3\)), recent studies have demonstrated substantial gains from large convolution kernels. This shift was driven by ConvNeXt [1] with \(7 \times 7\) kernels and further pushed by RepLKNet [2] (\(31 \times 31\) kernels), SLaK [3] (\(51 \times 51\) kernels), and UniRepLKNet [4] (\(13 \times 13\) kernels).

However, large kernel convolutions face fundamental computational challenges that severely limit their practical deployment. Standard convolution implementations must gather scattered input neighborhoods for each output position, causing increasingly irregular memory access patterns as kernel size grows.

The core problem. Our comprehensive operator-level benchmarks (4.1) reveal that standard depthwise convolutions slow down by 71–78% when kernel size increases from \(5 \times 5\) to \(13 \times 13\), and by 92–94% for \(27 \times 27\) kernels. This degradation is consistent across all batch sizes (4, 16, 64, 128, 256) and feature map resolutions (\(\geq 28 \times 28\)), demonstrating that the performance penalty is inherent to the memory access pattern rather than implementation-specific.

1 contrasts the two computation paradigms at a high level. As 2 illustrates in more detail, the root cause is that gather-based depthwise convolution is memory-bound: each output collects \(k^2\) values from non-contiguous rows, exceeding L1/L2 cache and yielding only \(O(1)\) FLOPs per loaded element. WBMM eliminates this irregularity by reading each window as one contiguous block, shifting the operator from memory-bound to compute-bound (full arithmetic-intensity analysis in 3.4).

The limitation of existing acceleration. While the Large Kernel Acceleration (LKA) CUDA kernels from RepLKNet/UniRepLKNet mitigate this degradation on small feature maps during training, their fixed tiling no longer matches large feature maps, where they become even slower than the non-accelerated baseline. At \(224 \times 224\) resolution and batch=128, LKA with \(7 \times 7\) kernels is 80% slower than baseline, and \(13 \times 13\) kernels are 89% slower. This severely limits deployment of large kernels in tasks that process high-resolution feature maps, such as semantic segmentation, object detection, and super-resolution.

1.1 Our Approach: Traversing Parameters Instead of Data↩︎

We propose Windowed Batch Matrix Multiplication (WBMM), which changes the computation paradigm by traversing the parameter table instead of input data. As illustrated in 1:

• Depthwise convolution: For each output position, gather \(k^2\) scattered input values \(\rightarrow\) irregular memory access.

• WBMM: Partition input into contiguous windows, index parameter table to construct weight matrix \(\rightarrow\) regular memory access.

Concretely, WBMM partitions input into contiguous non-overlapping windows (e.g., \(7 \times 7\)) and constructs a weight matrix \(M \in \mathbb{R}^{C \times 49 \times 49}\) by indexing into a compact relative position bias table \(R \in \mathbb{R}^{C \times 169}\); the output is computed via batched matrix multiplication on contiguous memory blocks. This “partition-then-multiply” design is why WBMM’s efficiency does not degrade as feature map resolution or window size grows—the batched matrix multiplication fully leverages GPU parallelism and actually improves with more windows.

1.2 Contributions↩︎

Our main contributions are:

• Novel computation paradigm: We replace the gather-based memory access of large-kernel depthwise convolution with a “partition-then-index” scheme: WBMM indexes a compact relative-position bias table to form a per-channel weight matrix, then applies it via batched matrix multiplication on contiguous windows. Throughput improves with larger windows and stays consistent across resolutions—versus the 80–89% LKA slowdown at \(224 \times 224\) (batch=128)—making WBMM suitable for classification, detection, segmentation, and super-resolution.

• Batch-independent weight construction: Unlike window attention, which builds an input-dependent matrix per sample and window (\(O(B \cdot N_h N_w \cdot d^2)\)), WBMM constructs a single channel-dependent matrix (\(O(C \cdot d^2)\)), scaling better with batch size and resolution.

• Three lightweight components: (i) inter-block \(3 \times 3\) depthwise convolutions for cross-window communication, reaching comparable or higher accuracy at higher speed than intra-block alternatives; (ii) inference-time weight-matrix caching that removes redundant indexing on small feature maps; and (iii) hierarchical window reparameterization, a dual-scale fusion improving dense-prediction accuracy without slowing inference.

• Hardware-agnostic efficiency: Without specialized kernels, WBMM is consistently faster on GPU, CPU, and edge devices (16); across ImageNet-1K, COCO, and ADE20K it delivers a 1.31–1.88\(\times\) training speedup and, with hierarchical reparameterization, surpasses UniRepLKNet in both mIoU and AP.

2 Related Work↩︎

2.1 Large Convolution Kernels in CNNs↩︎

The resurgence of large convolution kernels began with ConvNeXt [1], which showed that \(7 \times 7\) depthwise convolutions could match Transformer performance. This trend continued with RepLKNet [2], which extended kernels to \(31 \times 31\). SLaK [3] reduced parameters through sparse rectangular decomposition, reaching \(51 \times 51\) kernels, though with suboptimal computational efficiency. PeLK [5] achieved linear parameter scaling through parameter sharing, enabling \(101 \times 101\) kernels but still faced quadratic computational complexity in \(k\).

UniRepLKNet [4] reported that \(13 \times 13\) is a strong kernel size across various modalities, based on extensive ablations that balance performance and computational cost. However, its structural reparameterization adds training overhead and relies on specialized GPU kernels that can become counterproductive during inference on large feature maps.

2.2 Convolution Implementations and Their Limitations↩︎

Im2col+GEMM. The classic im2col+GEMM approach [6], [7] converts convolution to matrix multiplication by extracting each \(k \times k\) receptive field as columns. For input \(X \in \mathbb{R}^{B \times C \times H \times W}\), this creates an expanded matrix of size \(\mathbb{R}^{BC \times HW \times k^2}\). The \(k^2\) memory expansion becomes prohibitive for large kernels, and gathering scattered neighbors creates irregular memory access that worsens with kernel size.

Standard Depthwise Convolution (DW-Std). Modern deep learning frameworks like PyTorch [8] use implicit GEMM via cuDNN [9] for depthwise convolutions. While implicit GEMM avoids explicit im2col materialization by computing indices on-the-fly, reducing memory overhead from \(O(k^2)\) to \(O(1)\), the fundamental limitation remains the gather-based computation pattern: for each output position, the implementation must compute indices to gather \(k^2\) scattered input values. As kernel size increases, these memory accesses become increasingly non-coalesced, causing the 71–78% degradation we observe from \(5 \times 5\) to \(13 \times 13\) kernels.

Large Kernel Acceleration (LKA). RepLKNet [2] and UniRepLKNet [4] provide specialized CUDA kernels (we call these LKA throughout, distinct from VAN’s Large Kernel Attention [10]) that tile the depthwise convolution so a working block stays in on-chip shared memory, improving data reuse. The gather-based access pattern itself is unchanged—each output still collects \((2k_h+1)(2k_w+1)\) scattered inputs—so this reuse holds only while the block stays on-chip, i.e.on small feature maps (\(\leq 14 \times 14\)). On larger maps its fixed compile-time tiling no longer matches the workload and, lacking cuDNN’s adaptive algorithm selection, the kernel runs slower than the standard baseline—likely a fixed-configuration effect: at \(224 \times 224\) and batch=128, LKA with \(13 \times 13\) kernels is 89% slower than DW-Std \(5 \times 5\) (12).

WBMM (Ours). Unlike all above methods that traverse input data to gather scattered neighborhoods, WBMM traverses only the parameter table; input data remains in contiguous window blocks, providing regular memory access regardless of effective kernel size or feature map resolution. The arithmetic-intensity analysis explaining why this enables acceleration with larger windows is deferred to 3.4.

2.3 Window-Based Processing in Vision Models↩︎

Window-based processing has become a key efficiency technique in vision models. Vision Transformer (ViT) [11] successfully adapts transformers to images, while Swin Transformer [12] introduces efficient window-based self-attention with non-overlapping windows. For cross-window information exchange, Swin employs shifted window partitioning in alternate layers, and has been applied to various tasks [13]–[15].

Critical distinction from window attention. Unlike Swin-style window attention that computes input-dependent matrices and applies a softmax per sample/window, WBMM constructs a single input-independent weight matrix \(M\) shared across the whole batch; \(M\) can be cached at inference (3.4.5) and tiled with a simple purely-linear kernel. We empirically compare against a parameter-matched non-overlapping window attention in 9.

2.4 Structural Reparameterization↩︎

Structural reparameterization enhances neural networks by adding branches during training that are merged at inference. RepVGG [16] pioneered this approach by adding \(1 \times 1\) convolutions and identity shortcuts during training and fusing them at inference. ACNet [17] incorporated asymmetric convolutions, while RepLKNet [2] transformed large \(31 \times 31\) kernels into functionally equivalent smaller \(5 \times 5\) kernels. RepMLPNet [18] introduced locality injection for fully-connected layers.

Our WBMM method introduces a structural reparameterization tailored to window-based batch matrix multiplication with embedded position bias. We apply this minimal reparameterization only to the smallest model variants (WBMM-P and WBMM-N) at the final stage (S4) for image classification tasks, avoiding the heavy multi-branch training overhead of UniRepLKNet.

2.5 Position Encoding in Vision Networks↩︎

Position encoding is crucial for modeling spatial relationships. Existing approaches include absolute position encoding [11], relative position bias [12], [19], and conditional position encoding [20]. Unlike Swin [12], which shares a single relative position bias across all samples, heads, and windows, our approach adopts channel-specific position-aware weights, enhancing representational capacity while maintaining computational efficiency.

3 Method↩︎

3.1 Notation and Definitions↩︎

We establish notation used throughout this paper: \(X \in \mathbb{R}^{B \times C \times H \times W}\) denotes the input feature map (batch \(B\), channels \(C\), spatial \(H \times W\)); \((h, w)\) are spatial position indices, not to be confused with the window size \(w_h \times w_w\) or the feature-map width \(W\); \(N = B \cdot N_h \cdot N_w\) is the total number of windows, where \(N_h = H / w_h\), \(N_w = W / w_w\); \(d = w_h \cdot w_w\) is positions per window; \(R \in \mathbb{R}^{C \times (2w_h-1)(2w_w-1)}\) is the relative position bias table; \(I \in \mathbb{Z}^{d \times d}\) is the relative position index matrix; and \(M \in \mathbb{R}^{C \times d \times d}\) is the weight matrix constructed from \(R\). To avoid clash between \(W\) (feature-map width) and depthwise kernel weights, we denote the kernel by \(K_{c,p,q}\) in 1 below; likewise, to avoid clash with the batch index \(b\), we write the per-channel bias as \(\beta_c\).

3.2 Problem Formulation: Why Large Kernels Are Slow↩︎

Standard depthwise convolution with kernel size \((2k_h+1) \times (2k_w+1)\) computes: \[Y_{b,c,h,w} = \sum_{i=-k_h}^{k_h} \sum_{j=-k_w}^{k_w} K_{c,\,i+k_h,\,j+k_w} \cdot X_{b,c,h+i,w+j} + \beta_c \label{eq:depthwise95conv}\tag{1}\] where \(K_c\) is the depthwise kernel for channel \(c\) and \(\beta_c\) is a per-channel bias.

The key computational challenge is that for each output position \((h,w)\), the implementation must gather \((2k_h+1)(2k_w+1)\) input values \(X_{b,c,h+i,w+j}\) from scattered memory locations. As kernel size \(k\) increases, these memory accesses become increasingly non-coalesced, causing severe performance degradation regardless of the underlying implementation strategy.

3.3 Theoretical Foundation: Convolution–Matrix Equivalence↩︎

We establish that depthwise convolutions on bounded feature maps can be exactly represented as matrix multiplication with position-dependent weights.

Design choice: from exact equivalence to a windowed approximation. 2 is exact but impractical to instantiate globally: at \(56 \times 56\), \(M_c\) would carry \(HW = 3{,}136\) rows and \(\sim\)​9.8M entries per channel. Following the locality strategy of Swin Transformer [12], we apply 2 inside non-overlapping \(w_h \times w_w\) windows, which shrinks each matrix to \(\mathbb{R}^{d \times d}\) with \(d = w_h w_w\) (e.g., \(49 \times 49\) for \(7 \times 7\) windows). WBMM is therefore not an exact realization of global convolution, but a windowed design inspired by the equivalence. The three subsequent components each address a trade-off introduced by this restriction: inter-block \(3 \times 3\) depthwise mixing restores cross-window connectivity (3.5); hierarchical reparameterization recovers multi-scale context (3.6); and the input-independence of \(M\) enables zero-overhead inference caching (3.4.5).

3.4 Windowed Batch Matrix Multiplication↩︎

Building on the equivalence theorem and the above design choice, we design WBMM to efficiently process arbitrary-sized feature maps through window partitioning and relative position indexing. 3 presents the complete algorithm.

3.4.1 Handling Arbitrary Input Sizes↩︎

A key advantage of WBMM is its ability to handle arbitrary input resolutions through zero-padding. For any input with spatial dimensions \(H \times W\), we compute the required padding as: \[\begin{align} \text{pad}_h &= (w_h - H \bmod w_h) \bmod w_h, \\ \text{pad}_w &= (w_w - W \bmod w_w) \bmod w_w. \end{align}\] This ensures the padded dimensions \(H' = H + \text{pad}_h\) and \(W' = W + \text{pad}_w\) are divisible by the window size. After processing, the padding is removed to restore the original output dimensions. 12 reports an empirical comparison against reflection and replication padding, showing that zero-padding is the most accurate among the three.

3.4.2 Relative Position Bias Table and Indexing↩︎

The core of WBMM is the relative position bias table \(R \in \mathbb{R}^{C \times (2w_h-1)(2w_w-1)}\), storing learnable weights for all unique relative offsets within a window.

We precompute the relative position index matrix \(I \in \mathbb{Z}^{d \times d}\): \[\begin{align} \delta_h &= h_i - h_j, \quad \delta_w = w_i - w_j, \\ I[i, j] &= (\delta_h + w_h - 1) \cdot (2w_w - 1) + (\delta_w + w_w - 1). \label{eq:index95computation} \end{align}\tag{4}\]

The weight matrix is constructed via index selection: \[M = R[:,\, I.\texttt{flatten()}].\texttt{view}(C, d, d). \label{eq:index95select}\tag{5}\]

3.4.3 Why WBMM Accelerates with Larger Windows↩︎

A critical finding is that WBMM speeds up when window size increases from \(7 \times 7\) to \(14 \times 14\)—the opposite behavior to depthwise convolutions. The reason lies in arithmetic intensity. For a window of size \(w \times w\): data reads are \(O(w^2)\) elements per window, computation is \(O(w^4)\) multiply-add operations, yielding arithmetic intensity of \(O(w^2)\) FLOPs per loaded element (the single weight matrix \(M\) is shared across all \(N\) windows in the batch, so its \(O(w^4)\) load cost is amortized over \(N\) windows and becomes negligible relative to the per-window input read of \(O(w^2)\); hence the per-window arithmetic intensity is dominated by input reads). Larger windows provide higher arithmetic intensity, allowing GPU compute units to be more fully utilized.

3.4.4 Why WBMM Scales with Feature Map Size↩︎

Unlike LKA-accelerated convolutions whose fixed tiling no longer matches large feature maps, WBMM’s efficiency improves with larger feature maps because: (1) more windows provide better GPU occupancy and parallelism; (2) the batched matrix multiplication amortizes fixed overhead across more windows; (3) memory access remains regular regardless of total feature map size.

3.4.5 Inference Caching Mechanism↩︎

During training, \(M\) must be reconstructed at each forward pass because gradients flow through \(R\). During inference, \(M\) is constant and can be computed once and cached, providing additional speedup on small feature maps where index operation overhead is relatively higher. We refer to the two modes as WBMM-NC (no-cache; \(M\) reconstructed each forward pass, used in training) and WBMM-C (cached; \(M\) computed once at inference); they produce identical outputs and differ only in whether \(M\) is recomputed.

3.5 Inter-Block Cross-Window Communication↩︎

Non-overlapping windows cannot directly exchange information. We address this by placing lightweight \(3 \times 3\) depthwise convolutions as separate blocks between WBMM blocks in an inter-block design (schematically, \(\text{Block}_n = \text{WBMM}\), \(\text{Block}_{n+1} = \text{DWConv}_{3 \times 3}\)); the actual stage-wise patterns used in our models (e.g., \([\text{W,D,W}]\), \([\text{W}{\leftrightarrow}\text{D}]{\times}9\), \([\text{W,W,W}]\)) are task-specific and listed in 19. How densely these \(3 \times 3\) blocks are interleaved with WBMM, and at which stages, is task-specific; we determine the patterns empirically in 4.2.2. Keeping WBMM and \(3 \times 3\) convolutions in separate blocks—rather than pairing a \(3 \times 3\) with WBMM inside every block—lets each operator run as a clean, highly optimized kernel, improving training and inference speed through more regular memory access and better operator fusion while retaining strong accuracy.

3.6 Hierarchical Window Reparameterization↩︎

For dense prediction tasks requiring multi-scale context, we introduce hierarchical window reparameterization combining global and local window processing (abbreviated “Hier” in tables; equivalently “G+L” in 3). The detailed algorithm is in 19.

Dual-Scale Bias Tables. We maintain two learnable tables: \(R_g \in \mathbb{R}^{C \times (2w_g-1)^2}\) for large windows and \(R_l \in \mathbb{R}^{C \times (2w_l-1)^2}\) for smaller sub-windows, where \(w_g = 2w_l\).

Weight Fusion. Each global window decomposes into a \(2 \times 2\) grid of local sub-windows, indexed by \(s \in \{0,1,2,3\}\). During training the two scales are processed independently, each with an identity shortcut; at inference they collapse into a single matrix that adds the local pattern \(M_l\) to the four diagonal blocks of the global matrix \(M_g\): \[M_{\text{fused}}[c,\, B_{s,s}] = M_g[c,\, B_{s,s}] + M_l[c], \quad s \in \{0,1,2,3\},\] while the off-diagonal cross-sub-window blocks are left to \(M_g\). The residual shortcuts are absorbed as identity terms in the fused matrix; the exact form and indexing are given in 19. The fusion is a one-time tensor addition performed when the model is loaded, after which \(M_{\text{fused}}\) is cached; hierarchical reparameterization therefore adds no runtime cost over single-scale WBMM at the same window size (\(w = 14\)).

3.7 Multi-Kernel Fusion for Minimal Feature Maps↩︎

For the S4 stage of Pico and Nano variants processing \(7 \times 7\) feature maps, we introduce parallel depthwise convolution paths fused with the WBMM main branch: \(Y = \text{WBMM}(X) + \text{BN}_1(\text{DW}_5(X)) + \text{BN}_2(\text{DW}_3(X))\), where \(\text{DW}_k\) denotes kernel size \(k \times k\). At inference, the parallel DW paths are fused into the WBMM weight matrix \(M\), achieving a 0.2 percentage point improvement in Top-1 over single-path WBMM at zero additional inference cost.

4 Experiments↩︎

We conduct experiments on ImageNet-1K [21] classification, ADE20K [22] semantic segmentation with UPerNet [23], and COCO [24] object detection with Cascade Mask R-CNN [25].

Implementation details. All experiments use FP32 precision; training runs on 8\(\times\) NVIDIA A800 GPUs (80GB) and inference benchmarks on a single A800 unless otherwise specified. The implementation follows the UniRepLKNet codebase with identical training protocols, hyperparameters, and data preprocessing; our only changes are replacing large kernels with WBMM and adding inter-block mixing. We use 300-epoch ImageNet training with AdamW [26], 160k iterations for ADE20K, and a \(3\times\) schedule for COCO. Ablation metrics are mean\(\pm\)std over three runs; full configurations are in 17.

Why compare with UniRepLKNet. UniRepLKNet [4] identified \(13 \times 13\) as the optimal kernel size through extensive cross-modal ablations, making it a rigorous state-of-the-art baseline for validating WBMM. Comparisons against further baselines—SLaK-style decomposition, non-overlapping window attention, and a ConvNeXt-T drop-in replacement—are in 8.

4.1 Comprehensive Operator-Level Benchmark↩︎

To validate WBMM’s efficiency, we conduct operator-level benchmarks on single-layer feature maps. This isolated setup makes the conclusions architecture-agnostic: the speedup characteristics hold regardless of whether WBMM is deployed in UniRepLKNet, ConvNeXt, or other CNN architectures.

All measurements use 256 channels, FP32 precision on a single NVIDIA A800 GPU (80GB), with 5 independent runs and IQR-based outlier removal. We present batch=128 results here; complete results across all batch sizes (4, 16, 64, 128, 256) are provided in 14.

4.1.1 Key Findings from Benchmark Analysis↩︎

4 presents the comprehensive benchmark results across four configurations (batch=128), revealing several important findings.

4.1.1.1 Finding 1: Standard Convolution Degrades 71–78%.

Standard depthwise convolutions (DW-Std, using implicit GEMM [9] via PyTorch [8]/cuDNN) consistently degrade by 71–78% when kernel size increases from \(5 \times 5\) to \(13 \times 13\), across all tested batch sizes and feature map resolutions (\(\geq 28 \times 28\)). For \(27 \times 27\) kernels, degradation reaches 92–94%. This degradation is inherent to the gather-based computation pattern and represents the baseline efficiency that any acceleration technique must improve upon.

4.1.1.2 Finding 2: LKA Becomes Counterproductive on Large Feature Maps.

The Large Kernel Acceleration CUDA kernels (LKA), employed by RepLKNet and UniRepLKNet for training workloads, use a fixed tiling tuned for small feature maps that no longer matches the workload at large resolution, so the kernels fall behind the standard baseline. At \(224 \times 224\) resolution and batch=128, LKA \(7 \times 7\) is 80% slower than baseline, and LKA \(13 \times 13\) is 89% slower. These kernels should be avoided for inference on large feature maps.

4.1.1.3 Finding 3: WBMM Accelerates with Larger Windows.

Unlike depthwise convolutions that degrade with kernel size, WBMM accelerates with larger window sizes. Compared to the DW-Std \(5 \times 5\) baseline, WBMM-C \(14 \times 14\) (window) provides a \(7.8\times\) larger per-layer receptive field (196 vs. positions per layer) while achieving faster speed on feature maps \(\geq 28 \times 28\) at batch \(\geq 16\) (and \(\geq 56 \times 56\) at all tested batch sizes; see 1).

4.1.1.4 Finding 4: WBMM’s Advantage Scales with Batch Size.

1 shows that WBMM-C \(14 \times 14\) (window) achieves increasing speedup over DW-Std \(5 \times 5\) as batch size grows: from \(0.45\times\) at batch=4 to \(\mathbf{3.72\times}\) at batch=256 on \(14 \times 14\) feature maps (note: window size and feature-map size coincide in this column), confirming WBMM’s batch-independent weight construction advantage.

3pt

Note: “–” indicates OOM at batch=256 for DW-Std.

4.2 Ablation Studies↩︎

We establish UniRepLKNet-T as our baseline model and systematically examine critical design components.

4.2.1 Feature Extraction and Window Interaction↩︎

2.5pt

¿tbl:tab:ablation95methods? presents our findings. For feature extraction, incorporating position-aware weights through relative position bias significantly improved performance (82.71\(\pm\)​0.04% Top-1, 45.50\(\pm\)​0.07% mIoU) compared to basic full connection (80.89\(\pm\)​0.02% Top-1, 44.02\(\pm\)​0.05% mIoU). The substantial performance gap (nearly 2 percentage points) highlights the importance of spatial relationship modeling.

For window interaction, we compare three ways of combining a \(3 \times 3\) depthwise convolution with WBMM: a parallel branch, a serial arrangement in which the \(3 \times 3\) convolution follows WBMM, and channel splitting in which half the channels take the \(3 \times 3\) path. The serial arrangement attains the best mIoU (48.01\(\pm\)​0.03%), ahead of parallel (47.90\(\pm\)​0.06%) and well above channel splitting (46.21\(\pm\)​0.04%). On Top-1, serial (83.21\(\pm\)​0.03%) and channel splitting (83.22\(\pm\)​0.01%) are essentially tied and both edge out parallel (83.11\(\pm\)​0.02%). Since channel splitting falls clearly behind on dense prediction while serial leads parallel on both metrics, the serial integration of local and global features provides more effective information flow than fusing them in parallel. Applying a \(3 \times 3\) after every WBMM, however, increases both training and inference time, so we do not adopt this per-block serial design; instead, the mixed block patterns in 2 distribute WBMM and \(3 \times 3\) across stages, which—tuned per task—match the per-block serial design’s classification accuracy (the \(0.01\)% Top-1 gap is negligible) while surpassing its mIoU (\(48.32\) vs.\(48.01\)), all at lower training and inference cost.

4.2.2 Optimal Architecture Design↩︎

2.5pt

Note: S1–S4 represent network stages at different resolutions. “\(\leftrightarrow\)” indicates alternating operations.

2 explores these mixed block patterns across stages. For image classification, performance peaked (83.21% Top-1) with configuration \([\text{D,D,D}]_{S1} \to [\text{W,D,W}]_{S2} \to [\text{W}{\leftrightarrow}\text{D}]_{S3} \to [\text{W,W,W}]_{S4}\). For semantic segmentation, optimal performance (48.32% mIoU) was achieved with \([\text{W,D,W}]_{S1} \to [\text{W,D,W}]_{S2} \to [\text{W}{\leftrightarrow}\text{D}]_{S3} \to [\text{W,W,W}]_{S4}\). This task-dependent divergence reflects different downstream use of S1 features: classification uses only the globally pooled deepest feature (favoring local \(3 \times 3\) extraction at S1), whereas FPN/UPerNet-style decoders fuse S1 directly with upsampled deeper features, benefiting from the broader spatial context WBMM provides at the highest-resolution stage.

4.2.3 Hierarchical Window Reparameterization↩︎

3pt

3 shows that applying hierarchical reparameterization (G+L) selectively to S1–S2 stages achieves 48.8\(\pm\)​0.06% mIoU, providing 0.8% improvement over the \(7 \times 7\) baseline and 0.9% over pure \(14 \times 14\) configuration.

4.3 ImageNet Classification↩︎

1.5pt

4 compares WBMM and UniRepLKNet on ImageNet-1K. WBMM matches or exceeds accuracy while offering significant efficiency gains: training speedup of 1.31–1.88\(\times\), memory reduction of 8–12%, and comparable accuracy (within 0.1%).

4.4 Semantic Segmentation on ADE20K↩︎

We evaluate WBMM on ADE20K semantic segmentation using UPerNet [23] with both Tiny and Small model variants. Semantic segmentation processes high-resolution feature maps with variable-sized inputs, which WBMM handles seamlessly via its zero-padding mechanism (3.4).

3pt

Note: “Hier” denotes hierarchical reparameterization combining \(14 \times 14\) windows with \(7 \times 7\)

sub-windows.

5 shows that WBMM with hierarchical reparameterization achieves higher mIoU than UniRepLKNet on both Tiny and Small variants. The pure \(7 \times 7\) variant matches FLOPs and parameters of UniRepLKNet within \(0.3\) mIoU, and adding hierarchical reparameterization closes this gap without slowing inference.

4.4.1 Inference Speed Analysis↩︎

We provide comprehensive throughput measurements for both Tiny and Small models across multiple input resolutions (512\(\times\)​512, 512\(\times\)​1024, 1024\(\times\)​1024) and batch sizes (2, 4, 8, 16) in 15, demonstrating three key findings: (1) WBMM with \(7 \times 7\) windows approaches the theoretical upper bound of Conv \(3 \times 3\) speed while enabling dense intra-window interactions; (2) Both WBMM configurations—\(7 \times 7\) windows and hierarchical reparameterization with \(14 \times 14\) windows—deliver consistent speedup over UniRepLKNet across all tested resolutions and batch sizes; (3) WBMM excels particularly at high resolution (e.g., \(1024 \times 1024\)) where LKA-accelerated UniRepLKNet suffers from severe performance degradation.

4.5 Object Detection on COCO↩︎

We further evaluate WBMM on COCO object detection with Cascade Mask R-CNN [25], using the same dense-prediction architecture as segmentation. WBMM with hierarchical reparameterization achieves higher AP than UniRepLKNet: 51.9%/45.1% (box/mask) for Tiny and 53.1%/46.1% for Small (6), while maintaining competitive inference speed.

3pt

Note: FLOPs computed at \(1280 \times 800\)

resolution.

4.6 Hardware Generalization↩︎

We evaluate WBMM on diverse hardware platforms; detailed benchmarks are in 16. On a single NVIDIA A800 GPU (FP32, batch=128), WBMM achieves 1.01–1.28\(\times\) over LKA-accelerated UniRepLKNet and 1.23–1.41\(\times\) over the non-accelerated version. On Intel Core i7-13700K CPU (FP32), where specialized large kernel acceleration is unavailable, WBMM achieves 1.03–1.48\(\times\) via highly optimized BLAS libraries. On NVIDIA Jetson Orin Nano (8GB), WBMM achieves 1.19–3.12\(\times\) on GPU (FP16) and 1.33–2.44\(\times\) on CPU (INT8), confirming compatibility with quantization.

5 Conclusion↩︎

We presented Windowed Batch Matrix Multiplication (WBMM), which accelerates large-kernel convolutions by traversing a compact parameter table instead of gathering scattered input neighborhoods, turning a memory-bound gather into a compute-bound batched matrix multiplication. WBMM thus speeds up with larger windows and with batch size (up to \(3.72\times\) at batch=256), and—unlike LKA-accelerated convolutions that slow down on large feature maps—maintains consistent efficiency while providing a \(7.8\times\) larger per-layer receptive field than DW-Std \(5 \times 5\). Combined with inter-block communication, inference caching, and hierarchical reparameterization, WBMM matches UniRepLKNet on ImageNet-1K and outperforms it on ADE20K and COCO with a 1.31–1.88\(\times\) training speedup, while generalizing across GPU, CPU, and edge devices without specialized kernels. We discuss limitations and future extensions in 20.

Acknowledgements↩︎

This work was supported in part by the Anhui Province Major Science and Technology Project under Grant JZ2024AKKZ0025. The numerical computations were carried out on the HPC Platform of Hefei University of Technology. The authors thank Professor Wengang Zhou of the University of Science and Technology of China for his valuable guidance on the writing of this manuscript, and the anonymous ICML 2026 reviewers for their constructive comments that substantially improved the paper.

Impact Statement↩︎

This paper presents work advancing the field of Machine Learning. The efficiency gains from WBMM could reduce energy consumption and carbon footprint in training and deploying vision models, potentially democratizing access to large-receptive-field models for researchers with limited computational resources.

6 Appendix Overview↩︎

This appendix provides comprehensive supplementary materials organized as follows:

• 7: Fair comparison protocol ensuring reproducibility.

• 8: Controlled cross-operator comparison (SLaK, DW variants) under identical architecture.

• 9: WBMM vs.non-overlapping window attention under matched parameters.

• 10: ConvNeXt-T drop-in replacement experiment.

• 11: Base-scale (~100M) ImageNet-1K results.

• 12: Padding-strategy ablation on ADE20K.

• 13: Analysis of learned weight matrices (\(M\)).

• 14: Complete operator-level benchmarks across all batch sizes.

• 15: Detailed segmentation inference speed analysis.

• 16: Hardware generalization benchmarks on GPU, CPU, and edge devices.

• 17: Complete training configurations.

• 18: Model architecture specifications.

• 19: Hierarchical reparameterization algorithm details.

• 20: Limitations and future work.

7 Fair Comparison Protocol↩︎

To ensure fair evaluation, we maintain identical framework configurations: same backbone structure, training protocols, hyperparameters, data preprocessing, and optimization strategies. The modifications include: (1) replacing large depthwise convolution kernels with WBMM operations, and (2) adapting the mixing strategy of how WBMM blocks and \(3 \times 3\) depthwise convolution blocks are arranged across different layers. The implementation follows the UniRepLKNet codebase.

Note on Structural Reparameterization. UniRepLKNet employs extensive structural reparameterization with multiple parallel branches during training, adding significant computational overhead. WBMM uses minimal reparameterization (only applied to S4 stage for Pico/Nano variants in classification tasks). Therefore, reported training speedups represent conservative estimates.

8 Controlled Cross-Operator Comparison↩︎

To isolate WBMM’s contribution from architecture-level design choices, we conduct a strictly controlled experiment in which all configurations share identical architecture, training protocol, hyperparameters, data preprocessing, and optimization; only the spatial mixing operator inside the spatial blocks differs. To enable a single comparison that is meaningful for both classification and dense prediction, all entries in 7 use the dense-prediction architecture variant (Tiny\(^*\) in 19: \([\text{W,D,W}]\) at S1, \([\text{W,D,W}]\) at S2, \([\text{W}{\leftrightarrow}\text{D}]{\times}9\) at S3, \([\text{W,W,W}]\) at S4), in which extra spatial-mixing blocks are placed at the high-resolution stage S1 to favor segmentation/detection. This is why the WBMM-T (\(w{=}7\)) Top-1 in 7 (83.0) is slightly lower than the value reported in 4 of the main text (83.2): the classification-tuned architecture used there (Tiny in 19: \([\text{D,D,D}]\) at S1) replaces the S1 WBMM blocks with pure \(3 \times 3\) depthwise convolutions, which is empirically better for ImageNet classification (see also 2). The same architecture is applied to every operator variant in 7, so the comparison among operators remains apples-to-apples. Training uses \(8\times\) A800 GPUs; inference speed is measured on a single A800 (FP32, batch=128). “Hier” denotes hierarchical window reparameterization (see 3.6) with \(14 \times 14\) global windows and \(7 \times 7\) local sub-windows fused at inference into a single matrix without slowing inference. “LKA” denotes the Large Kernel Acceleration CUDA kernels of RepLKNet/UniRepLKNet (used consistently throughout this paper); “reparam” denotes the UniRepLKNet structural reparameterization with kernel_sizes=[5,7,3,3,3] and dilates=[1,2,3,4,5]. “SLaK-style” uses parallel rectangular decomposition (\(5 \times k + k \times 5 + 5 \times 5\) with \(k \in \{51,49,47,13\}\)).

4pt

Note: All entries here use the dense-prediction architecture (Tiny\(^*\) in [tbl:tab:stage95patterns]: \([\text{W,D,W}]\) at S1) for an apples-to-apples comparison between operators on both classification and segmentation. The classification-tuned architecture (Tiny in [tbl:tab:stage95patterns]: \([\text{D,D,D}]\) at S1) is used in [tbl:tab:imagenet] of the main text and yields the higher Top-1 (83.2) for WBMM-T (\(w{=}7\)). The slight Time/Mem difference vs.[tbl:tab:imagenet] (6:20/15.16 vs.

6:18/15.04) reflects this architectural difference and is consistent within run-to-run variation.

Key observations. (1) WBMM \(w{=}7\) vs.\(13 \times 13\) DW. Comparable classification (83.0 vs.) and identical mIoU (48.3), but WBMM is faster in training (6:20 vs.:43 with LKA / 9:11 without) and inference (1833 vs./1305 img/s) with lower memory. (2) WBMM \(w{=}14\) Hier vs.\(13 \times 13\) DW + reparam. UniRepLKNet’s complex reparameterization recovers \(48.7\) mIoU at a severe training cost (\(10{:}21\)–\(12{:}46\) vs.\(6{:}31\)) and \(+1.36\) GB memory; WBMM Hier still leads \(+0.1\) mIoU at much lower overhead. (3) WBMM vs.SLaK variants. Despite an effective receptive field of \(51\), both SLaK variants underperform on dense prediction (\(47.2\)/\(47.4\) vs.\(48.3\)/\(48.8\) mIoU) and are \(2.4\)–\(3.4\times\) slower in both training and inference. (4) WBMM vs.window attention. See 9.

9 WBMM vs.Non-Overlapping Window Attention↩︎

To compare input-independent (WBMM) and input-dependent (attention) weight construction at the operator level, we replace only the WBMM blocks with non-overlapping \(7 \times 7\) window self-attention blocks. The MLP expansion ratio of the attention variant is set to \(2.5\) so that the parameter count and FLOPs match WBMM. Everything else—overall architecture, the inter-block \(3 \times 3\) depthwise mixing, training schedule and hyperparameters—is held identical.

The relevant rows from 7: Win7 Transformer (MLP \(r{=}2.5\)) reaches \(31.2\) M params, \(5.1\) G FLOPs, Top-1 \(83.3\), mIoU \(48.1\), with per-epoch time \(8{:}38\) and inference speed \(1245.1\) img/s. WBMM-T at matched scale reaches Top-1 \(83.0\), mIoU \(48.3\) (\(w{=}7\)) / \(48.8\) (\(w{=}14\), Hier), with training \(6{:}20\)/\(6{:}31\) and inference \(1833.1\)/\(1842.2\) img/s. Window attention obtains a small classification gain (\(+0.3\) Top-1 under the dense-prediction architecture used in 7; see the note under that table), but falls behind on dense prediction (\(-0.2\) to \(-0.7\) mIoU). WBMM is 1.32–1.36\(\times\) faster in training and 1.47–1.48\(\times\) faster in inference, with \(\sim\)​0.8–1.5 GB lower memory.

Where the gap comes from. Window attention computes query–key products for every sample and every window at every forward pass (\(O(B \cdot N_h N_w \cdot d^2)\), with \(N_h N_w\) windows per sample), and its softmax couples positions in a way that requires online normalization when tiled. WBMM’s weights are input-independent and constructed once at \(O(C \cdot d^2)\), shared across all batches and windows; cached at inference, the cost is amortized to zero. The empirical gap is consistent with this asymptotic difference, and grows with batch size. 8 summarizes these property differences.

6pt

When input-dependence may help. Content-adaptive routing tasks (e.g., VQA, referring expression) and tasks with extreme intra-image variation (e.g., satellite imagery) may benefit from input-dependent weights. A lightweight input-dependent gate over WBMM’s cached \(M\) would offer a hybrid design and is a natural extension.

10 Cross-Architecture Validation on ConvNeXt-T↩︎

To verify that WBMM’s efficiency is not specific to the UniRepLKNet backbone, we perform a drop-in replacement on ConvNeXt-T: all \(7 \times 7\) depthwise convolutions are replaced with WBMM blocks (\(w{=}7\)) interleaved with \(3 \times 3\) depthwise mixing blocks. No other change is made to the ConvNeXt training recipe.

6pt

As shown in 9, classification accuracy is preserved exactly (\(82.1\) Top-1), segmentation is slightly higher (\(46.2\) vs.\(46.0\) mIoU), and training is \(1.66\times\) faster (\(5{:}08\) vs.\(8{:}30\)). This confirms that WBMM’s efficiency gains are architecture-agnostic and require no task-specific tuning.

11 Base-Scale Results↩︎

We extend the comparison to the Base scale (\(\sim\)​100 M parameters) on ImageNet-1K (300 epochs, identical hyperparameters), with results summarized in 10:

6pt

WBMM-B achieves \(1.38\times\) training speedup (11:24 vs.:46 per epoch) with 9.1% lower memory (26.38 vs. GB) and \(+0.1\)% higher accuracy (\(83.9\) vs.\(83.8\)). Notably, UniRepLKNet-B does not improve over its Small variant (\(83.8\) vs.\(83.9\)) on ImageNet-1K alone, suggesting saturation at this data scale—the original paper uses ImageNet-22K pre-training for Base models. WBMM-B maintains Small-level accuracy on ImageNet-1K alone, showing slightly better scaling. ImageNet-22K pre-training experiments are planned future work.

12 Padding-Strategy Ablation↩︎

WBMM relies on padding to align spatial dimensions with the window size. On ADE20K with UPerNet (WBMM-T Hier, \(w{=}14\), variable input sizes), we compare three padding strategies. Zero-padding is empirically the best of the three strategies tested. Several factors contribute to this result: (1) four-stage downsampling (\(4\times\) to \(32\times\)) confines the padded region to \(<0.1\%\) of the S4 feature map; (2) UPerNet’s multi-scale fusion dilutes any boundary artifacts; (3) the inter-block \(3 \times 3\) depthwise mixing smooths boundary discontinuities; (4) reflection and replication padding introduce correlated values that interfere with the position-aware weights in \(R\), whereas zeros provide a clean neutral signal. Combined with hierarchical reparameterization, which reinforces local patterns, boundary artifacts are kept negligible. The quantitative results are summarized in 11.

6pt

The gap between zero-padding and the alternatives is small (\(0.2\)–\(0.4\) mIoU) but consistent across three independent runs, and we therefore use zero-padding throughout the paper. This comparison is performed with the hierarchical \(14 \times 14 + 7 \times 7\) configuration at the high-resolution stages, where the padding fraction is highest; for smaller windows the absolute differences shrink further because less padding is required.

13 Analysis of Learned Weight Matrices↩︎

To understand what WBMM learns, we visualize the effective weight matrix \(M\) at inference and observe three interpretable patterns; 5 summarizes the visualization.

(a) Locality. Across stages, \(M\) exhibits strong diagonal dominance with diagonal-to-off-diagonal ratios of \(41\)–\(58\times\). The mean off-diagonal magnitude \(|M(p,q)|\) decays by over \(90\%\) within Euclidean distance \(2\), confirming that WBMM automatically learns a locality prior despite having no explicit regularization toward locality.

(b) Channel specialization. Different channels exhibit clearly distinct preferences—horizontal, vertical, and diagonal—that resemble oriented edge detectors. This diversity arises naturally from the per-channel parameterization of \(R\), in contrast to attention which shares weights across channels within a head.

(c) Frequency selectivity. Channels split into low-pass and high-pass families (e.g., low-frequency energy ratio LF\(=83\%\) vs.LF\(=9\%\)). Shallow stages (S2) bias toward high-pass, while deeper stages (S3/S4) are more balanced. The \((2w-1)^2 = 169\) parameters per channel are clearly sufficient to encode this diversity.

These patterns emerge without explicit regularization, suggesting WBMM provides an efficient and expressive basis for learning spatial operators.

14 Complete Operator-Level Benchmark Across All Batch Sizes↩︎

We present comprehensive operator-level benchmark results across all tested batch sizes (4, 16, 64, 128, 256). All measurements use 256 channels, FP32 precision on a single NVIDIA A800 GPU (80GB), with 5 independent runs and IQR-based outlier removal. These results provide the complete data underlying 4 and 1 in the main text.

Why Single-Layer Operator Benchmarks. We conduct operator-level benchmarks on single-layer feature maps for several important reasons:

1. Isolation of computational properties: Single-layer testing isolates the fundamental computational characteristics of WBMM vs depthwise convolutions, free from confounding factors like network depth, skip connections, or layer interactions.

2. Architecture-agnostic conclusions: Results from isolated operator benchmarks generalize across different CNN architectures (ConvNeXt, ResNet, etc.), not just UniRepLKNet.

3. Controlled experimental conditions: Fixed channel count (256) and systematic variation of batch size and feature map resolution enable precise characterization of scaling behavior.

4. Reproducibility: Isolated benchmarks are easier to reproduce and verify independently.

1.8pt

Methods: DW-Std = Standard depthwise convolution (implicit GEMM via PyTorch/cuDNN); DW-LKA = Large Kernel Acceleration kernels of RepLKNet/UniRepLKNet; WBMM-C = WBMM with cached weight matrix (inference mode, see [sec:sec:caching]); WBMM-NC = WBMM with on-the-fly weight construction (training mode). WBMM-C and WBMM-NC differ only in whether \(M\) is recomputed each forward pass; they produce identical outputs. Highlighted rows show baseline (\(5\times5\)) and best WBMM configuration (\(14\times14\)

cached).

14.1 Benchmark Analysis Summary↩︎

The comprehensive benchmark results in 12 confirm the following architecture-agnostic properties:

Key Finding 1: Standard Convolution Degradation is Inherent. The 71–78% throughput degradation of DW-Std from \(5 \times 5\) to \(13 \times 13\) (and 92–94% for \(27 \times 27\)) holds across all tested batch sizes and feature map resolutions \(\geq 28 \times 28\), confirming that this penalty is inherent to the gather-based computation pattern (cf.4.1).

Key Finding 2: LKA Has a Crossover Point. LKA shows effectiveness only on small feature maps. At batch\(\geq\)​64, LKA provides speedup only for feature maps \(\leq 14 \times 14\). On large feature maps (\(\geq 112 \times 112\)), LKA is 63–99% slower than baseline.

Key Finding 3: WBMM Scaling. WBMM-C \(14\times14\) achieves \(+56\%\)–\(+108\%\) speedup over baseline on feature maps \(\geq 28\times28\) at batch\(\geq\)​16, peaking at \(+272\%\) on \(14\times14\) maps at batch=256.

Key Finding 4: Caching Benefit. WBMM-C provides up to \(\sim\)​3\(\times\) speedup over WBMM-NC on the smallest feature maps (e.g., a single \(14 \times 14\) window at small batch), where weight construction overhead is proportionally larger.

15 Segmentation Inference Speed Analysis↩︎

We provide comprehensive throughput measurements for both Tiny and Small model variants on ADE20K semantic segmentation. All experiments are conducted on a single NVIDIA A800 GPU (80GB) with FP32 precision. These results support the claims made in 4.4.

15.1 Tiny Model Results↩︎

13 reports the throughput measurements for the Tiny models.

4pt

15.1.0.1 Key Observations for Tiny Model.

(1) WBMM approaches theoretical upper bound: At \(512\times512\) with batch=16, WBMM-T with \(7 \times 7\) windows achieves 104.6 FPS, only 2.0% slower than the Conv \(3 \times 3\) upper bound (106.7 FPS) despite providing a \(5.4\times\) larger per-layer receptive field (49 vs. positions per layer). This demonstrates that WBMM’s regular memory access pattern nearly eliminates the performance penalty traditionally associated with large receptive fields.

(2) WBMM excels at high resolution: At \(1024\times1024\) resolution—critical for dense prediction quality—hierarchical WBMM achieves 30.4 FPS (batch=16), outperforming even Conv \(5 \times 5\) (29.8 FPS) while providing a \(7.8\times\) larger per-layer receptive field (196 vs. positions per layer). This validates WBMM’s core design principle: efficiency improves rather than degrades with larger feature maps.

(3) Consistent speedup over UniRepLKNet: WBMM delivers 1.03–1.08\(\times\) speedup over UniRepLKNet (without LKA) across all tested configurations. Against LKA-accelerated UniRepLKNet, WBMM’s advantage is even more pronounced at high resolution with small batch sizes where LKA’s custom kernels are least efficient.

(4) LKA severely degrades on high-resolution: UniRepLKNet with LKA shows dramatic performance degradation on large feature maps, confirming our operator-level finding that LKA becomes counterproductive for dense prediction tasks.

15.2 Small Model Results↩︎

14 reports the throughput measurements for the Small models.

4pt

15.2.0.1 Key Observations for Small Model.

WBMM-S with \(7\times7\) windows approaches theoretical maximum efficiency: 93.9 FPS vs 94.5 FPS for Conv \(3\times3\) at \(512\times512\) (only 0.6% gap), and matches Conv \(3\times3\) exactly at \(1024\times1024\) (26.9 vs 26.8 FPS). The hierarchical variant demonstrates strong scalability, achieving 26.3 FPS at \(1024\times1024\) (B16) which exceeds Conv \(5\times5\) by 2.7% despite a \(7.8\times\) larger per-layer receptive field (196 vs. positions per layer). Compared to UniRepLKNet-S, WBMM-S achieves consistent speedup across all configurations.

16 Hardware Generalization: CPU and Edge Device Benchmarks↩︎

To demonstrate WBMM’s platform-independent efficiency, we conduct comprehensive experiments on diverse hardware configurations beyond the main GPU experiments. These results support the brief hardware generalization summary in 4.6.

16.1 Experimental Setup↩︎

Hardware Configurations:

• High-end GPU: NVIDIA A800 (80GB), FP32 precision.

• Desktop CPU: Intel Core i7-13700K, FP32 precision.

• Edge Device: NVIDIA Jetson Orin Nano (8GB), GPU with FP16 precision and CPU with INT8 quantization.

16.2 Desktop GPU and CPU Results↩︎

15 reports inference speed on desktop GPU and CPU.

0.8pt

Key findings: (1) On GPU, WBMM achieves 1.01–1.28\(\times\) speedup over accelerated UniRepLKNet (UA) and 1.23–1.41\(\times\) over non-accelerated (UN). (2) On CPU where specialized large kernel acceleration is unavailable, WBMM achieves 1.03–1.48\(\times\)

speedup by leveraging standard optimized matrix operations.

16.3 Edge Device Results↩︎

16 reports inference speed on the edge device.

0.4pt

Key findings: (1) On edge GPU with FP16, WBMM achieves 1.19–3.12\(\times\) speedup over accelerated UniRepLKNet, with larger gains at higher resolutions. (2) On edge CPU with INT8 quantization, WBMM achieves 1.33–2.44\(\times\)

speedup, demonstrating excellent compatibility with quantization.

16.4 Analysis and Discussion↩︎

1. Platform-Independent Efficiency: WBMM consistently outperforms UniRepLKNet across all tested hardware configurations because it relies on standard batched matrix multiplication operations that are universally optimized.

2. Scaling with Input Resolution: On GPU, WBMM’s advantage increases with higher input resolution (e.g., 1.04\(\times\) at 224\(\times\)​224 to 1.28\(\times\) at 1024\(\times\)​1024 for WBMM-P vs UA).

3. CPU Efficiency: On CPU where LKA acceleration is not available, WBMM achieves substantial speedups (up to 1.48\(\times\)) by leveraging highly optimized BLAS libraries.

4. Edge Device Compatibility: WBMM demonstrates excellent efficiency with both FP16 and INT8 precision, confirming compatibility with quantization techniques essential for edge deployment.

17 Training Configurations↩︎

This section provides complete training configurations for all experiments reported in the main text.

17.1 ImageNet-1K Classification↩︎

17 lists the ImageNet-1K training configuration.

17.2 COCO Object Detection↩︎

Following the \(3\times\) training schedule:

• Framework: MMDetection [27] with Cascade Mask R-CNN.

• Training epochs: 36.

• Batch size: 16.

• Optimizer: AdamW with learning rate 1e-4.

• Image scales: Short side 480–800, long side \(\leq\) 1333.

• Pre-training: ImageNet-1K weights.

17.3 ADE20K Semantic Segmentation↩︎

• Framework: OpenMMLab semantic segmentation toolbox (MMSegmentation) with UPerNet.

• Training iterations: 160k.

• Batch size: 16.

• Optimizer: AdamW with learning rate 1e-4.

• LR schedule: Polynomial decay (power=1.0).

• Crop size: \(512 \times 512\).

• Pre-training: ImageNet-1K weights.

18 Model Architecture Specifications↩︎

18 summarizes the stage depths, channel dimensions, and computational costs of all model variants.

3pt

\(^*\)For dense prediction variants, S1 and S2 use hierarchical \(14 \times 14 + 7 \times 7\) windows, S3 uses a single \(14 \times 14\) window, and S4 uses a single \(7 \times 7\)

window (its feature map being the smallest in the network).

19 Hierarchical Reparameterization Algorithm Details↩︎

This section provides the detailed algorithms for hierarchical reparameterization, referenced from 3.6.

19.1 Training Mode: Independent Multi-Scale Learning↩︎

19.2 Inference Mode: Zero-Cost Matrix Fusion↩︎

Zero overhead property. At deployment, the merged \(M_{\text{fused}}\) is a single \(C \times 196 \times 196\) tensor, so spatial mixing is carried out by one batched matrix multiplication over \(14 \times 14\) windows—no more arithmetic per block than single-scale WBMM at \(w = 14\), while expressing both global and local patterns. The two tables \(R_g\) and \(R_l\) are learned independently and both remain as parameters of the deployed model; their fusion is a one-time tensor addition at load and is invisible to the forward pass.

Training–inference equivalence. [alg:training_hierarchical,alg:inference_fused] are functionally identical: the training mode computes \(Y_w = X_w M_g + X_w + X_{sub} M_l + X_{sub}\) (with \(X_{sub} M_l + X_{sub}\) reshaped back to global windows), while the inference mode computes \(Y_w = X_w M_{\text{fused}}\). Because \(X_{sub}\) is just \(X_w\) repartitioned into \(2 \times 2\) sub-windows, the contribution \(X_{sub} M_l\) is structurally equivalent to multiplying \(X_w\) by a block matrix that contains \(M_l\) on its four diagonal blocks (and zeros off-diagonal), and the two identity shortcuts collapse into the \(196 \times 196\) identity. Therefore \(M_{\text{fused}} = M_g + \text{diag\_blocks}(M_l, M_l, M_l, M_l) + 2\,I_{196}\) reproduces the training-mode output exactly.

Per the segmentation/detection architectural specification (19), this hierarchical fusion is used at the high-resolution stages S1 and S2 where multi-scale context contributes most; deeper stages run a single window per block, with S3 operating at \(14 \times 14\) for spatial coverage and S4 using a single \(7 \times 7\) window at the network’s smallest feature map.

20 Limitations and Future Work↩︎

We note WBMM’s current limitations along several dimensions.

Memory footprint at very large windows. The constructed weight matrix \(M \in \mathbb{R}^{C \times d \times d}\) grows as \(O(C \cdot w^4)\) with window size \(w\). At \(w = 14\) (\(d = 196\)), \(M\) takes on the order of tens of MB across the network’s stages (specifically, \(\sim\) \(37.5\) MB per 256-channel layer at FP32, lower for earlier stages with fewer channels)—negligible compared to activations. For windows beyond \(28 \times 28\), however, materializing the full \(M\) becomes memory-prohibitive. Crucially, the bias table \(R\) remains compact (\(O(C \cdot (2w{-}1)^2)\)), so this is an engineering rather than a representational limitation: a Flash-Attention-style tiling kernel that stores \(R\) in SRAM and computes \(M\) blocks on-the-fly via the index map (4 5 ) would avoid full materialization entirely. Because WBMM is purely linear (no softmax), such tiling is strictly simpler than Flash Attention’s online softmax. Triton-based tiled kernels are a planned engineering effort.

Window size vs.task. Pure large windows are not uniformly optimal: 3 shows that pure \(14 \times 14\) underperforms pure \(7 \times 7\) on ADE20K. The hierarchical reparameterization of 3.6 resolves this by combining global and local patterns without slowing inference, but task-specific window selection still requires design effort. Automated window-size search is a natural next step.

Input independence. WBMM constructs a single weight matrix shared across all samples and windows. This is a feature for efficiency on dense prediction (9: WBMM matches or beats parameter-matched window attention on segmentation while being noticeably faster), but content-adaptive routing tasks may benefit from input-dependent weights. A lightweight input-dependent gate \(\sigma(\cdot)\) over the cached \(M\) is a promising hybrid we leave to future work.

Generalization to 1D and 3D data. This paper focuses on 2D vision. WBMM’s core abstraction—a compact relative-position bias table indexed into a per-channel weight matrix, applied via batched MatMul to contiguous windows—naturally extends to 1D temporal sequences and 3D spatio-temporal video. The optimal window shape (e.g., \(T \times H \times W\) or factorized variants) and cross-window strategies in these domains require dedicated study and constitute a promising future direction.

References↩︎