TaskTok: Delving into Task Tokens for Task-driven Image Restoration
For practitioners of task-driven image restoration, this work offers a computationally efficient method that avoids indiscriminate updates, though it is an incremental improvement over existing generative prior approaches.
TaskTok selectively restores only task-relevant tokens in the latent space, achieving significant improvements in downstream task performance (e.g., classification, segmentation, detection) while maintaining high computational efficiency.
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