WHTMix: Efficient Stereo Depth Estimation via Walsh-Hadamard Token Mixing

Abstract

Stereo depth estimation for driving, robotics and augmented reality must run at high resolution under tight latency budgets, yet in transformer-based matchers the global self-attention that aggregates scene context grows quadratically with the number of pixels and comes to dominate runtime. We show that the joint self-attention stage of a stereo transformer, whose role is to spread context across both views, can be replaced by a data-independent Walsh-Hadamard token mixer that mixes tokens globally in the transform domain at log-linear cost, while the data-dependent cross-attention that performs left-right correspondence is retained. On synthetic driving data the mixer matches the attention baseline in end-point error while reducing model compute by a factor of 2.46 and single-image inference latency by a factor of 2.65. A complexity analysis shows the benefit is governed by the ratio of sequence length to channel width, which explains why high-resolution stereo matching is a particularly favorable setting and why classification transformers are not; we confirm this token-to-channel scaling on non-stereo long-sequence benchmarks. Furthermore, we introduce a hybrid log-disparity loss function designed to up-weight small-disparity pixels corresponding to long-range objects. This approach reduces the error on distant objects without incurring any additional computational overhead.

I Introduction

Stereo depth estimation recovers a dense disparity map from a rectified image pair and underpins robotics, autonomous driving and augmented reality. Transformer-based matchers such as STTR [12] attain competitive accuracy by replacing a fixed disparity search with attention, but their global self-attention scales as $O(N^{2})$ in the number of feature tokens $N$, making them slow at the high resolutions these applications demand.

Attention plays two distinct roles inside a stereo transformer layer. A joint self-attention first spreads global context within and across the two views; a subsequent cross-attention performs the actual left-right correspondence. Only the second is inherently content dependent: matching a pixel to its epipolar counterpart requires comparing image content. Context aggregation, by contrast, can be performed by a fixed, data-independent operator. This motivates replacing the joint self-attention with a cheap spectral token mixer while keeping cross-attention for matching.

We use the orthonormal Walsh-Hadamard transform (WHT), a multiply-free, real-valued relative of the Fourier transform computable in $O(N\log N)$ by a butterfly. Our mixer, WHTMix, applies learnable per-frequency gains in the WHT domain over both channels and tokens, giving a global receptive field at log-linear cost. Our contributions are:

• A drop-in Walsh-Hadamard token mixer that replaces the joint self-attention of a stereo transformer, cutting model compute $2.46\times$ and latency $2.65\times$ at matched accuracy on synthetic data.

• A complexity analysis giving a closed-form speedup of $2+N/C$, with a crossover at $N=2C$ and an end-to-end ceiling near $3\times$, that predicts exactly where the mixer helps, matches measured latency, and is confirmed on non-stereo Long-Range Arena tasks.

• A hybrid log-disparity loss that roughly halves error on far objects at no latency cost, and an honest study of the synthetic-to-real accuracy trade-off.

II Related Work

Stereo matching. Cost-volume networks with 3D convolutions [3] and iterative refinement methods such as RAFT-Stereo [13], IGEV-Stereo [21] dominate standard benchmarks, while foundation models push zero-shot generalization [20]. STTR [12] recast matching as a sequence-to-sequence problem using alternating self- and cross-attention along epipolar lines. Closest to us, HART [4] builds a stereo transformer on Hadamard-product attention; we instead replace the joint self-attention with the Walsh-Hadamard transform [2, 1] as a fixed token mixer, keeping dot-product cross-attention for matching. We adopt an STTR-style backbone as our testbed because it is the paradigm in which global self-attention is the bottleneck.

Efficient and spectral token mixers. Dot-product self-attention [19] is quadratic in sequence length; efficient variants use windowed attention [14], or selective state spaces [9], while FlashAttention [5] lowers attention’s memory but not its quadratic compute. Spectral mixers are a distinct line: FNet [11] replaces attention with an unparameterized 2D Fourier transform; GFNet [17] learns global filters in the Fourier domain; AFNO [10] adds a channel mixer in frequency space; and Walsh-Hadamard layers have cut convolution and projection cost in deep networks [16, 6]. Our mixer is closest to GFNet but uses the real, multiply-free Walsh-Hadamard transform and, unlike these classification-oriented works, is applied selectively to the context-mixing stage of a dense correspondence model while retaining data-dependent cross-attention.

III Method

III-A Backbone

We follow an STTR-style pipeline: a shared CNN feature extractor downsamples each image by $8\times$; a fixed 2D sinusoidal positional encoding is added; a stack of $L$ transformer layers alternates joint self-attention (over the concatenated left-right tokens, length $N=2\,(H/8)(W/8)$) with two cross-attentions; a normalized correlation volume and a soft-argmin regression head produce the disparity. Channel width is $C$. We replace only the joint self-attention.

III-B Walsh-Hadamard Token Mixer

Let $\mathbf{X}\in\mathbb{R}^{N\times C}$ be the joint token sequence and let $\mathcal{H}$ denote the orthonormal WHT (the Sylvester-Hadamard matrix normalized by $\sqrt{n}$, defined for $n$ a power of two and applied by an $O(n\log n)$ butterfly). WHTMix computes

$$\mathbf{Z} =\mathbf{X}\mathbf{W}_{\mathrm{in}},$$ (1)

$$\mathbf{Z} \leftarrow\mathcal{H}_{C}^{-1}\!\big(e^{\mathbf{g}_{c}}\odot\mathcal{H}_{C}(\mathbf{Z})\big),$$ (2)

$$\mathbf{Z} \leftarrow\mathcal{H}_{N}^{-1}\!\big(e^{\mathbf{g}_{t}}\odot\mathcal{H}_{N}(\mathbf{Z})\big),$$ (3)

$$\mathbf{Y} =\mathbf{Z}\mathbf{W}_{\mathrm{out}},$$ (4)

where $\mathcal{H}_{C}$ and $\mathcal{H}_{N}$ act over channels and tokens respectively, and $\mathbf{g}_{c},\mathbf{g}_{t}$ are learnable log-domain per-frequency gains (initialized to zero, i.e. unit gain). The token axis is zero-padded to the next power of two before $\mathcal{H}_{N}$. Both transforms are orthonormal, hence magnitude preserving, so layers stack stably; as in FNet/GFNet the block is linear and the layer feed-forward network supplies the nonlinearity. WHTMix mirrors the attention call signature and is a literal drop-in for the self-attention module. Fig. 1 contrasts the two token mixers.

Figure 1: Architectural difference between self-attention and WHTMix

III-C Hybrid Log-Disparity Loss

Standard masked $\ell_{1}$ is dominated by near objects with large disparities, so far objects (small disparities) are under-supervised. We weight each valid pixel $p$ by

$$w_{p}=\mathrm{clip}\!\left(\tau/d_{p},\;1,\;w_{\max}\right),$$ (5)

with ground-truth disparity $d_{p}$, so pixels below $\tau$ px are up-weighted up to $w_{\max}$, and minimize $\frac{1}{\sum_{p}w_{p}}\sum_{p}w_{p}\,|d_{p}-\hat{d}_{p}|$. We use $\tau=16$, $w_{\max}=32$.

III-D Complexity Analysis

Counting multiply-accumulates (MACs; FLOPs $=2\times$MACs), a self-attention block costs $4NC^{2}$ (projections) $+\,2N^{2}C$ (the quadratic core), whereas WHTMix costs $\approx 2NC^{2}$ (two projections; the butterfly is add-only, contributing no MACs). The token-mixer speedup is therefore

$$S_{\mathrm{mod}}=\frac{4NC^{2}+2N^{2}C}{2NC^{2}}=2+\frac{N}{C},$$ (6)

with a crossover at $N=2C$: below it attention is projection-bound and replacing it saves little; above it the quadratic core dominates and the gain grows linearly in $N/C$. The benefit is thus governed by the token-to-channel ratio. For our stereo model at $512\times 1024$, $N=16{,}384$ and $C=64$ give $N/C=256$; a classification ViT [7] ($N\!\approx\!196$, $C\!=\!768$) gives $N/C\!\approx\!0.26$ and no benefit.

Because the retained cross-attentions (each on $N/2$ tokens) and the feed-forward network are shared by both models, the end-to-end layer speedup is capped:

$$S_{\mathrm{e2e}}=\frac{16NC^{2}+3N^{2}C}{14NC^{2}+N^{2}C}=\frac{16C+3N}{14C+N}\;\xrightarrow[N\to\infty]{}\;3.$$ (7)

Keeping genuine matching is the price of this ceiling, and it is worth paying since a fixed filter cannot perform correspondence.

Figure 2: Qualitative comparison on CARLA (input, ground truth, and the three models). WHTMix produces disparity maps visually indistinguishable from the attention baseline; the log-disparity variant sharpens distant structure.

IV Experiments

IV-A Setup

We pre-train on a synthetic CARLA [8] stereo set (about 2,214 triplets, 90/10 split) with AdamW ($\mathrm{lr}=10^{-4}$, weight decay $0.01$), batch size 32, bf16 autocast, at $512\times 1024$, and fine-tune on KITTI 2015 [15] (200 triplets) at $384\times 1248$ with a frozen-backbone warm-up. All models use $C=64$, $8$ heads, $L=3$ layers, maximum disparity $256$. We compare three variants that differ by one factor at a time: Baseline (softmax self-attention), WHTMix (L1) (our mixer, $\ell_{1}$ loss) and WHTMix (log) (our mixer, log-disparity loss). Cross-attention is identical in all three. We report end-point error (EPE) over valid pixels and a per-disparity-bin breakdown, parameters, GFLOPs measured with fvcore (with a custom handler that counts fused attention), and batch-1 latency (warm-up plus CUDA-synchronized timing) on an RTX 4090.

IV-B Accuracy and Efficiency

Table 1 reports the main results. On CARLA, WHTMix matches the attention baseline in EPE ($2.27$ vs. $2.29$, within run-to-run noise) while using $2.46\times$ fewer FLOPs and running $2.65\times$ faster. On real KITTI data, where only 200 frames are available for adaptation, the fixed mixer is less accurate than adaptive attention, a trade-off we report honestly. Efficiency gains are essentially identical in both domains.

Table 1: Accuracy and efficiency on CARLA and KITTI.

Model	Params (M)	GFLOPs	Lat. (ms)	FPS	EPE (px)
CARLA ($512\times 1024$)
Baseline (attn)	0.783	348.9	95.7	10.4	2.294
WHTMix (L1)	0.759	141.9	36.0	27.8	2.274
KITTI 2015 ($384\times 1248$)
Baseline (attn)	0.783	294.6	81.8	12.2	2.204
WHTMix (L1)	0.759	121.6	31.8	31.5	2.670

Table 2 breaks EPE down by disparity. The log-disparity loss cuts error on the far, small-disparity bins by roughly half (2–4 px: $3.29\!\to\!1.60$; 4–8 px: $3.36\!\to\!1.94$) at the cost of a modest overall regression, isolating the effect of the loss from that of the mixer. Qualitatively (Fig. 2), WHTMix disparity maps are visually indistinguishable from the attention baseline.

Table 2: Per-disparity-bin EPE (px) on CARLA. Smaller the bins, farther the objects.

Model	0–2	2–4	4–8	8–16	16–32	32–64	64–128	128$+$
Baseline	0.64	2.96	3.04	2.89	2.79	3.62	3.12	1.70
WHTMix (L1)	0.57	3.29	3.36	2.94	2.91	3.62	3.04	1.68
WHTMix (log)	0.56	1.60	1.94	2.49	3.03	4.01	3.44	1.96

IV-C Latency Scaling

Table 3 and Fig. 3 sweep input resolution. End-to-end, the WHTMix speedup grows with resolution, from $1.1\times$ at $256\times 512$ to $2.85\times$ at $896\times 1792$, approaching the ceiling of Eq. (7) (predicted $2.90\times$ vs. measured $2.65\times$ at $512\times 1024$; $2.97\times$ vs. $2.85\times$ at $896\times 1792$): both curves are near-quadratic (fitted slopes $1.84$ and $1.50$) because the shared cross-attention dominates. Isolating the token mixer alone (Fig. 4) exposes the true asymptotics. Self-attention scales near-quadratically (log-log slope $1.70$) while WHTMix stays nearly flat (slope $0.25$); at the deployed $N{=}16{,}384$ the standalone mixer is $33\times$ cheaper, reaching $162\times$ at $N{=}65$k, though its fixed overhead leaves it no faster than attention at the smallest size ($N\!=\!1024$). The end-to-end and isolated views are complementary: the first is the deployable, Amdahl-capped result, the second the mechanism behind it.

Table 3: End-to-end batch-1 latency vs. input resolution. $N$ is the joint token count; speedup is baseline over WHTMix.

Resolution	$N$	Baseline (ms)	WHTMix (ms)	Speedup
$256\times 512$	4,096	8.6	7.5	$1.15\times$
$384\times 768$	9,216	33.6	14.8	$2.27\times$
$512\times 1024$	16,384	95.6	36.1	$2.65\times$
$640\times 1280$	25,600	226.3	80.9	$2.80\times$
$768\times 1536$	36,864	453.5	166.5	$2.72\times$
$896\times 1792$	50,176	848.8	298.0	$2.85\times$

Figure 3: End-to-end batch-1 latency vs. input resolution (CARLA aspect). Left: linear, with per-point speedup. Right: log-log; fitted slopes approximate the empirical complexity exponent.

Figure 4: Isolated token-mixer latency vs. token count $N$ (batch-1, power-of-two $N$ so no padding, RTX 4090); only the mixer that differs between the two models is timed.

IV-D Beyond Stereo: Long-Range Arena

To test whether the benefit is stereo-specific or depends only on the token-to-channel ratio, we insert WHTMix as a drop-in self-attention replacement into a plain sequence classifier ($C{=}64$, $L{=}4$, $8$ heads) and evaluate on two Long-Range Arena [18] tasks unrelated to stereo: byte-level image classification (CIFAR-10 grayscale, $N{=}1024$) and byte-level text sentiment (IMDB, $N{=}4096$). Only the token mixer differs between paired runs. Table 4 and Fig. 5 reproduce the stereo trend: the compute reduction grows from $4.3\times$ at $N/C{=}16$ to $13.8\times$ at $N/C{=}64$, and the wall-clock advantage is absent at the small size, where WHTMix is launch-overhead-bound and slower than the fused attention kernel, emerges at the larger one ($1.9\times$ faster). Accuracy follows the same regime: on the long-sequence text task WHTMix is within $1.6$ points of attention, while on the shorter image task the fixed mixer trails by $6.9$ points. On unrelated data, then, the mixer’s utility is governed by $N/C$ rather than by the task; these are single-seed runs and we defer multi-seed confirmation to future work.

Table 4: Beyond stereo: WHTMix vs. self-attention on two Long-Range Arena [18] tasks (byte-level), identical config ($C{=}64$, $L{=}4$), single seed; only the mixer differs. Accuracy, GFLOPs, and batch-1 latency on an RTX 4090.

Task ($N/C$)	Mixer	Acc.	GFLOPs	Lat. (ms)	FLOP$\downarrow$
Image (16)	attn	0.383	1.48	1.47	—
Image (16)	WHTMix	0.314	0.34	3.22	$4.3\times$
Text (64)	attn	0.632	18.81	6.10	—
Text (64)	WHTMix	0.617	1.37	3.14	$13.8\times$

Figure 5: Accuracy vs. compute on two non-stereo Long-Range Arena tasks. Arrows mark the self-attention$\to$WHTMix compute reduction: the FLOP saving grows with $N/C$ ($4.3\times$ image, $13.8\times$ text) while the accuracy gap shrinks, WHTMix nearly matching attention on the long-sequence text task.

IV-E Discussion

The results support a scoped claim: for high-resolution, low-channel transformers ($N\!\gg\!C$), dense-prediction vision in particular, replacing an $O(N^{2})$ global self-attention with an $O(N\log N)$ mixer yields large, resolution-growing compute savings at parity accuracy, while genuine cross-attention is retained for correspondence. The favorable regime is exactly the one Eq. (6) predicts; classification transformers ($N\!<\!C$) fall outside it. The same $N/C$ dependence holds on two non-stereo Long-Range Arena tasks (Table 4), indicating the effect is regime-driven rather than task-specific. The KITTI gap indicates that a fixed mixer adapts less well than attention when real training data are scarce, motivating larger-scale real-data pre-training as future work.

V Conclusion

We replaced the joint self-attention of a stereo transformer with a learnable Walsh-Hadamard token mixer, retaining data-dependent cross-attention for matching. The mixer matches attention accuracy on synthetic data at $2.46\times$ less compute and $2.65\times$ lower latency, a benefit our $2+N/C$ analysis both predicts and bounds. A hybrid log-disparity loss further improves distant-object accuracy for free. Future work includes multi-seed confirmation, larger real-data pre-training, and applying the mixer to other long-sequence dense-prediction transformers.