PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient JPEG Restoration in the Wavelet Domain via Mean Flows

Stefan-Alexandru Asandei, Mihai-Alexandru Radu

arXiv:2608.28730v1eess.IVcs.CV

TL;DR

JPEG restoration must handle detail lost under aggressive compression, while strong pretrained diffusion restorers can be costly to deploy. This paper trains a compact wavelet-domain MeanFlow restorer from scratch; its reported throughput and LPIPS results position it as a quality-efficiency trade-off, with pretrained priors retaining an advantage under severe compression.

  • Problem

    Aggressive JPEG quantization irretrievably removes high-frequency coefficients, leaving restoration ill-posed, while strong pretrained diffusion restorers retain costly latent pipelines.

    Source: §1 · p. 1

  • Method

    A from-scratch conditional MeanFlow model predicts a clean residual in an exactly invertible two-level Haar domain, using a RELA-based linear DiT that estimates JPEG severity internally.

    Source: §1 · pp. 1–2, §3.1 · p. 2

  • Results

    8.05 images/s at 1024^2 with 65.32M parameters; the reported model is not uniformly best in perceptual quality.

    Source: §6 · p. 7

  • Takeaways & Limitations

    Two-update inference gives the experiments’ best all-operating point, while a single-evaluation variant is available when throughput is the priority.

    Source: §1 · p. 1

    At QF 5 and on MANIQA and DISTS, the model’s limited semantic knowledge trails pretrained text-to-image systems, especially when compression removes evidence for object-specific detail.

    Source: §5 · p. 7

Abstract

from arXiv · show

Latest JPEG restoration systems achieve strong quality with large models, yet often remain too slow and expensive for efficient on-device deployment. We present a 65M-parameter generative restorer that attains the lowest LPIPS at QF 10 and 20 on LIVE-1, Urban100, and DIV2K-val while sustaining 8.05 images/s at $1024\times1024$ on a single RTX 3090, roughly $4.9\times$ the reported throughput of one-step SODiff at one-twentieth of its parameters. Trained from scratch, the model replaces the learned VAE encoder-decoder with an exactly invertible two-level Haar transform, predicts a clean wavelet-domain residual through a rank-enhanced linear-attention DiT that estimates compression severity internally, and is optimized with an improved MeanFlow objective that enables inference in one or two network evaluations without distillation. Large pretrained priors remain stronger under severe compression (QF 5), whereas our model prioritizes throughput for deployment-constrained restoration.

1 INTRODUCTION

JPEG restoration must recover uncertain high-frequency detail while balancing perceptual quality against the cost of large pretrained pipelines. This work presents a compact, from-scratch MeanFlow restorer and evaluates its quality-efficiency trade-offs.

1 INTRODUCTION

  • At aggressive compression, quantized high-frequency coefficients are unrecoverable, making restoration ill-posed because detail must be suppressed or inferred without unsupported hallucination.Artifacts include blocking, ringing, color discontinuities, and lost texture.Source: §1 · p. 1
  • Pretrained diffusion priors can deliver strong perceptual quality, but SODiff retains a Stable Diffusion 2.1 pipeline that is difficult to deploy under memory and latency constraints.Its inference cost comes from the UNet, VAE encoding and decoding, and semantic conditioning rather than a long sampling trajectory [4].Source: §1 · p. 1
  • The 65.32M-parameter model is trained from scratch without a text encoder, pretrained generator, or VAE, using an exactly invertible two-level Haar representation and improved MeanFlow.Two network evaluations yield the best all-operating point in the experiments; a single-evaluation variant is available when throughput is the priority.Source: §1 · p. 1
  • 8.05 images/s at 1024 × 1024 on one RTX 3090, while matching or improving LPIPS at QF 10 and 20 across three datasets.The authors frame this as a practical Pareto point, not uniformly superior restoration.Source: §1 · p. 1
  • Its restoration backbone combines RELA with a SANA-style linear DiT and estimates unknown JPEG severity internally.The contribution list also describes clean-residual prediction in the two-level Haar domain.Source: §1 · pp. 1–2
  • An exact segmented JVP reduces MeanFlow training memory from 15.75 to 6.16 GiB in a controlled comparison.The quality-efficiency study spans three datasets, six inference configurations, resolution scaling, and deployment-oriented performance.Source: §1 · p. 2

2 RELATED WORK

Prior JPEG restoration methods address unknown compression severity through quality prediction, coefficient reasoning, or degradation prompts, while generative restorers reuse pretrained image priors. Fast-flow and transformer work motivates an alternative that trades some semantic prior strength for a smaller direct system.

2.1 JPEG artifact removal

  • FBCNN predicts JPEG quality to condition restoration strength [9], while JDEC reasons over cosine coefficients and JPEG decoding [10] and PromptCIR learns degradation prompts [11].These methods preserve image structure well, but severe quantization leaves multiple plausible high-frequency reconstructions that deterministic regression may over-smooth.Source: §2.1 · p. 2

2.2 Generative restoration priors

  • DiffBIR and SUPIR reuse pretrained diffusion priors for restoration [2] [3], while SODiff performs one-step JPEG restoration with Stable Diffusion 2.1 and retains its UNet and frozen VAE [4].The paper contrasts that strong semantic prior with a smaller direct image-to-image system.Source: §2.2 · p. 2

2.3 Fast flows and efficient transformers

  • MeanFlow models average velocity for single-update traversal without teacher distillation [7], and its improved variant stabilizes training [8]; JiT and pixel MeanFlow motivate direct prediction [5] [6].For efficient transformers, SANA uses linear attention [16], while RELA restores local feature diversity with a depthwise-convolution branch while retaining linear token scaling [17].Source: §2.3 · p. 2

3 METHOD

The method restores JPEG images by predicting residuals in a two-level Haar wavelet representation and training a conditional model with MeanFlow. A rank-enhanced linear DiT estimates compression severity, while segmented exact JVPs reduce training memory.

3.1 Problem and wavelet representation

  • The exactly invertible two-level Haar transform encodes the image into a compact, frequency-organized grid, and the model predicts a residual while preserving observed low-frequency content on a fixed skip path.The transform maps 3 × H × W to 48 × H/4 × W/4 without a learned autoencoder.Source: §3.1 · p. 2

3.2 Conditional pixel MeanFlow

  • MeanFlow learns average velocity over an interval, allowing restoration with one or two network evaluations without teacher distillation.The network predicts a clean residual endpoint and converts it to average velocity; sampling uses one evaluation per interval.Source: §3.2 · pp. 2–3
  • Improved MeanFlow trains the model by regressing a JVP-corrected velocity onto the conditional target with adaptive weighting.The objective also applies LPIPS to decoded image estimates for t < 0.8.Source: §3.2 · p. 3

3.3 Rank-enhanced linear DiT

  • The 16-block linear DiT uses a depthwise-convolution branch to add a local, potentially full-rank path and internally predicts JPEG severity for later blocks.The severity estimate is predicted after the first three blocks and conditions the remaining 13.Source: §3.3 · p. 3

3.4 Memory-efficient exact JVP

  • Segmenting the forward pass enables exact JVP evaluation while releasing intermediate tangent graphs earlier, without changing the tangent or loss numerically.The primal path remains fully differentiable for the training loss.Source: §3.4 · p. 3

4 EXPERIMENTS

Experiments compare full-image restoration across three datasets and quality factors, then examine flow updates, resolution scaling, deployment changes, and JVP implementations.

4.1 Datasets, training, and evaluation

  • Training uses DF2K and LSDIR, one checkpoint without dataset-specific tuning, and full-image evaluation on LIVE-1, Urban100, and DIV2K-val at QF 5, 10, and 20.The final checkpoint follows small-crop training and a large-crop continuation; the evaluation QF is not given to the network, which estimates severity internally.Source: §4.1 · pp. 3–4

4.2 Main results

  • At QF 5, SODiff improves LIVE-1 LPIPS from 0.2864 to 0.2229; it also dominates DISTS, while CLIPIQA favors our model on Urban100 and DIV2K-val at every tested QF.Large pretrained restorers are generally stronger on MUSIQ and MANIQA.Source: §4.2 · p. 4
  • Our model matches SODiff on LIVE-1 at QF 10 and has the lowest LPIPS on the other QF 10/20 dataset pairs.At QF 20, LPIPS falls relative to SODiff by 25.9% on LIVE-1, 36.4% on Urban100, and 23.2% on DIV2K-val.Source: Table 1 · p. 4, §4.2 · p. 4

4.3 Flow objective ablation

  • Two-update MeanFlow gives the best LPIPS at every QF using two evaluations, versus 20–100 for RF solvers, but RF remains stronger on DISTS and several no-reference measures.The second update reduces LPIPS relative to one update by 22.3% at QF 10, 31.5% at QF 20, and 6.6% at QF 5; the main comparison uses uniform-2.Source: §4.3 · pp. 4–5

4.4 Efficiency and resolution scaling

  • The complete two-update system reaches 8.05 images/s at 1024 × 1024; its reported 4.9× throughput comparison with one-step SODiff is cross-paper, not controlled.Its 0.1242 s latency is reported as 2.2× lower than FBCNN and 4.9× lower than SODiff; these comparisons combine backbone size, representation overhead, and update count.Source: Table 3 · p. 5, §4.4 · p. 5
  • At 1440², RELA is 2.9× faster than full-attention DiT; Restormer uses 2.79 GiB versus RELA’s 0.81 GiB.The controlled one-forward-pass benchmark uses batch-one BF16 on an RTX 3090; RELA’s depthwise branch costs 5.7% throughput versus SANA linear attention at 1440².Source: Table 4 · p. 6, §4.4 · p. 6, Figure 2 · p. 6

4.5 Deployment optimizations

  • On 16 LIVE-1 images at QF 10, uniform-2 raises throughput from 1.50 to 12.61 images/s while LPIPS improves from 0.177425 to 0.160386.torch.compile adds a further 2.1× throughput gain; weight-only INT8 reduces checkpoint size from 784 MB to 109 MB but slightly reduces throughput.Source: Table 5 · p. 6, §4.5 · p. 6

4.6 JVP implementation

  • The segmented JVP matches torch.func.jvp in tangent output and loss while reducing peak memory from 15.75 to 6.16 GiB, a 60.9% reduction.Runtime is effectively unchanged at 286.31 versus 286.39 ms, so segmentation is a training-capacity optimization.Source: Table 6 · p. 7, §4.6 · p. 7

5 LIMITATIONS AND FUTURE WORK

The model’s efficiency comes with limits in semantic detail, training scale, and the comparability of efficiency measurements. The authors propose larger-scale training and call for more standardized system benchmarks.

5 LIMITATIONS AND FUTURE WORK

  • At QF 5 and on MANIQA and DISTS, the quality gap reflects limited semantic knowledge; outputs are not faithful recovery of information absent from the bitstream.The model was trained from scratch on 13,450 images and learns local structures but cannot reliably infer object-specific detail when compression removes supporting evidence.Source: §5 · p. 7
  • Training used one RTX 3090 for roughly 101k updates on a small fraction of LSDIR.The authors propose distributed training on full LSDIR with progressively increased crop resolution to test whether the remaining quality gap is data- and compute-limited or intrinsic to lacking a pretrained prior.Source: §5 · p. 7
  • Larger models and longer large-crop training may improve texture knowledge while preserving the VAE-free pipeline.A progressive-resolution schedule may improve generalization to large evaluation images.Source: §5 · p. 7
  • Cross-paper efficiency comparisons inherit software-stack differences, while the architecture study measures isolated forward passes rather than trained quality for every backbone.The authors call for end-to-end latency, energy, and peak reserved memory measurements under a single public harness across methods and accelerators.Source: §5 · p. 7

6 CONCLUSION

The conclusion presents a compact generative JPEG restorer designed for the quality-efficiency frontier. It combines a VAE-free wavelet pipeline with two-update MeanFlow inference, achieving high throughput without matching pretrained semantic priors under severe compression.

6 CONCLUSION

  • An exact two-level Haar transform removes the need for a learned VAE, while a conditional residual formulation preserves observed content and a RELA-based linear DiT combines global interaction with a local full-rank path.Clean-residual prediction parameterizes an improved MeanFlow objective for two-update inference without distillation from a pretrained generator.Source: §6 · p. 7
  • At 1024^2, the 65.32M-parameter model reaches 8.05 images/s, and its LPIPS results demonstrate high-throughput generative restoration without a full text-to-image latent pipeline.The conclusion frames the model around the quality-efficiency frontier, rather than claiming uniformly best perceptual quality.Source: §6 · p. 7
  • Large semantic priors retain an advantage under severe compression, so the model is not uniformly best in perceptual quality.Source: §6 · p. 7

A QUALITATIVE RESULTS

The qualitative results show generated predictions for 512 × 512 LSDIR crops, grouped by JPEG quality factor and inference-step count.

A QUALITATIVE RESULTS

Loading 2608.28730v1…