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

Source: [TaskTok: Delving into Task Tokens for Task-driven Image Restoration](<https://www.sota2.com/research/paper/6a3e07b9206542251d7cb607>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2026-07-23T05:00:05.031Z

Authors: Hongjae Lee, Sojung Kang, Jaeseong Yu, Seung-Won Jung

Published: 2026-06-25T05:20:01.000Z

## Abstract

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

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2606.26615>)

[PDF](<https://arxiv.org/pdf/2606.26615.pdf>)
