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

2026-06-25

AI Summary

This work addresses the high computational cost and semantic distortion commonly incurred by existing task-driven image restoration methods. To this end, the authors propose TaskTok, a novel framework that reveals—for the first time—that task-relevant information in the latent token space exhibits index-level specificity. Leveraging this insight, they introduce a learnable token gating mechanism coupled with a lightweight refinement module to selectively restore only the subset of tokens critical to downstream tasks. Token importance is assessed through generative priors, enabling the method to achieve significant performance gains across diverse tasks—including image classification, semantic segmentation, and object detection—while maintaining high computational efficiency.

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

Problem

• Task-Driven Image Restoration
• Latent Tokens
• Computational Efficiency
• Semantic Alteration
• Downstream Vision Tasks

Innovation

• task-driven image restoration
• latent token selection
• computational efficiency
• generative priors
• token refinement