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

While traditional image restoration focuses on perceptualquality, Task-Driven Image Restoration (TDIR) aims to maximize theperformance of downstream high-level vision tasks. Recent approachesleveraging generative priors have shown promise for TDIR; however,they typically suffer from computational inefficiency and potential se-mantic alteration by indiscriminately updating all latent tokens. In thispaper, we posit that not all visual information is equally important formachine perception. Through an analysis of the latent token space, weobserve that task-relevant cues are unevenly distributed across the to-ken sequence, exhibiting index-wise specialization. This suggests thatselectively refining a subset of tokens can be sufficient for task-drivenobjectives. Leveraging this insight, we propose TaskTok, a novel frame-work that selectively restores only task-relevant tokens via a learnabletoken switch and a lightweight token refinement module. Extensive ex-periments across image classification, semantic segmentation, and objectdetection demonstrate that TaskTok significantly enhances task perfor-mance with high computational efficiency. The source code is availableat our project page https://github.com/jimmy9704/TaskTok.