RestoreVAR: Visual Autoregressive Generation for All-in-One Image Restoration

In plain words

Old photos and camera frames often suffer from haze, rain, snow, darkness or blur, and a single model that undoes all of these is called all-in-one restoration. The best-looking recent versions use image generators that start from noise and refine it step by step. They look good, but each image takes many slow passes.

RestoreVAR takes a different generator. It is built on a model that draws an image as a series of ever-finer layers: first a rough sketch of the scene, then more and more detail. The authors noticed that in this series the rough layers mostly carry the damage, such as haze or snow, while the fine layers carry the actual scene. So the restoration task becomes "get the rough layers right", and the rest follows.

The model reads the damaged picture, so it knows what to fix, and it writes the clean picture one layer at a time. A small add-on sharpens the internal representation before it is turned back into pixels, and the decoder that does this turn was adjusted to keep fine edges. The result is faster than the diffusion-based systems, with more faithful detail. It is still a generative model, and it does not match the best non-generative restorers on pixel scores.

TL;DR

RestoreVAR is a generative all-in-one restoration model built on a pretrained visual autoregressive (VAR) transformer rather than a latent diffusion model. The authors report higher PSNR, SSIM and LPIPS than the diffusion-based generative baselines on five restoration datasets, and about 10 times faster inference (0.201 s per image against 2.04 s for Diff-Plugin). It still trails the non-generative models on pixel metrics, and its fidelity ceiling is set by the VAE decoder. The released code has no license, and its training schedule differs from the paper in several places.

Problem

Five degradations are targeted: haze, snow, rain, low light and blur §4.2. Latent diffusion restorers need many denoising passes per image, which is too slow for time-sensitive uses such as video surveillance or navigation §1. Their VAE decoders are tuned for generative variety rather than exact pixel reconstruction, which loses fine structure §1. Non-generative all-in-one restorers are fast and accurate on their training distribution but generalise poorly to real degradations §1. The authors observe that VAR is fast, but only two earlier works used it for restoration, and one of those covered super-resolution only §1.

Method

RestoreVAR has three parts §3.3: (1) a VAR transformer trained directly for restoration, with the degraded image injected through cross-attention; (2) a fine-tuned VAE decoder that works on continuous latents; (3) a non-generative latent refinement transformer (LRT) that turns predicted discrete latents into continuous ones before decoding §3.3, §3.3.2.

VAR background. VAR encodes an image with a multi-scale VQ-VAE. Its residual quantiser produces one index map per scale, using a shared codebook, and the number of scales is K = 10 §3.1. The transformer is a decoder-only model with a block-wise causal mask, so each scale attends only to earlier scales. Training uses teacher forcing and cross-entropy on the ground-truth index maps. Inference predicts one scale at a time, upsamples and embeds the result, and feeds it to the next scale, using key-value caching §3.1.

Scale analysis. Coarse scales (k = 1 to 5) hold most of the degradation; fine scales (k = 6 to 10) hold scene detail. Swapping coarse scales between degraded and clean reconstructions moves the degradation, while swapping fine scales does not §3.2, fig. 2. The authors argue that restoration therefore reduces to predicting the early, small scales correctly §3.2.

Degraded-image conditioning. The start token is a fixed label embedding plus a global summary of the degraded latent. That summary is a mean over a learned projection, scaled by a gain initialised at 0 §3.3.1 app. F. The continuous (pre-quantisation) latent of the degraded image is the key and value source for a cross-attention added to every block. Its output is multiplied by a per-block gate initialised at 0, so the pretrained behaviour is preserved at the start of training §3.3.1 fig. 3.

# RestoreVAR inference  [§3.3] [§3.3.2] [app. F]
# Input:  degraded image, 512 x 512 x 3
# Output: restored image, 512 x 512 x 3

def restore(degraded):
    # Continuous latent of the degraded input, used for conditioning  [§3.3.1]
    cond_latent = vae_encode_continuous(degraded)    # (C, 32, 32)
    # Start token: fixed label embedding plus a gated global summary  [§3.3.1]
    sos = label_embedding + gain * mean(project(cond_latent))
    # Predict scales coarse to fine, taking the most likely index  [§3.1]
    for scale in [1, 2, 3, 4, 6, 9, 13, 18, 24, 32]:
        # Each block adds cross-attention to the degraded latent  [§3.3.1]
        hidden = transformer_blocks(tokens, cond_latent, scale)
        indices = argmax(predict_logits(hidden))
        # Embed, upsample and accumulate; feed forward to the next scale  [§3.1]
        feed_next_scale(indices)
    # Refiner adds a residual from the final hidden states  [§3.3.2]
    discrete_latent = accumulated_latent()
    continuous = discrete_latent + lrt(discrete_latent, all_hidden)
    # Decode with the fine-tuned VAE decoder  [§3.3.2]
    return vae_decode(continuous)

Decoder fine-tuning. The VAE encoder and quantiser are frozen. The decoder is trained on pairs of (continuous ground-truth latent, clean image), with a loss that is a weighted sum of L1, SSIM, perceptual and adversarial terms (a PatchGAN discriminator) §3.3.2 §4.1. The authors report 28.14 dB PSNR and SSIM 0.842 for this decoder on average over 1000 samples, against 22.59 dB and 0.679 for the VAR VQ-VAE and 26.48 dB and 0.804 for HART's decoder §3.3.2. The dataset of those 1000 samples is not stated.

Latent refinement transformer (LRT). The LRT predicts a residual that is added to the discrete latent, giving a continuous latent §3.3.2. It receives the final-block hidden states of the restoration transformer as guidance, which the authors call critical for performance §3.3.2 §4.4. It is trained with an L1 loss against the ground-truth continuous latent §3.3.2. The authors report about 3% extra inference overhead, against about 20% for HART's diffusion-based refiner §3.3.2.

# Refiner training (VAR frozen)  [§3.3.2] [app. F]
# Input:  degraded image, clean image
# Output: L1 loss on continuous latents

def refiner_loss(degraded, clean):
    # Target: continuous latent of the clean image  [§3.3.2]
    target_latent = vae_encode_continuous(clean)
    # Only refiner weights are trainable; transformer frozen
    # [code: train_refiner.py#L115-L117]
    refined = greedy_inference_with_refiner(degraded)
    # L1 between predicted and ground-truth continuous latents  [§3.3.2]
    return l1(refined, target_latent)

Transformer training. The transformer is trained with teacher forcing on clean index maps, conditioned on the degraded image's continuous latent, using cross-entropy averaged over all tokens §3.1 §3.3.1.

Architecture changes. Absolute position embeddings are replaced by 2D rotary embeddings so the model can run at 512 x 512 instead of VAR's 256 x 256 §3.3.1 app. D. The authors remove AdaLN layers from the transformer, saving about 100M parameters with little effect on quality §3.3.1. The main transformer has depth 16, width 1024 and 16 heads; the refiner has depth 12, 6 heads, width 384 and an MLP of 768 app. F.

Training regime. Each component is trained separately to keep their objectives apart §4.1. The release's refiner script loads a trained transformer checkpoint and a separate decoder checkpoint train_refiner.sh, so the refiner is trained after the other two. The VAR backbone was pretrained on ImageNet-1K, about 1M images §4.3.

Hyperparameters

Values are from the paper unless the Where column says "code". "Not stated" means neither the paper nor the code gives it.

Results

Author-reported. Setup: five restoration tasks, each with its own training and test sets §4.2. Baselines: non-generative PromptIR, InstructIR, AWRaCLe, DCPT and DFPIR; latent-diffusion Diff-Plugin, AutoDIR and PixWizard §4.3. PromptIR and AWRaCLe were retrained by the authors; the others used public checkpoints. AutoDIR is evaluated without its structure-correction module §4.3 app. I. Cells show PSNR (dB) / SSIM / LPIPS tab. 1.

Table 1 tab. 1 (within-distribution, five datasets). The authors note that non-generative methods score higher on these pixel metrics, which they attribute to the VAE decoder's limit §4.3.

Speed §4.3 app. B: time per image and steps. Diff-Plugin 2.04 s (20 steps); AutoDIR 8.477 s (100 steps); PixWizard 8.247 s (60 steps); RestoreVAR 0.201 s (10 steps). Parameters: 859.50 M, 859.50 M, 2011.40 M and 296.95 M. TFLOPs: 16.08, 67.80, 19.27 and 1.05. The authors report a 10x speed-up over Diff-Plugin and over 40x over AutoDIR and PixWizard. Non-generative methods are faster (InstructIR 0.023 s, DCPT 0.024 s, PromptIR 0.162 s, AWRaCLe 0.188 s, DFPIR 0.117 s) app. B.

Generalisation §4.3 tab. 2: no-reference MUSIQ / CLIPIQA on unseen and mixed datasets. Averages are over the six datasets.

Against the non-generative methods, RestoreVAR has the best average. Against the diffusion methods, AutoDIR has a higher average MUSIQ (56.767 vs 56.217) and PixWizard a higher average CLIPIQA (0.394 vs 0.383) tab. 2. The authors describe the result as competitive and attribute the gap to backbone pretraining data: Stable Diffusion's training set is much larger than VAR's ImageNet-1K pretraining §4.3.

User study §4.3 tab. 3: 50 real-world scenes, 36 responses, 5-point scale covering scene consistency, restoration quality and overall appeal. Means: PromptIR 2.11, InstructIR 2.93, AWRaCLe 2.33, DCPT 2.42, DFPIR 2.35, AutoDIR 3.68, RestoreVAR 4.36 tab. 3. The paired tests reported in app. E find the gains significant against AutoDIR and four non-generative baselines. The Qualtrics survey details are in app. E.

LRT ablation tab. 4 §4.4. Columns: refiner time per image, parameters, PSNR / SSIM. No refiner: no time, no parameters, 21.71 / 0.690. HART refiner: 0.0455 s, 36.06 M, 23.48 / 0.777. LRT without final-block guidance: 0.0036 s, 14.61 M, 21.23 / 0.660. Proposed LRT: 0.0061 s, 22.97 M, 24.67 / 0.821. The caption does not name the dataset; the proposed row matches the RESIDE column of Table 1. The HART refiner is about 7 times slower than the LRT §4.4. With no LRT, MUSIQ / CLIPIQA is 66.56 / 0.472; with LRT it is 66.45 / 0.481 §4.4. The dataset for these two scores is not named.

Other reported results. Continuous conditioning beats discrete conditioning on validation accuracy over 15 epochs §4.4 fig. 7. Adding the adversarial term sharpens decoder output §4.4 fig. 8. Ablations on the SSIM and perceptual loss terms appear in app. D.

Paper vs code

Official code: https://github.com/sudraj2002/RestoreVAR at commit 4ffe8880ab5af5bc4469848bceba7ecbbe0aadf5. The paper's reproducibility statement says the code would be released after review. I searched for it in three steps: the paper's own links (the project page, which links this repository); the GitHub API search for "RestoreVAR", which returned the repository owned by the first author's account (the README links that account to the author's homepage); and no other candidates. The repository has no LICENSE file; the GitHub API reports no license. Code is cited only, not quoted.

Discrepancies found, each with both citations:

• Transformer training length. Paper: 100 epochs §4.1. Script: --ep=50 train.sh.
• Refiner batch and length. Paper: batch 96, 100 epochs §4.1. Script: batch 48 and 50 epochs train_refiner.sh.
• GPUs. Paper: 8 x RTX A6000 §4.1. Script: 2 processes train.sh. The batch is global, and the code scales the learning rate linearly with batch, so the effective peak rate is 1e-4 in both arg_util.py.
• Decoder fine-tuning. The paper gives the loss weights, discriminator and schedule §3.3.2 §4.1. The release contains no decoder training code; the README points to a pretrained decoder checkpoint README.md. The weights and schedule cannot be checked against code.
• AdaLN. The paper says AdaLN layers are removed §3.3.1. The transformer blocks carry no AdaLN basic_var.py basic_var.py, but the output head still uses the conditioning vector through an AdaLN head var.py.
• Conditioning latent. The paper uses the continuous latent of the degraded image §3.3.1. The code keeps four encoder levels, then uses only the final one, resized to 32 x 32 with bicubic interpolation vqvae.py var.py.
• Refiner guidance. The paper says the guidance is the final block's output §3.3.2. The code concatenates the final-block output at every scale into one sequence var.py var.py, which matches the paper's L x D description.
• Inference decoding. The paper does not say how indices are chosen. The code takes the most likely index at each scale, so output is deterministic var.py.
• Learning-rate warm-up. The code's warm-up runs one epoch at the full rate arg_util.py lr_control.py, so there is none in effect. The paper gives no schedule.
• Weight decay and clipping. The code applies a constant weight decay of 0.05 and clips gradients at norm 2.0 arg_util.py arg_util.py lr_control.py. The paper gives neither.
• Evaluation resizing. Test images are resized to 512 x 512 without keeping the aspect ratio, and the metrics are computed on that resized pair data.py infer.py. Training instead resizes the shorter side to 576 and takes a random 512 crop data.py. The paper does not state either choice.
• Refiner training. The paper says each part is trained separately §4.1. The code freezes everything except the refiner train_refiner.py, and computes its loss through the greedy autoregressive pass trainer_refiner.py.

No discrepancies were found for: VAR depth 16 and width 1024 app. F; the 2D rotary embedding var.py; the zero-initialised gates basic_var.py var.py; the refiner dimensions vae_refiner_rope.py; the patch schedule to 32 arg_util.py.

Pitfalls

• Teacher forcing against greedy decoding. Training uses ground-truth indices, while inference uses the model's own argmax at each scale. Errors in early scales then feed every later scale. The scale analysis §3.2 implies early scales matter most for restoration; this is an inference from the paper, not a measured failure rate.
• Fidelity ceiling. The authors state that pixel metrics are bounded by the VAE decoder and that ground-truth latents would do better §4.3 app. J. The fine-tuned decoder alone reaches 28.14 dB on average §3.3.2, so the decoder limits results even with perfect prediction.
• Refiner cannot restore alone. Feeding degraded latents through the LRT leaves the output hazy, so the LRT only refines the transformer's output app. D. Do not use the refiner as a standalone restorer.
• Severe degradation. Very dark inputs (SID) produce output clearly worse than ground truth app. J.
• Fixed resolution. The model is trained and evaluated at 512 x 512 §3.3.1 app. F. Other sizes need the position scheme changed, and the test transform distorts non-square inputs data.py.
• Baseline choices change the picture. AutoDIR is run without its structure-correction module, which the authors show acts as an independent restorer app. I. The paper's speed-up figures compare against baselines with their own step counts and resolutions, which the paper does not fully state.
• Diffusion with fewer steps. With DDIM, diffusion baselines match RestoreVAR's time only at about two steps, and then lose more than 3 dB PSNR app. B.
• Asymptotic speed argument. The O(log n) advantage over latent diffusion assumes quadratic self-attention cost and a fixed number of forward passes app. C. It is a model of cost, not a measurement.
• Small margins on some metrics. On RESIDE the PSNR lead over AutoDIR is 0.19 dB (24.67 vs 24.48), and on the generalisation averages AutoDIR and PixWizard score higher on MUSIQ and CLIPIQA respectively tab. 1 tab. 2.
• Unnamed datasets. The decoder reconstruction numbers §3.3.2, the LRT ablation tab. 4 and the ablation MUSIQ / CLIPIQA numbers §4.4 do not name a dataset.
• Missing baselines. InstructIR and DCPT do not report Snow100k results, and DFPIR does not either tab. 1. The table has no entry for those cells.
• Licensing of code. The repository has no license, so code can be cited but not quoted or reused. Implement from the paper and cite the code's behaviour by path and line.

Lineage

• Builds on: VAR, visual autoregressive modelling with next-scale prediction §3.1. The arXiv identifier is not in the paper's bibliography, so the card does not link it.
• Builds on: HART's latent refinement idea, with a diffusion-based refiner that the authors replace by a non-generative one §2.2 §3.3.2.
• Prior VAR-for-restoration works: VarSR (super-resolution only) and Varformer (non-generative restorer guided by VAR features) §1.
• Compared against: PromptIR, InstructIR, AWRaCLe, DCPT, DFPIR; Diff-Plugin, AutoDIR, PixWizard §4.3.
• Superseded by: not known from this card's sources.