Wavelet Domain Masked Image Modeling

• W-MIM is a self-supervised framework that reconstructs wavelet coefficients to capture hierarchical, frequency-aware image and video representations.
• It aligns encoder layers with specific wavelet decomposition levels to efficiently recover structural details while minimizing redundant fine texture.
• In HDR video reconstruction, wavelet-domain masking pre-trains networks for color-consistent, high-fidelity restoration at a fraction of the computational cost.

Searching arXiv for the cited papers and closely related MIM baselines. Wavelet-domain Masked Image Modeling (W-MIM) denotes a class of masked reconstruction methods in which the supervisory signal, the masking operation, or both are defined in a wavelet representation rather than directly in pixel space. In the image-representation setting of "Wavelet-Driven Masked Image Modeling: A Path to Efficient Visual Representation" (Xiang et al., 2 Mar 2025), wavelet coefficients at multiple decomposition levels are used as compact reconstruction targets for masked image modeling. In the HDR video setting of "Wavelet-Domain Masked Image Modeling for Color-Consistent HDR Video Reconstruction" (Zhang et al., 7 Feb 2026), masking is applied directly to wavelet subbands to pre-train an encoder for color-consistent reconstruction. Across both formulations, the central premise is that pixel-based MIM overemphasizes redundant fine detail, whereas wavelet decomposition supplies frequency-aware, spatially localized, and multi-level structure that can be aligned with hierarchical neural architectures.

1. Conceptual definition and scope

Masked Image Modeling (MIM) is presented as a self-supervised paradigm for learning scalable visual representations, but the wavelet-domain variants argue that images inherently contain abundant redundant information, causing pixel-based reconstruction to focus excessively on finer details such as textures and thereby prolong training times unnecessarily (Xiang et al., 2 Mar 2025). W-MIM addresses this by shifting reconstruction from raw pixels to a compact representation of features in the wavelet domain.

The key distinction from Fourier-based frequency analysis is explicit in the underlying formulation: wavelet transform not only offers frequency information but also preserves spatial characteristics and multi-level features of the image (Xiang et al., 2 Mar 2025). In the video-reconstruction formulation, this property is extended to color and detail restoration, where wavelet-domain masking is used to reconstruct brightness, color, and fine structure from low dynamic range inputs (Zhang et al., 7 Feb 2026).

Two concrete instantiations define the current scope of W-MIM. The first is WaMIM, an image pre-training framework that assigns different wavelet bands to different encoder depths and optimizes a multi-scale reconstruction loss (Xiang et al., 2 Mar 2025). The second is WMNet, an HDR video reconstruction network in which W-MIM is the Phase I self-reconstruction pre-training strategy preceding temporal modeling with T-MoE and DMM (Zhang et al., 7 Feb 2026). This suggests that W-MIM is better understood as a design pattern—wavelet-structured masking and reconstruction—than as a single fixed architecture.

2. Wavelet decomposition as the representational substrate

In WaMIM, a 2D image $x\in\mathbb{R}^{M\times N}$ undergoes a one-level discrete wavelet transform (DWT) that produces one low-frequency approximation coefficient map $w^0_{\phi}$ and three high-frequency detail maps $w^0_{\psi,i}$, $i\in\{H,V,D\}$ (Xiang et al., 2 Mar 2025). In practice, DWT is implemented by separable filter banks, for example Haar, and recursively applied to the low-frequency branch. After $J$ levels, the decomposition is written as

$$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$

where each $w^j_{\psi}$ bundles the three detail maps $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$ (Xiang et al., 2 Mar 2025).

In WMNet, the corresponding frame-level formulation uses a discrete 2-D Haar wavelet transform $W(\cdot)$ at $N$ levels for an LDR frame $w^0_{\phi}$0. At each level $w^0_{\phi}$1, the transform yields four subbands: one approximation $w^0_{\phi}$2 and three detail bands $w^0_{\phi}$3. With separable filters

$w^0_{\phi}$4

the single-level transform is defined by

$w^0_{\phi}$5

$w^0_{\phi}$6

with higher levels recursively applied to the preceding $w^0_{\phi}$7 branch (Zhang et al., 7 Feb 2026).

A common misconception is that frequency-domain modeling necessarily discards localization. The wavelet formulations used here do not: the image paper explicitly contrasts wavelets with the commonly used Fourier transform on the grounds that wavelets preserve spatial characteristics and multi-level features, while the video paper operationalizes that claim by reconstructing masked wavelet coefficients back into an image through inverse DWT before encoding (Xiang et al., 2 Mar 2025, Zhang et al., 7 Feb 2026).

3. Multi-scale target assignment and loss design in WaMIM

WaMIM selects $w^0_{\phi}$8 encoder layers $w^0_{\phi}$9 and sets $w^0_{\psi,i}$0. The assignment rule is depth-to-frequency aligned: shallow layers reconstruct high-frequency details, while deep layers reconstruct low-frequency content. For layer $w^0_{\psi,i}$1, the target level is

$w^0_{\psi,i}$2

The target generator is defined as

$w^0_{\psi,i}$3

where $w^0_{\psi,i}$4 denotes spatial-channel concatenation (Xiang et al., 2 Mar 2025).

The masking operator is $w^0_{\psi,i}$5, the encoder is $w^0_{\psi,i}$6, and the layer-wise decoder is $w^0_{\psi,i}$7. If

$w^0_{\psi,i}$8

and $w^0_{\psi,i}$9, the total loss is

$i\in\{H,V,D\}$0

An equivalent form is

$i\in\{H,V,D\}$1

where $i\in\{H,V,D\}$2 (Xiang et al., 2 Mar 2025).

Architecturally, the framework plugs in any ViT or Swin backbone, masks patches at the input with mask ratio $i\in\{H,V,D\}$3, extracts hidden features $i\in\{H,V,D\}$4 from selected layers, and attaches to each chosen layer a light-weight decoder consisting of one Transformer block plus a small linear head (Xiang et al., 2 Mar 2025). The multi-scale targets are fed independently to each decoder; there is no pooling or mixing across scales. This clean separation is explicitly intended to let each decoder focus on one frequency band while fusing spatial structure with a frequency-domain target.

4. Training protocol, ablations, and reported performance of WaMIM

The reported WaMIM training setup uses ImageNet-1K at resolution $i\in\{H,V,D\}$5, block-wise random masking at $i\in\{H,V,D\}$6 of patches, Haar wavelet DWT with $i\in\{H,V,D\}$7 decomposition levels, and selected levels $i\in\{H,V,D\}$8 as targets for $i\in\{H,V,D\}$9 (Xiang et al., 2 Mar 2025). The optimizer is AdamW with weight decay $J$0, $J$1, $J$2, batch size $J$3, cosine learning-rate schedule, and warmup of $J$4–$J$5 epochs depending on total pre-train epochs. The base learning rate is $J$6 for ViT and $J$7 for Swin. Pre-training is reported for $J$8 epochs with fixed loss weights $J$9 and no curriculum on weights (Xiang et al., 2 Mar 2025).

The ablation results identify several operating points. The best mask ratio is $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$0; too low is described as under-constrained, and too high ($$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$1) as under-learned. Testing $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$2, $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$3, and $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$4 decomposition levels shows $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$5 to be best because it covers a broader range of scales and frequencies. Light-weight decoders with $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$6–$$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$7 dimensions and $$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$8–$$[\,w^0_{\phi};\; w^0_{\psi},\,w^1_{\psi},\,\dots,\,w^{J}_{\psi}\;] = f_w(x),$$9 heads suffice, matching heavier decoders’ accuracy. Uniform spacing of feature locations, exemplified by $w^j_{\psi}$0 in ViT-B and $w^j_{\psi}$1 in Swin-B, outperforms other heuristics. Emphasizing low-frequency, high-level targets moderately with $w^j_{\psi}$2 yields the best results, whereas reversing that ordering degrades accuracy (Xiang et al., 2 Mar 2025).

Setting	Baseline	WaMIM
ImageNet-1K, ViT-B	MAE, 1600 epochs, 752 GPU hrs, 83.6 Top-1	400 epochs, 96 GPU hrs, 83.8 Top-1
ImageNet-1K, ViT-S	MAE, 300 epochs, 96 GPU hrs, 80.9 Top-1	300 epochs, 45 GPU hrs, 82.0 Top-1
ImageNet-1K, Swin-B	SimMIM, 800 epochs, 656 GPU hrs, 84.0 Top-1	400 epochs, 160 GPU hrs, 84.1 Top-1

For downstream transfer, COCO object detection and instance segmentation with Mask R-CNN and Swin-B report AP$w^j_{\psi}$3=$w^j_{\psi}$4 and AP$w^j_{\psi}$5=$w^j_{\psi}$6 for WaMIM, compared with $w^j_{\psi}$7 and $w^j_{\psi}$8 for SimMIM, at $w^j_{\psi}$9 versus $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$0 pre-train hours. ADE20K semantic segmentation with UperNet and ViT-B reports mIoU $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$1 for WaMIM versus $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$2 for MAE, at $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$3 versus $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$4 pre-train hours (Xiang et al., 2 Mar 2025). The paper summarizes these findings by stating that WaMIM matches or exceeds pixel-based MIM accuracy at a small fraction ($\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$5–$\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$6) of the GPU time.

5. Wavelet-domain masking and two-phase training in WMNet

In WMNet, W-MIM is not a generic representation-learning objective but the pre-training phase of an HDR video reconstruction pipeline (Zhang et al., 7 Feb 2026). The masking strategy introduces two binary masks. A random mask $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$7 is applied to the deepest $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$8, covering a proportion $\{w^j_{\psi,H},w^j_{\psi,V},w^j_{\psi,D}\}$9 of its coefficients. A full-zero mask $W(\cdot)$0 is applied to all high-frequency bands $W(\cdot)$1, forcing the network to reconstruct edges, textures, and chromatic variations from the remaining approximation. Formally,

$W(\cdot)$2

where $W(\cdot)$3, and the masked coefficients are

$W(\cdot)$4

The masked frame is then

$W(\cdot)$5

Phase I trains an asymmetric encoder–decoder $W(\cdot)$6–$W(\cdot)$7 to reconstruct the original LDR frame from $W(\cdot)$8 with the $W(\cdot)$9 loss

$N$0

The encoder is specified as $N$1 ResBlocks of width $N$2, and the decoder is lightweight (Zhang et al., 7 Feb 2026). Difficulty is increased gradually through curriculum learning: $N$3 so the low-frequency mask ratio rises from zero to $N$4 over the first $N$5 iterations or epochs.

Phase II initializes the video encoder $N$6 with the pre-trained weights and adds a full-sized video decoder $N$7, the Temporal Mixture of Experts (T-MoE), and the Dynamic Memory Module (DMM). For an input clip $N$8, the network computes $N$9, conditions the features temporally, and reconstructs HDR output as

$w^0_{\phi}$00

The reconstruction loss is

$w^0_{\phi}$01

with $w^0_{\phi}$02, and no extra regularizers beyond weight decay (Zhang et al., 7 Feb 2026).

6. Architectural integration, empirical effects, and interpretive implications

WMNet inserts W-MIM before any 2-D convolutional layer: each input frame is wavelet-decomposed, masked, inverse-transformed to produce $w^0_{\phi}$03, and then sent into a shared encoder defined as $w^0_{\phi}$04 with hidden dimension $w^0_{\phi}$05 (Zhang et al., 7 Feb 2026). The decoder uses upsampling by PixelShuffle and convolution layers to reconstruct spatial resolution. T-MoE splits the encoder’s $w^0_{\phi}$06 ResBlocks into $w^0_{\phi}$07 groups, applies a small expert network $w^0_{\phi}$08 over a temporal window of three adjacent frames, and fuses group outputs as

$w^0_{\phi}$09

DMM maintains a per-scene queue $w^0_{\phi}$10 of up to $w^0_{\phi}$11 memory features, computes cross-attention between current features and scene memory, and updates memory through convolutional refinement (Zhang et al., 7 Feb 2026).

The ablation on HDRTV4K-Scene isolates the contribution of W-MIM. The baseline without W-MIM reports PSNR $w^0_{\phi}$12 dB, SSIM $w^0_{\phi}$13, and $w^0_{\phi}$14. Adding W-MIM only yields PSNR $w^0_{\phi}$15 dB, SSIM $w^0_{\phi}$16, and $w^0_{\phi}$17. A low-frequency-only variant reports PSNR $w^0_{\phi}$18 and $w^0_{\phi}$19 dB, while a high-frequency-only variant reports $w^0_{\phi}$20, described as best color. A spatial-domain MIM with the same overall mask rate yields PSNR $w^0_{\phi}$21 and $w^0_{\phi}$22. With T-MoE and DMM added, PSNR reaches $w^0_{\phi}$23 dB and $w^0_{\phi}$24 drops to $w^0_{\phi}$25, with state-of-the-art across seven metrics: PSNR, SSIM, SR-SIM, $w^0_{\phi}$26, HDR-VDP3, LPIPS, and $w^0_{\phi}$27 (Zhang et al., 7 Feb 2026).

Configuration	Reported outcome	Interpretation
No W-MIM	PSNR 35.15, SSIM 0.9612, $w^0_{\phi}$28	Reference point
+ W-MIM only	PSNR 36.00, SSIM 0.9627, $w^0_{\phi}$29	Better fidelity and color
+ W-MIM + T-MoE + DMM	PSNR 36.23, $w^0_{\phi}$30	Best full-system result

Taken together, the image and video formulations support a unified interpretation. WaMIM uses multi-scale wavelet coefficients as compact, frequency-aware reconstruction targets and aligns encoder depth with wavelet scale (Xiang et al., 2 Mar 2025). WMNet instead uses wavelet-domain masking to pre-train color and detail restoration before temporal reconstruction (Zhang et al., 7 Feb 2026). A plausible implication is that W-MIM is most effective when the decomposition hierarchy is explicitly coupled to the architecture or task structure rather than treated as a purely alternative signal representation.