Three-Stage Framework
Overview of SLAIR: (I) Degradation-Aware Latent Separation, (II) Conditional Flow Matching Transport, (III) High-Frequency Refinement Decoder.
Degradation-Aware Latent Separation
A fundamental challenge in all-in-one restoration is that the latent codes of degraded and clean images are highly entangled in a vanilla VAE, making it difficult to define a clear transport direction. We address this by fine-tuning the encoder with two auxiliary losses: a binary separation loss that maximally separates the distributions of clean and degraded latents, and a degradation-type classification loss over ten degradation categories. This structured latent space ensures that the flow model always transports in a well-defined and consistent direction, without requiring explicit degradation-type labels at inference time.
Conditional Flow Matching in Latent Space
Given paired degraded and clean images, SLAIR encodes them into the structured VAE latent space and learns a conditional vector field that deterministically transports $z_{\mathrm{deg}}$ to $z_{\mathrm{gt}}$. As described in Section 3.3 of the SLAIR paper, the path is trained with bridge-noised intermediate states, velocity regression, and cosine alignment, then restored by solving an ODE from $t=0$ to $t=1$.
High-Frequency Refinement Decoder
Latent-space restoration may lose high-frequency structures during decoding due to the low-pass nature of VAE decoders. To recover fine-grained texture, edges, and perceptual detail, we augment the VAE decoder with a High-Frequency Refinement Module (HRM). The HRM injects multi-scale encoder features via lightweight fusion blocks at each decoder resolution, and applies a discrete wavelet transform (DWT) loss on the decoded output to directly supervise high-frequency subbands (LH, HL, HH). This prevents the decoder from reintroducing degradation artifacts (rain streaks, haze, moiré) while recovering sharp structures. The decoder remains frozen after Stage 1; only the fusion blocks and HRM parameters are trained.
