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

• URL: http://arxiv.org/abs/2606.26615v1
• Date: Thu, 25 Jun 2026 05:20:01 GMT
• Title: TaskTok: Delving into Task Tokens for Task-driven Image Restoration
• Abstract summary: Task-Driven Image Restoration (TDIR) aims to maximize the performance of downstream high-level vision tasks.<n>Recent approaches leveraging generative priors have shown promise for TDIR.<n>We propose TaskTok, a novel framework that selectively restores only task-relevant tokens.
• Score: 14.471297434908209
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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

Related papers

• Downstream Task Inspired Underwater Image Enhancement: A Perception-Aware Study from Dataset Construction to Network Design [53.9340120911759]
  We propose a Downstream Task-Inspired Underwater Image Enhancement (DTI-UIE) framework to enhance images effectively for underwater vision tasks.<n>Specifically, we design an efficient two-branch network with task-aware attention module for feature mixing.<n>We show that DTI-UIE significantly improves task performance by generating preprocessed images that are beneficial for downstream tasks.
  arXiv  Detail & Related papers  (2026-03-02T11:50:09Z)
• SPOT: Sparsification with Attention Dynamics via Token Relevance in Vision Transformers [9.24617274894642]
  We present SParsification with attentiOn dynamics via Token relevance (SPOT) as a framework for early detection of redundant tokens within Vision Transformers (ViTs)<n>SPOT informs token sparsification and facilitates the elimination of such tokens, improving computational efficiency without sacrificing performance.<n> Empirical evaluations demonstrate significant efficiency gains of up to 40% compared to standard ViTs.
  arXiv  Detail & Related papers  (2025-11-13T16:56:24Z)
• ZSPAPrune: Zero-Shot Prompt-Aware Token Pruning for Vision-Language Models [7.7352936204066]
  We propose a novel, zero-shot method to model visual token pruning as a balance between task relevance and information diversity.<n>Our method achieves performance that matches or surpasses the state-of-the-art with only minimal accuracy loss.<n>These gains are accompanied by significant reductions in GPU memory footprint and inference latency.
  arXiv  Detail & Related papers  (2025-10-20T06:18:47Z)
• VisionThink: Smart and Efficient Vision Language Model via Reinforcement Learning [95.89543460132413]
  Vision-language models (VLMs) have improved performance by increasing the number of visual tokens.<n>However, most real-world scenarios do not require such an extensive number of visual tokens.<n>We present a new paradigm for visual token compression, namely, VisionThink.
  arXiv  Detail & Related papers  (2025-07-17T17:59:55Z)
• ToDRE: Visual Token Pruning via Diversity and Task Awareness for Efficient Large Vision-Language Models [59.47738955960352]
  ToDRE is a two-stage and training-free token compression framework.<n>It achieves superior performance by pruning tokens based on token Diversity and token-task RElevance.
  arXiv  Detail & Related papers  (2025-05-24T15:47:49Z)
• End-to-End Vision Tokenizer Tuning [73.3065542220568]
  The vision tokenizer optimized for low-level reconstruction is to downstream tasks requiring varied representations and semantics.<n>The loss of the vision tokenization can be the representation bottleneck for target tasks.<n>We propose ETT, an end-to-end vision tokenizer tuning approach that enables joint optimization between vision tokenization and target autoregressive tasks.
  arXiv  Detail & Related papers  (2025-05-15T17:59:39Z)
• Multi-Task Label Discovery via Hierarchical Task Tokens for Partially Annotated Dense Predictions [44.78165979575075]
  We propose a novel approach to optimize a set of compact learnable hierarchical task tokens.<n>The global task tokens are designed for effective cross-task feature interactions in a global context.<n>A group of fine-grained task-specific spatial tokens for each task is learned from the corresponding global task tokens.<n>The learned global and local fine-grained task tokens are further used to discover pseudo task-specific dense labels at different levels of granularity.
  arXiv  Detail & Related papers  (2024-11-27T23:53:27Z)
• ToSA: Token Selective Attention for Efficient Vision Transformers [50.13756218204456]
  ToSA is a token selective attention approach that can identify tokens that need to be attended as well as those that can skip a transformer layer. We show that ToSA can significantly reduce computation costs while maintaining accuracy on the ImageNet classification benchmark.
  arXiv  Detail & Related papers  (2024-06-13T05:17:21Z)
• Leveraging Auxiliary Tasks with Affinity Learning for Weakly Supervised Semantic Segmentation [88.49669148290306]
  We propose a novel weakly supervised multi-task framework called AuxSegNet to leverage saliency detection and multi-label image classification as auxiliary tasks. Inspired by their similar structured semantics, we also propose to learn a cross-task global pixel-level affinity map from the saliency and segmentation representations. The learned cross-task affinity can be used to refine saliency predictions and propagate CAM maps to provide improved pseudo labels for both tasks.
  arXiv  Detail & Related papers  (2021-07-25T11:39:58Z)
• Automatic Shortcut Removal for Self-Supervised Representation Learning [39.636691159890354]
  In self-supervised visual representation learning, a feature extractor is trained on a "pretext task" for which labels can be generated cheaply, without human annotation. Much work has gone into identifying such "shortcut" features and hand-designing schemes to reduce their effect. We show that this assumption holds across common pretext tasks and datasets by training a "lens" network to make small image changes that maximally reduce performance in the pretext task.
  arXiv  Detail & Related papers  (2020-02-20T16:00:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.