# TaskTok: Delving into Task Tokens for Task-driven Image Restoration

arXiv: https://arxiv.org/abs/2606.26615 · HTML: https://mlsanity.com/paper/2606.26615/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 47 references that are in the map.
- First posted on arXiv 25 Jun 2026.

## Builds on (most cited first, 30 of 34)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks](https://mlsanity.com/paper/1506.01497/) (2015)
- [Visual Instruction Tuning](https://mlsanity.com/paper/2304.08485/) (2023)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction](https://mlsanity.com/paper/2404.02905/) (2024)
- [OpenAI GPT-5 System Card](https://mlsanity.com/paper/2601.03267/) (2025)
- [SeeSR: Towards Semantics-Aware Real-World Image Super-Resolution](https://mlsanity.com/paper/2311.16518/) (2023)
- [An Image is Worth 32 Tokens for Reconstruction and Generation](https://mlsanity.com/paper/2406.07550/) (2024)
- [Towards Robust Blind Face Restoration with Codebook Lookup Transformer](https://mlsanity.com/paper/2206.11253/) (2022)
- [Image-Adaptive YOLO for Object Detection in Adverse Weather Conditions](https://mlsanity.com/paper/2112.08088/) (2021)
- [Task-Driven Super Resolution: Object Detection in Low-resolution Images](https://mlsanity.com/paper/1803.11316/) (2018)
- [Describing Differences in Image Sets with Natural Language](https://mlsanity.com/paper/2312.02974/) (2023)
- [Democratizing Text-to-Image Masked Generative Models with Compact Text-Aware One-Dimensional Tokens](https://mlsanity.com/paper/2501.07730/) (2025)
- [UniRestore: Unified Perceptual and Task-Oriented Image Restoration Model Using Diffusion Prior](https://mlsanity.com/paper/2501.13134/) (2025)
- [Connecting Image Denoising and High-Level Vision Tasks via Deep Learning](https://mlsanity.com/paper/1809.01826/) (2018)
- [FlowTok: Flowing Seamlessly Across Text and Image Tokens](https://mlsanity.com/paper/2503.10772/) (2025)
- [URIE: Universal Image Enhancement for Visual Recognition in the Wild](https://mlsanity.com/paper/2007.08979/) (2020)
- [Is Image Super-resolution Helpful for Other Vision Tasks?](https://mlsanity.com/paper/1509.07009/) (2015)
- [Does Haze Removal Help CNN-based Image Classification?](https://mlsanity.com/paper/1810.05716/) (2018)
- [Beyond Image Super-Resolution for Image Recognition with Task-Driven Perceptual Loss](https://mlsanity.com/paper/2404.01692/) (2024)
- [Highly Compressed Tokenizer Can Generate Without Training](https://mlsanity.com/paper/2506.08257/) (2025)

## Built on by (most cited first, 3 of 3)

- [TaskGuard: Task-Conditioned Restoration Utility for Risk-Aware Object Detection](https://mlsanity.com/paper/2609.08011/) (2026)
- [TaskIR: Task-Driven Image Restoration via Degradation Adaptation and Task Feedback](https://mlsanity.com/paper/2609.31170/) (2026)
- [FireNav: Decision-Oriented Reconstruction Evidence Routing for Vision-Language Navigation in Indoor Fire Environments](https://mlsanity.com/paper/2610.15595/) (2026)

## Abstract (arXiv)

While traditional image restoration focuses on perceptual quality, Task-Driven Image Restoration (TDIR) aims to maximize the performance of downstream high-level vision tasks. Recent approaches leveraging generative priors have shown promise for TDIR; however, they typically suffer from computational inefficiency and potential semantic alteration by indiscriminately updating all latent tokens. In this paper, we posit that not all visual information is equally important for machine perception. Through an analysis of the latent token space, we observe that task-relevant cues are unevenly distributed across the token sequence, exhibiting index-wise specialization. This suggests that selectively refining a subset of tokens can be sufficient for task-driven objectives. Leveraging this insight, we propose TaskTok, a novel framework that selectively restores only task-relevant tokens via a learnable token switch and a lightweight token refinement module. Extensive experiments across image classification, semantic segmentation, and object detection demonstrate that TaskTok significantly enhances task performance with high computational efficiency. The source code is available at https://github.com/jimmy9704/TaskTok