You Only Need 90K Parameters to Adapt Light: A Light Weight Transformer for Image Enhancement and Exposure Correction

In plain words

Photos taken in dark rooms come out too dark, and photos of bright scenes can come out washed out. Both look wrong to people, and they also make it harder for software that looks at the picture to recognise what is in it. A camera turns the raw sensor data into the final picture with a chain of colour and tone steps, and those steps are what goes wrong here.

The idea is to split the fix into two jobs. A small network adjusts the brightness of each pixel with a multiply and an add. A second, global part predicts a colour correction matrix and a gamma curve, which are the same kinds of settings a camera's image processor uses for colour and tone. A transformer-style attention block with a few learned query slots produces those global settings from the whole image.

The model is tiny, about 90 thousand parameters, and the authors report it runs in about 0.004 seconds per image. They also report that it helps standard detectors and segmenters work in poor light. All of these results are the authors' own, on public benchmarks.

TL;DR

• IAT restores sRGB images from low light and from under- or over-exposure. A local branch predicts per-pixel multiply and add maps, and a global branch predicts a 3x3 colour matrix and a gamma §3.1 §3.2 eq. 4.
• The local branch works at full resolution with depth-wise convolutions in place of self-attention, so it can take images of any size §3.2.
• The global branch uses 10 learned queries, one for gamma and nine for the colour matrix, attending to a low-resolution encoding of the image §3.2 fig. 3.
• About 0.09M parameters and about 0.004 s per image on the authors' setup, as reported in their tables tab. 1 tab. 2.
• Author-reported: LOL-V2-real PSNR 23.50 and SSIM 0.824 tab. 1; MIT-Adobe FiveK PSNR 25.32 and SSIM 0.920 tab. 2; exposure correction best averages on every metric the authors report tab. 3.
• Caveats: in the authors' own LOL-V1 table, MAXIM scores higher on PSNR and SSIM tab. 1. The released code also differs from the paper in several training details (see Paper vs code).

Problem

Real-world light is hard: low-light images carry limited photon counts and sensor noise, while sunlit scenes saturate the sensor, and shadows can make the contrast ratio 1000:1 or higher within one image §1 fig. 1. Such images degrade both human viewing and downstream tasks such as detection and segmentation §1.

Earlier learned methods often work on raw sensor data, which is not available in everyday use. Methods that work directly on sRGB images are typically large U-Net or transformer models. The paper cites 115.63M parameters for IPT and 14.14M for MAXIM as examples §2 tab. 1. The authors aim for a model small enough for mobile and edge devices §2.

Method

Formulation. The paper models the camera pipeline as a rendering function G applied after an unknown reconstruction F that maps the input back to raw data §3.1 eq. 1. It replaces this with two parts: a pixel-wise local component f, and a global ISP component g with a simplified form §3.1 eq. 2 eq. 3.

Local branch. A 3x3 convolution expands the RGB input to 16 channels §3.2. Two independent branches, one for the multiply map M and one for the add map A, each stack three Pixel-wise Enhancement Modules (PEM) at full resolution §3.2 §4 fig. 3. Each PEM encodes position with a 3x3 depth-wise convolution, then applies a pointwise, depth-wise, pointwise sequence and two 1x1 convolutions. It uses a light normalisation with a learnable per-channel scale and bias, plus a learnable channel-mixing matrix initialised to the identity, in place of layer normalisation §3.2 fig. 3. Layer Scale multiplies each block's output by a small constant. The paper does not give that constant's value §3.2. Each branch adds a residual skip from its input. A final 3x3 convolution gives M with a ReLU and A with a Tanh §3.2.

Global branch. Two convolutions act as a lightweight encoder that lowers the resolution and raises the channel count §3.2. A Global Prediction Module then uses learnable queries Q that attend to keys and values from the encoded features. The keys and values get a positional encoding from a depth-wise convolution, so the module handles different input sizes §3.2 fig. 3. Unlike DETR, the queries have no self-attention among themselves, and they start at zeros in the paper §3.2. A feed-forward network follows §3.2. Ten queries are used: nine give the entries of the 3x3 colour matrix W, and one gives gamma §3.2. The output is initialised so that W is the identity and gamma is one §3.2.

Combined model. The local maps are applied first, as an affine per-pixel change, then the colour matrix, then the gamma curve. A small epsilon keeps the base of the power non-negative §3.2 eq. 4. The paper's supplement gives the simplified ISP as a white-balance matrix, then a colour matrix, then gamma §7 eq. 5.

Training. The paper trains on a single RTX 3090 with batch size 8, the Adam optimiser, initial learning rate 2e-4 and weight decay 1e-4, with a cosine schedule and horizontal and vertical flips §4.1. The loss is L1 for MIT-Adobe FiveK and exposure correction §4.1 §4.2. For LOL, the paper uses smooth L1 plus a VGG perceptual loss §4.1. An optional raw-RGB supervision term adds an L1 loss on an intermediate representation, with weight 0.1 §7 eq. 7. It needs raw data synthesised with an invertible ISP §7.

High-level tasks. IAT can be placed in front of a detector or segmenter. The paper trains the whole pipeline end to end, minimising the downstream task loss through IAT §4.3 §8 eq. 8.

# IAT forward pass  [§3.2] [eq. 4]
# Input:  low-light or mis-exposed sRGB image in [0, 1]  (B, 3, H, W)
# Output: corrected sRGB image  (B, 3, H, W)

def iat_forward(image):
    # Expand to 16 channels, full resolution  [§3.2]
    # [code: IAT_main.py#L46-L47]
    local_features = leaky_relu(conv3x3(image))  # (B, 16, H, W)
    # Two independent stacks of PEM blocks, each with a residual  [§3.2] [§4]
    multiply_map = relu(stack_m(local_features) + local_features)
    add_map = tanh(stack_a(local_features) + local_features)
    # Global colour matrix and gamma from the whole image  [§3.2]
    gamma, colour_matrix = global_prediction(image)  # (B,), (B, 3, 3)
    # Pixel-wise affine adjustment  [§3.2] [eq. 4]
    adjusted = image * multiply_map + add_map
    # Apply the colour matrix to each pixel  [eq. 4]
    coloured = apply_colour(colour_matrix, adjusted)
    # Floor the base at epsilon before the power  [§3.2] [eq. 3]
    # The code also clamps to 1.0; the paper does not  [code: IAT_main.py#L109]
    base = clamp(coloured, 1e-8, 1.0)
    # Gamma correction  [eq. 4]
    return base ** gamma

# Global prediction module  [§3.2] [fig. 3]
# Input:  image (B, 3, H, W)
# Output: gamma (B,), colour_matrix (B, 3, 3)

def global_prediction(image):
    # Lightweight encoder at lower resolution  [§3.2]
    # [code: global_net.py#L64-L80]
    features = encoder(image)  # (B, C, H/4, W/4)
    # Ten learned queries: index 0 for gamma, 1 to 9 for the matrix  [§3.2]
    # Paper initialises them to zeros; code uses ones  [code: global_net.py#L17]
    queries = learned_queries  # (B, 10, C)
    # Keys and values: features plus depth-wise positional term  [§3.2]
    keys_values = layer_norm(features + positional(features))  # (B, HW, C)
    # Cross-attention only: no query-to-query step, no residual  [§3.2]
    # [code: global_net.py#L24-L37, #L56-L59]
    queries = attend(queries, keys_values)
    # Feed-forward block with residual  [§3.2] [code: global_net.py#L60]
    queries = queries + feed_forward(layer_norm(queries))
    # Gamma starts at one  [§3.2] [code: global_net.py#L86-L89, #L118]
    gamma = linear_gamma(queries[:, 0]) + 1.0
    # Matrix starts at the identity  [§3.2] [code: global_net.py#L90, #L120]
    offsets = linear_matrix(queries[:, 1:])  # (B, 9)
    colour_matrix = reshape(offsets, (B, 3, 3)) + identity_3x3
    return gamma, colour_matrix

Key equations:

I_t = G(F(I_i)) eq. 1

I_t = g_t(f(I_i)) eq. 2

g_t(\cdot) = \big(\max\big(\textstyle\sum_{c_j} W_{c_i,c_j}(\cdot),\ \epsilon\big)\big)^{\gamma},\quad c_i, c_j \in \{r,g,b\} eq. 3

I_t = \Big(\max\Big(\textstyle\sum_{c_j} W_{c_i,c_j}\big(I_i \odot M + A\big),\ 0\Big)\Big)^{\gamma} eq. 4

G(\cdot) = \mathrm{Gamma}\big(W_{ccm}(W_{wb}(\cdot))\big) eq. 5, supplement

Here W is the product of the white-balance and colour-matrix transforms, and γ is the gamma exponent. Equation 3 is restated in the supplement as eq. 6.

L_{total} = L_{rgb} + \lambda\, L_{raw},\quad L_{rgb} = L_1\big(g_t(f(I_i)), I_t\big),\quad L_{raw} = L_1\big(f(I_i), I_{raw}\big),\quad \lambda = 0.1 eq. 7, supplement

\min_{i \in \mathbb{I},\, d \in \mathbb{D}} L_{obj}(\hat{t}, t),\qquad \hat{t} = \mathbb{D}\big(\mathbb{I}(I_{low})\big) eq. 8, joint training

Equation 8 is the paper's detector objective with the input passed through IAT, written here in words: the detector loss is minimised over both the IAT parameters and the detector parameters.

Hyperparameters

Results

Author-reported. Hardware for timings is not stated in the paper. Times for baselines come from each method's own code platform, marked M (Matlab), T (TensorFlow) or P (PyTorch) in tab. 1.

Low-level enhancement, LOL tab. 1. Values are PSNR and SSIM on LOL-V1 and LOL-V2-real, plus efficiency. Best values are red and second-best are blue in the source.

The source table lists the KIND time as "0e38 (T)", which is most likely a typo for 0.038, but the paper does not say so. Treat that cell as unreadable.

Low-level enhancement, MIT-Adobe FiveK tab. 2. Expert C ground truth §4.1.

Exposure correction tab. 3. Test set has five experts' targets. Per-expert PSNR and SSIM are in the source; the averages and PI are shown here. PI is lower-is-better.

IAT per-expert PSNR and SSIM, from the same table: Expert A 19.90 and 0.817; B 21.65 and 0.867; C 21.23 and 0.850; D 19.86 and 0.844; E 19.34 and 0.840. The paper's text also cites about 7M parameters for MSEC and 0.5 s per image §4.2. Neither number appears in the table.

High-level tasks tab. 4. Detector and segmenter are trained on enhanced data. EXDark mAP is for YOLOv3, ACDC mIoU is for DeepLabV3+, and TYOL mAP is for YOLOv3. Time is per image in seconds.

Joint training tab. 6. The table compares pre-enhancement with joint training, and the initial weights used for IAT.

Ablations. Results on LOL-V2-real (PSNR, SSIM) and EXDark (mAP) tab. 7:

Number of PEM blocks for M and A, PSNR on LOL-V2-real tab. 8:

Channel and block trade-off on LOL-V2-real tab. 9: 12 channels with 4 blocks gives 22.60 PSNR, 0.807 SSIM and 86.22K parameters. 24 channels with 2 blocks gives 22.70, 0.815 and 101.03K. The chosen 16 channels with 3 blocks gives 23.50, 0.824 and 91.15K.

Raw-RGB supervision tab. 5. With raw supervision, results change as follows, each as PSNR and SSIM per expert §7. Expert A 19.98 and 0.822 (without: 19.90 and 0.817); B 22.03 and 0.885 (21.65 and 0.867); C 21.16 and 0.843 (21.23 and 0.850); D 19.94 and 0.852 (19.86 and 0.844); E 19.48 and 0.841 (19.34 and 0.840).

Paper vs code

Code: cuiziteng/Illumination-Adaptive-Transformer at c76472265247f47cea57649af28b15018bb64cb1 (Apache-2.0, copyright University of Tokyo). The repository is linked from the paper's abstract §1.

1. Query initialisation. The paper initialises global queries to zeros §3.2. The code initialises them to ones global_net.py.
2. Colour clamp. The paper's Eq. 4 floors the base at 0 and the text states an epsilon of 1e-8 eq. 4 §3.2. The code clamps the base to [1e-8, 1.0] IAT_main.py. The upper clamp at 1.0 is not in the paper.
3. Middle depth-wise kernel. The paper describes a 3x3 depth-wise position encoding followed by pointwise, depth-wise, pointwise §3.2. The code's middle depth-wise convolution is 5x5 blocks.py.
4. Local-branch activation. The code applies LeakyReLU with slope 0.2 after the first convolution IAT_main.py. The paper does not mention it.
5. Global encoder. The paper says two convolutions §3.2. The code uses two stride-2 convolutions with batch normalisation and GELU global_net.py.
6. Exposure crops. The paper trains on 512x512 crops and resizes test images to a maximum dimension of 512 §4.2. The code resizes all exposure images to 1200x900 exposure.py, crops training images to half the resized height and width exposure.py, and transposes portrait images to landscape first exposure.py.
7. LOL-V1 augmentation. The paper uses horizontal and vertical flips §4.1. The code takes a 256x256 random crop and picks one of nine variants, including rotations and flips. Two of the nine variants leave the crop unchanged lol_v1_new.py.
8. LOL loss. The paper uses smooth L1 plus VGG for LOL §4.1. The LOL-V1 scripts use plain L1 train_lol_v1_patch.py. The LOL-V2 script uses smooth L1 plus 0.04 times a VGG16 feature loss train_lol_v2.py.
9. Weight decay. The paper states 1e-4 §4.1. The code uses 4e-4 for LOL-V1 patch training train_lol_v1_patch.py, 5e-4 for LOL-V2 train_lol_v2.py, 1e-4 for LOL-V1 fine-tuning train_lol_v1_whole.py, and 2e-4 for exposure train_exposure.py.
10. Fine-tuning length. The paper says 100 epochs of fine-tuning on 600x400 images §4.1. The code's default is 400 epochs train_lol_v1_whole.py.
11. LR schedule step. The scripts create CosineAnnealingLR(T_max=num_epochs) train_lol_v1_patch.py and call scheduler.step() inside the per-batch loop train_lol_v1_patch.py. The cosine therefore runs over iterations, not epochs. The paper says only "cosine" §4.1.
12. Exposure gamma. The code builds the model with type='exp', which freezes the global gamma base train_exposure.py. The paper does not say so. A commented-out alternative gives the global branch a 0.1x learning rate train_exposure.py.
13. Raw-RGB supervision. The paper describes a raw loss using invertible-ISP data §7. No raw loss appears in the released training scripts, and the invertible-ISP weights are not in the repository.
14. Detection weight decay. The paper states 1e-4 §4.3. The EXDark YOLOv3 config uses 5e-4 yolov3_IAT_lol.py.
15. Segmentation schedule. The paper states 20,000 iterations, batch 8, lr 0.05 §4.3. The released configs use lr 0.005 [code: IAT_high/IAT_mmseg/configs/base/schedules/schedule_40k.py#L2] or 0.01 [schedule_160k.py#L2], and crops of 512x512 [code: configs/base/datasets/ACDC_night.py#L7] or 512x1024 with histogram equalisation ACDC_night_with_HE.py. Which setting produced tab. 4 was not checked.

Pitfalls

• Claims versus the authors' own tables. The text says IAT achieves SOTA "on both image quality and model efficiency" on LOL §4.1. In tab. 1, MAXIM scores higher on LOL-V1 PSNR (23.43 vs 23.38) and SSIM (0.863 vs 0.809). IAT's LOL-V2-real SSIM (0.824) is above MAXIM's (0.818).
• Average does not match. The average of the five IAT per-expert PSNR values is 20.40 (101.98 / 5), not the reported 20.34 tab. 3. The SSIM average (0.844) matches.
• Final ablation row. The full model's EXDark mAP is 77.1 in the ablation tab. 7 and 77.2 in the main table tab. 4.
• Block-count claim. The text says the same PEM count for M and A helps §9.2. In tab. 8, M=2 with A=3 (22.85) beats M=2 with A=2 (22.10). The 3-3 setting (23.50) and 4-4 (23.48) are close, and parameter counts for the grid are not given.
• Speed claim on EXDark. The text says IAT is best on accuracy and speed §4.3. The unenhanced baseline is faster (0.033 s vs 0.040 s) tab. 4. The speed lead is only over the enhancement methods.
• ACDC. IAT's ACDC mIoU (62.1) is below the unenhanced baseline (63.3) tab. 4. The paper's fix is joint training, and the gain depends on initialisation: LOL-initialised joint training gives 63.8, random init gives 61.5 tab. 6.
• Timing and FLOPs context. The paper does not state the hardware for the 0.004 s timing, the input size used for FLOPs, or whether timings include pre- or post-processing. Baseline times were measured on each method's own platform tab. 1.
• Supplement numbering. The source is one LaTeX file with the supplement after the references. Section, equation and table numbers beyond the main text (§7 onward, eqs. 5 to 8, tab. 5 onward) follow the source order. They were not checked against a compiled PDF; cite by title if the numbers differ.
• Numerically sensitive step. Gamma is applied to a base floored at 1e-8, and the code also clamps to 1.0. Without the floor, the power is undefined or unstable at zero. Without the upper clamp, values above one would be raised to a power that can push them beyond the display range.
• Dataset-size and augmentation details. The exposure training crop differs between paper and code (see Paper vs code 6), so results from either setup may not be directly comparable.
• Model file size. The "417 KB" figure refers to the main IAT structure without the detector §8. The paper's comparison is against YOLO-V3 at 237 MB.

Lineage

• Builds on: the convolution-based transformer blocks in the code are copied from UniFormer blocks.py. The paper draws on DETR-style learned queries §3.1, CaiT's Layer Scale §3.2 and ResMLP's normalisation, which is used as an ablation baseline §9.1.
• Compared against: low-level methods MAXIM, IPT, Uformer, 3D-LUT, DRBN, KIND, MBLLEN, RetinexNet, RCT, Zero-DCE and Deep-UPE, and for exposure correction MSEC and DPED tab. 1 tab. 2 tab. 3. Detection and segmentation baselines use the unenhanced data, MBLLEN, DeepLPF and Zero-DCE tab. 4.