TTTIR: Unlocking Instance-Specific State Evolution via Test-Time Training for Image Restoration

Kaihang Zheng^*, Jun Li^*, Hang Guo, Hongyu Chi, Zimo Liu, Tao Dai, Jinpeng Wang^†, and Yaowei Wang^†

⭐ If TTTIR is helpful to your projects, please help star this repo. Thanks! 🤗

We sincerely invite readers to refer to our team’s other work MambaIR.

News

• 2026-08-05: The initial codebase is released. We are excited to release the first general image reconstruction model based on TTT. Our model achieves SOTA performance with fewer than 1M parameters.
• 2026-08-24: Our paper is still on hold at arXiv. We’ll update the paper link once the arXiv ID is available. 🌹
• 2026-08-26: For easier reproducibility, we are working to unify the codebase across all tasks.

Abstract: Image restoration is inherently challenging due to the diverse and highly input-dependent nature of real-world degradations. While recent architectures such as Transformers and state-space models have advanced the field, they predominantly rely on static, globally shared parameters, which struggle to fully accommodate instance-specific degradation patterns. To address this limitation, we propose TTTIR, a framework that reformulates image restoration as an instance-specific state evolution process. Progressive State Sequence Generation (PSSG) constructs complementary spatial-frequency target states that define what to recover, while State Transition Evolution (STE) adapts lightweight transition operators through a restoration-oriented test-time training inner loop that determines how features should evolve. Extensive experiments demonstrate strong performance across low-light enhancement, rain removal, and image dehazing benchmarks with favorable computational scalability.

Contents

1. Dependencies
2. Datasets
3. Training
4. Testing
5. Model Zoo
6. Results

Dependencies

Our experiments in the paper are conducted on NVIDIA RTX 3080 Ti GPUs.

git clone https://github.com/Elysiaaaaaaaa/TTTIR.git
cd TTTIR

conda create -n tttir python=3.7.12 -y
conda activate tttir

# The dependency file will be released soon.
pip install -r requirement.txt

Datasets

The datasets used for training and evaluation are listed below. They can be downloaded from Baidu Netdisk (extraction code: 7ghe).

Task	Training set	Testing set
Low-light enhancement	LOL-v1	LOL-v1, LOL-v2-Real
Low-light enhancement	LOL-v2-Synthetic	LOL-v2-Synthetic
Rain streak removal	Rain13K	Test100, Rain100H, Rain100L, Test1200, Test2800
Raindrop removal	Raindrop training set	Raindrop-A, Raindrop-B
Image dehazing	RESIDE-6K	RESIDE-6K
Image dehazing	Haze4K	Haze4K

Organize the datasets under a common root directory:

Click to expand the dataset directory structure

<dataset_root>/
|-- LOLv1/
|   |-- Train/
|   |   |-- input/
|   |   `-- target/
|   `-- Test/
|       |-- input/
|       `-- target/
|-- LOLv2/
|   |-- Real_captured/
|   |   |-- Train/
|   |   |   |-- Low/
|   |   |   `-- Normal/
|   |   `-- Test/
|   |       |-- Low/
|   |       `-- Normal/
|   `-- Synthetic/
|       |-- Train/
|       |   |-- Low/
|       |   `-- Normal/
|       `-- Test/
|           |-- low1/
|           `-- Normal/
|-- Rain13k/
|   |-- train/
|   |   |-- input/
|   |   `-- target/
|   `-- test/
|       |-- Rain100H/
|       |   |-- input/
|       |   `-- target/
|       |-- Rain100L/
|       |   |-- input/
|       |   `-- target/
|       |-- Test100/
|       |   |-- input/
|       |   `-- target/
|       |-- Test1200/
|       |   |-- input/
|       |   `-- target/
|       `-- Test2800/
|           |-- input/
|           `-- target/
|-- raindrop/
|   |-- train/
|   |   |-- data/
|   |   `-- gt/
|   |-- test_a/
|   |   |-- data/
|   |   `-- gt/
|   `-- test_b/
|       |-- data/
|       `-- gt/
|-- RESIDE-6K/
|   |-- train/
|   |   |-- haze/
|   |   `-- gt/
|   `-- test/
|       |-- haze/
|       `-- gt/
`-- Haze4K/
    |-- train/
    |   |-- haze/
    |   |-- gt/
    |   `-- trans/
    `-- test/
        |-- haze/
        |-- gt/
        `-- trans/

Training

Rain streak and raindrop removal

cd derain
CUDA_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node 2 --use_env --master_port 6198 main.py \
    --model_name Rain13k --mode train --num_epoch 300 --data_dir /data2/zhengkaihang/ttt/dataset/Rain13k \
    --learning_rate 1e-3 --save_freq 30 --valid_freq 1 --batch_size 4 --num_worker 6

Low-light enhancement

Training settings and dataset paths are defined in YAML files under lowlight/options/train/.

cd lowlight

# LOL-v1
python train.py -opt options/train/LOL-v1.yml

# LOL-v2-Synthetic
python train.py -opt options/train/LOL-Syn.yml

Testing

Rain streak and raindrop removal

cd derain

CUDA_VISIBLE_DEVICES=0 python main.py --mode test --data_dir /data2/zhengkaihang/ttt/dataset/Rain13k \
    --test_model /path/to/model.pkl --model_name Rain100H

Low-light enhancement

cd lowlight

# LOL-v1
python test.py -opt options/test/LOL-v1.yml

# LOL-v2-Real evaluated with the LOL-v1 model
python test.py -opt options/test/LOL-v2-Real-Based_v1.yml

# LOL-v2-Synthetic
python test.py -opt options/test/LOL-v2-Syn.yml

Model Zoo

Task	Training set	Evaluation set(s)	Checkpoint
Low-light enhancement	LOL-v1	LOL-v1, LOL-v2-Real	Models & results (code: a4r4)
Low-light enhancement	LOL-v2-Synthetic	LOL-v2-Synthetic	Models & results (code: a4r4)
Rain streak removal	Rain13K	Test100, Rain100H, Rain100L, Test1200, Test2800	Models & results (code: 8jf2)
Raindrop removal	Raindrop training set	Raindrop-A, Raindrop-B	Models & results (code: hs0z)
Image dehazing	RESIDE-6K	RESIDE-6K	Coming soon
Image dehazing	Haze4K	Haze4K	Coming soon

Results

Click each section below to view the quantitative and qualitative comparisons.

Low-light enhancement

Rain streak removal

Raindrop removal	Image dehazing
Table 3. Raindrop-A and Raindrop-B.	Table 4. RESIDE-6K and Haze4K.

Low-light enhancement

Rain removal

Citation

If you find our code useful or use the toolkit in your work, please consider citing:

@article{zheng2026tttir,
  title   = {TTTIR: Unlocking Instance-Specific State Evolution via Test-Time Training for Image Restoration},
  author  = {Zheng, Kaihang and Li, Jun and Guo, Hang and Chi, Hongyu and Liu, Zimo and Dai, Tao and Wang, Jinpeng and Wang, Yaowei},
  journal = {arXiv preprint},
  year    = {2026}
}