Efficient Diffusion as Low Light Enhancer

In plain words

Diffusion models can make very good enhanced photos from dark images, but they normally refine the picture over many small steps, and that is slow. This paper trains a fast "student" to copy the path of a slow "teacher" in just a few steps.

Copying the path alone loses quality when there are few steps. The authors name two causes: small errors in the teacher's predictions, and a mismatch between how the model was trained (on pure random noise) and what a dark photo needs. Their fix is to nudge the teacher's path toward a rough estimate of the clean picture, built from the image's reflectance, and to add two direct checks against the true bright photo. The authors report that the resulting models, with 2, 4 or 8 steps, match or beat earlier diffusion methods on several low-light benchmarks. The tables below give the exact numbers.

TL;DR

• Distils a multi-step diffusion low-light enhancer into 2, 4 or 8 sampling steps by trajectory matching against a teacher §3.1.
• Its RATR module shifts the teacher's end point toward a residual anchor built from an illumination map and a denoised noise estimate §3.2–§3.3.
• Adds a pixel L2 term and a perceptual term against ground truth §3.4.
• Abstract claim: 2 steps reach performance comparable to earlier diffusion methods that use more steps; 8 and 4 steps set new results [abstract].
• Author-reported: best PSNR on LOLv1, LOLv2-real and LOLv2-synthetic is 28.090, 31.250 and 30.166 §4.2, tab. 1.
• The teacher checkpoint is not released in the code repository, so the pipeline cannot be re-run end to end from the public material.

Problem

Diffusion-based low-light enhancement needs many denoising steps, which conflicts with real-time and edge use §1. Training-free accelerators and consistency-style distillation both lose quality as the step count falls §1. The authors attribute the loss to (1) fitting errors in the teacher's predicted score, which add undesired terms to the distillation target, and (2) an "inference gap" between the Gaussian noise the diffusion model was trained on and the more deterministic input a low-light enhancer needs §1, §3.2.

Method

The student and teacher share the conditioning on the low-light image y. The forward and reverse diffusion equations are the standard ones §3, eq. 1–2.

1. Trajectory decoder. The student's score ε_θ defines a deterministic jump from level t to level s < t §3.1, eq. 3.
2. Teacher trajectory through an intermediate level. The teacher's path from t to s is taken through a level u in [s, t), giving a second-order estimate §3.1, eq. 4.
3. Distillation target. The student's jump is matched to the teacher's jump, first as a noise-level match §3.1, eq. 5, then in clean-image space, where the target and estimate are the predicted clean images §3.1, eq. 6. The paper weights this loss by λ(t) = max(1, α_t²/σ_t²) §3.1, eq. 7.
4. Fitting-error refinement. If the teacher were ideal, the two teacher score terms would be equal §3.2, eq. 8–9. The authors keep the term that would otherwise conflict and replace it with a linear interpolation between the teacher's score at t and an ideal score, controlled by ω in (0, 1] §3.2, eq. 10. The resulting distillation loss is eq. 11 §3.2.
5. Residual shift (inference gap). The "ideal" noise is redefined as the offset of x_t from a chosen anchor x̃_0 at level t, so the diffusion starts from a point between the low-light and clean spaces §3.2, eq. 12. Corollary 1 states that the refined teacher trajectory equals ω · x^η_{s,u,t} + (1 − ω) · x̃_s, where x̃_s = α_s x̃_0 + σ_s ε_η §3.2, Corollary 1. The proof is deferred to the supplementary material.
6. RATR anchor. Choose x̃_0 from the reflectance rather than from the clean image, which would collapse the residual back to plain noise, or from the low-light input §3.3. The illumination is the channel-wise maximum of y, and the noise is z′ = |y − ψ(y)| with ψ a non-learned denoiser (not named). Then x̃_0 = (y − z′) / h′ §3.3, eq. 13.
7. Auxiliary losses. A pixel L2 term λ_pix‖x_0 − x_est‖² §3.4, eq. 14 and a perceptual term λ_per‖φ(x_0) − φ(x_est)‖² §3.4, eq. 15. Total: L_distill + L_per + L_pix §3.4, eq. 16. The weights λ_pix and λ_per and the feature extractor φ are not stated.

Figure 2 shows qualitative results for the method section fig. 2. No step-by-step inference sampler is described in the paper body.

# One distillation step, ReDDiT  [§3.1–§3.4]
# Input:  low-light y, ground truth x0, student and teacher networks
# Output: scalar loss for the student parameters

def distill_loss(y, x0, student, teacher, s, u, t):
    # Forward-diffuse the clean image to level t  [eq. 1]
    noise = random_normal(x0.shape)
    x_t = alpha_t * x0 + sigma_t * noise
    # Teacher's first jump, t to u  [eq. 3, teacher network]
    eps_t = teacher(x_t, y, t)
    x_u = (alpha_u / alpha_t) * x_t + (sigma_u - alpha_u * sigma_t / alpha_t) * eps_t
    eps_u = teacher(x_u, y, u)
    # Teacher's second-order trajectory from t to s via u  [eq. 4]
    x_teacher = (alpha_s / alpha_t) * x_t \
        + sigma_s * eps_u \
        + (alpha_s / alpha_u) * sigma_u * (eps_t - eps_u) \
        - (alpha_s / alpha_t) * sigma_t * eps_t
    # Residual anchor from the low-light input  [§3.3, eq. 13]
    x_tilde_0 = reflectance_estimate(y)
    # Pull the teacher end point toward the anchor  [§3.2, Cor. 1]
    x_tilde_s = alpha_s * x_tilde_0 + sigma_s * eps_u
    x_refined = omega * x_teacher + (1 - omega) * x_tilde_s
    # Student's jump from t to s  [eq. 3, student network]
    eps_student = student(x_t, y, t)
    x_student = (alpha_s / alpha_t) * x_t \
        + (sigma_s - alpha_s * sigma_t / alpha_t) * eps_student
    # Clean-image targets and estimates  [eq. 6]
    x_target = clean_from_level(x_refined, x_t, s, t)
    x_est = clean_from_level(x_student, x_t, s, t)
    # Weighted trajectory loss  [eq. 7]
    weight = max(1.0, alpha_t**2 / sigma_t**2)
    loss_distill = weight * squared_l2(x_target, x_est)
    # Pixel and perceptual terms against ground truth  [§3.4, eq. 14–15]
    loss_pix = lambda_pix * squared_l2(x0, x_est)
    loss_per = lambda_per * perceptual(x0, x_est)  # phi not stated
    return loss_distill + loss_per + loss_pix  # [eq. 16]


def clean_from_level(x_s, x_t, s, t):
    # Invert the level-s jump to a clean-image estimate  [eq. 6]
    ratio = sigma_s / sigma_t
    return (x_s - ratio * x_t) / (alpha_s - ratio * alpha_t)


def reflectance_estimate(y):
    # Illumination is the channel-wise maximum of the input  [§3.3]
    illumination = max_over_channels(y)  # (B, 1, H, W)
    # Noise is the distance to a non-learned denoised copy  [§3.3, eq. 13]
    noise = abs(y - denoise(y))  # denoise (psi) not named
    return (y - noise) / illumination

Hyperparameters

Values in the paper body are listed first. The code column gives the values the shipped config and scripts use. Where the paper and code disagree, see "Paper vs code".

The [src-only: ...] rows come from a block inside \begin{comment} in the authors' source. It is not compiled into the paper, so these values are treated as not stated in the paper and shown only for implementers.

Results

Author-reported. Setup: LOLv1 and LOLv2 (real and synthetic) paired test sets, PSNR, SSIM and LPIPS, against methods with the number of function evaluations (NFE) shown §4.2, tab. 1. Hardware and model size are not stated in the body.

Paired full-reference results on SID and SDSD tab. 2:

No-reference NIQE on DICM, LIME, MEF, NPE and VV, with the LOLv2-synthetic model applied without retraining §4.2, tab. 3. The paper does not state the step count for this row. This paper: 3.62 (DICM), 3.45 (LIME), 3.93 (MEF), 3.24 (NPE) and 3.00 (VV). Lower is better tab. 3.

Efficiency §4.5, tab. 4: teacher with 16 steps 0.611 s and 1.64 FPS, 17.43 M parameters. This paper: 8 steps 0.323 s, 3.09 FPS; 4 steps 0.162 s, 6.15 FPS; 2 steps 0.076 s, 13.1 FPS. All with 17.43 M parameters. PyDiff uses 97.89 M parameters and WCDM 22.08 M tab. 4. The hardware is not stated.

Ablations §4.4: removing RATR causes the largest drop, and the drop grows as steps fall fig. 6. On LOLv2-real, the 4-step model with ω = 0.8 is reported to exceed the 8-step model fig. 7. Using the reflectance as the residual anchor beats clean images, the low-light input, and a PairLIE-based anchor fig. 5. Removing either auxiliary loss costs less than removing RATR fig. 6.

Paper vs code

The abstract links a project page (https://mqleet.github.io/ReDDiT_Project/). That page links github.com/lgz-0713/ReDDiT, which we cloned and pinned at commit ba046491b7135850ddf74900ad2ff691d57b8331. The repository has no LICENSE file and no license statement, so we cite it only and copy no code. Its README acknowledges building on SR3 README.md.

• Teacher model not released. The README's TODO list has "Teacher Model" unchecked README.md. Both teacher and student are loaded from one checkpoint model.py. The paper's pretraining stage (2M iterations in the source-only text) has no code here.
• Loss shape. The paper sums a norm of the ε difference and a norm of the ε-to-ω-mix difference §3.2, eq. 11. The code takes the per-sample maximum of an x-space MSE and an ε-space MSE, then means it [code: model/ddpm_modules/diffusion.py#L393-L413]. Eq. 7's λ(t) is absent from the code (no λ-term anywhere in the loss) diffusion.py.
• Extra term not in paper. The code adds an MSE between the true forward noise and the student's ε, weighted by w_gt diffusion.py.
• Structure term. The code has a w_str hyper-parameter, but the term it feeds is hard-coded to zero, so w_str has no effect diffusion.py.
• Shipped weights differ. train_lol2_real.sh and train_lol2_syn.sh set w_gt = 0.0, which removes the pixel and noise terms, with w_snr = 0.4 and w_lpips = 0.6 train_lol2_real.sh train_lol2_syn.sh. train_lol1.sh is labelled LOLv1 but runs the LOLv2-real config and dataset, with w_gt = 1.0 and w_snr = 0.8 train_lol1.sh. The JSON default is w_gt 0.1, w_snr 0.5, w_str 0.1, w_lpips 0.2 lolv2_real_train.json.
• ω and w_snr are not the same knob. The code computes (1 − w_snr) · z_{t−2} + w_snr · refined point diffusion.py. Corollary 1 has ω multiplying the teacher point. Under that reading ω = 1 − w_snr, so the shipped w_snr = 0.4 corresponds to ω = 0.6. This is our reading, and the authors do not state the mapping.
• Refinement formula. The code normalises the input before subtracting noise, and it divides by (max channel + 0.1) [code: model/ddpm_modules/diffusion.py#L377-L386]. The paper divides the difference by the illumination only once, with no offset §3.3, eq. 13.
• Noise level sampling. The code draws one integer t for the whole batch, then draws a continuous noise level uniformly within each grid interval per sample diffusion.py. The paper does not describe this sampling.
• Stage loop and checkpoints. Each stage (16 → 8 → 4 → 2 → 1 teacher steps) restarts from the same checkpoint unless --progressive is passed, and the shipped scripts do not pass it train.py train.py. The loop runs range(len(PD_steps)) but reads PD_steps[i+1], so the last pass indexes past the list train.py. On our reading the run stops after the 1-step stage. We have not run it.
• Test-time sampler. The student is sampled with a deterministic DDIM-style update from Gaussian noise, clipping x0 to [−1, 1] at each step. The starting noise is drawn fresh with randn diffusion.py. The paper does not state the sampler.
• Brightness correction at evaluation. test.py and the training-time validation rescale the output so its mean gray value matches the ground truth test.py train.py. The paper body does not mention this.
• Model selection on the test split. The validation loader reads LOL's eval15 split LoL_dataset.py, and the checkpoint is saved whenever validation PSNR beats the best so far train.py. Reported numbers may therefore be selected on the same split.

Scoped check: the code's two-step teacher rollout equals eq. 4 when expanded, and the x-target construction matches eq. 6 diffusion.py.

Pitfalls

• Teacher is not in the repo. Without the teacher checkpoint you cannot reproduce the distillation. The code can only be run from a teacher you train yourself with the SR3-style setup the supplement describes (source-only).
• EMA is applied from the first step. The config key step_start_ema (5000) is never read, and the EMA update runs on every step model.py ema.py. After 5,000 updates at decay 0.9999 about 61 % of the EMA weight is still on the initial checkpoint (0.9999^5000 ≈ 0.607). The validated and saved models are EMA models model.py.
• Timestep wrap at the first even index. At t = 2 the lower bound for x_{t−2} uses index (t − 3)·stride, which is negative. Python wraps it to the end of the array, so the sampled noise level is drawn from an unrelated range diffusion.py. This happens with probability 1/N per iteration in a stage with N student steps, and every iteration in the 1-step stage. Our reading. We have not run it.
• Refinement normalisation may divide by zero. The loader scales inputs to [−1, 1] LoL_dataset.py LoL_dataset.py. The code then divides by (max channel + 0.1). That sum is zero where the max channel is 0.45 in [0, 1] units, so the ratio blows up there. Our reading. Check the refined anchor's range before using the code.
• ω versus w_snr. The mapping above is inferred, so the paper's ω = 0.8 result cannot be reproduced by setting w_snr = 0.8 without checking.
• Anchor choice in x̃_s. Corollary 1 writes x̃_s with ε_η, but the paper does not say which teacher evaluation; the code uses the teacher's noise at the intermediate level u.
• λ_pix, λ_per and φ are missing. Without them the auxiliary losses are not reproducible from the paper. The code uses LPIPS-VGG and weight 0.6 for perceptual, and 1.0 for pixel in the LOLv1-labelled script only.
• Text and table disagree on SID. The text says the SID PSNR is 25.32 dB §4.2. Table 2 shows 25.85 for 4 steps and 25.88 for 2 steps tab. 2.
• Not every SSIM is best. PyDiff at 0.930 on LOLv1 and GSAD at 0.944 on LOLv2-synthetic exceed this paper's SSIM values tab. 1.
• Step count for NIQE is missing. tab. 3 does not say which NFE is used.
• Float numbers are from source order. No compiled PDF was available, so figure and table numbers follow the order of the source. Check them against the PDF before quoting.
• Source-only text. The supplement (sec/X_suppl.tex) is not \input by main.tex. It says the teacher is trained for 2M iterations, and that the experiments ran on a single A100 with 80 GB. It also states the proof of Corollary 1. We do not treat this as part of the paper.
• Unlicensed code. No LICENSE file. Read it, cite it, and do not copy it.

Lineage

• Builds on: diffusion models §1, progressive and consistency-style distillation §1 §2, and CTM, which extends consistency to whole trajectories §2.
• Compared against: PyDiff (4 steps), WCDM and GSAD (10 steps) tab. 1, Retinexformer, LLFlow and other non-diffusion enhancers tab. 1, and accelerators such as DDIM and consistency distillation fig. 4.
• Superseded by: not known from this wiki.