Query-Aware Token Selection

• Query-aware token selection is a dynamic approach that prioritizes tokens based on the input query, reducing computational cost and memory usage in large-scale models.
• It employs techniques like dot-product scoring, softmax-normalized cross-attention, and conditional mutual information to efficiently filter tokens while maintaining accuracy.
• Empirical results demonstrate significant speedups and reduced inference load in long-context LLMs and vision models, making it valuable for scalable multimodal reasoning.

Query-aware token selection refers to a class of algorithms and architectural modules that dynamically select or prioritize a subset of tokens for further processing—conditioned explicitly on the input query or task—within large-scale models for language, vision, or multimodal reasoning. This approach addresses the computational and memory bottlenecks induced by the quadratic complexity of self-attention and the rapid growth of context (e.g., long text sequences or extended video streams), while preserving or even enhancing model performance on task-relevant outputs. Query-aware selection stands in contrast to query-agnostic (uniform or heuristic) pruning and has emerged as a central paradigm for efficient, scalable inference in long-context, retrieval, and reasoning tasks.

1. Core Principles and Motivation

The critical hypothesis underpinning query-aware token selection is that, for any given query (text prompt, question, or decision context), only a small, query-dependent subset of tokens in the input context or state bear significant influence on model predictions. This sparsity is highly dynamic: different queries may require distinct tokens, often scattered non-contiguously across the input sequence. Consequently, query-aware methods prioritize those tokens that maximize utility for the current task, minimizing redundancy in computation and maximizing content relevant to the query (Luo et al., 11 Mar 2025, Tang et al., 16 Jun 2024, Wu et al., 5 Nov 2024, Yuan et al., 23 May 2025).

Key motivating factors include:

• The quadratic scaling of self-attention in both LLMs and vision transformers.
• The rapid growth in sequence/context length in language (up to millions of tokens) and vision (e.g., patch tokens from long videos).
• The empirical observation that, for any query, softmax attention is dominated by a handful of tokens/pages/patches—a property exploited for efficient selection (Tang et al., 16 Jun 2024, Wu et al., 5 Nov 2024, Li et al., 14 Nov 2025).

2. Mathematical Formulations and Selection Mechanisms

Query-aware token selection methods operationalize token utility as a (possibly normalized) function of query–token interaction. Below are representative strategies from the literature:

• Dot-Product Scoring (Language):

For long-context LLMs, attention scores are computed as $s_{i,j} = q_i^\top k_j$ for query $q_i$ (at current step) and historical key $k_j$. The top-$k$ keys with maximal $|s_{i,j}|$ are selected for attention (Wu et al., 5 Nov 2024, Wang et al., 20 Feb 2025). Multi-head settings aggregate importance via per-head voting or softmax-weighted sum (Wu et al., 5 Nov 2024).

• Softmax-normalized Cross-attention (Multimodal):

In visual-LLMs and long-video LMMs, per-token relevance is often computed by running cross-attention from text query embeddings $Q$ to vision tokens $X$, extracting maximum attention weights as token relevance scores: $r_i = \max_{h, \ell} \alpha_{h,\ell,i}$, where $\alpha_{h,\ell,i}$ is the cross-attention output (Li et al., 14 Nov 2025, Jiao et al., 20 Nov 2024).

• Feature Compression and Cheap Proxies:

Token importance may also be computed in a compressed space via linear projections $f_{\theta_q}, f_{\theta_k}$ to enable efficient scoring: $f_s(m;c) \approx q'_c\,k'_m^\top$ (Wang et al., 20 Feb 2025).

• Bidirectional Scoring (Conditional Importance):

For reasoning chains, token relevance is computed as the conditional mutual information (CMI) between each token and the answer, conditioned on the question: $$r_i = -\log P(t_i|t_{<i}, x) + \log P(t_i|t_{<i}, x, y)$$, favoring tokens whose prediction is significantly helped by knowing the answer $y$ (Yuan et al., 23 May 2025).

These scores are then used to define a selection policy—often simply top-$n$ by relevance, possibly after budget prediction or further refinement (e.g., via soft gating or differentiable thresholding in training (Li et al., 14 Nov 2025)).

3. Representative Methods and Algorithms

Several representative and widely-cited methods have instantiated query-aware token selection in distinct domains:

Method	Domain	Core Mechanism/Score
TokenSelect	Long-text LLMs	Per-head Q–K dot-product, head-soft-vote top-$k$
Quest	Long-text LLMs	Upper-bound criticality via page min/max-K*Q
Efficient SA	Long-text LLMs	Q/K projection, low-dim scoring and proximity
QTSplus	Long-video MLLMs	Cross-attention, query-conditioned budget, top-$n$
QuoTA	Video LLMs	Frame scoring by lightweight LVLM, CoT decoupling
MCAT	Query→video search	Multi-tier query–guided cross-attention, class-tkn
LaVida Drive	Autonomous driving	Image–query alignment (MLP+cosine sim)
CTS	Reasoning traces	Conditional PPL drop (CMI) score, threshold top-$k$

• TokenSelect (Wu et al., 5 Nov 2024) implements per-query, per-head scoring and a soft voting mechanism to select a token-level subset of the KV-cache that minimizes attention output error relative to the full cache. It leverages a selection cache exploiting query similarity for further efficiency.
• Quest (Tang et al., 16 Jun 2024) partitions the KV cache into pages, tracking per-page min/max key values. For each query, it computes an upper bound on dot-products to cheaply rank/screen only the top-K pages as "critical," reducing memory movement and compute with negligible loss in performance.
• QTSplus (Li et al., 14 Nov 2025) computes cross-attention from text to vision tokens to derive per-token relevance, then predicts a query-adaptive budget, and applies differentiable gate or hard selection depending on phase (training/inference). A re-encoder preserves temporal order after token filtering.
• QuoTA (Luo et al., 11 Mar 2025) computes frame-level importance via lightweight LVLM scoring, with chain-of-thought (CoT) query decoupling to generate "entity" or "event" clues. Token allocation is dynamically computed using normalized frame scores under a strict budget, followed by per-frame token interpolation or merging.
• CTS (Yuan et al., 23 May 2025) in reasoning models compresses chains-of-thought by conditional importance scoring conditioned on both the query and answer, yielding significant token reductions without accuracy loss.
• MCAT (Mishra et al., 8 Apr 2025) for ultrasound video localization uses multi-tier, cross-attentional fusion guided by a visual query, followed by class-aware token selection (selecting only the task-relevant class token from a pool), leading to ∼96% token compression.

4. Application Domains and Integration Strategies

Query-aware token selection is widely applied in domains where context length, memory, or frame count is a limiting factor:

• Long-context LLMs: Select or load only the most query-relevant historical KV tokens, enabling accurate QA, retrieval, and copy-paste tasks at context lengths up to 1M tokens (Wu et al., 5 Nov 2024, Tang et al., 16 Jun 2024, Wang et al., 20 Feb 2025).
• Video and Multimodal LMMs: For long-form video QA, dynamically allocate vision tokens based on query relevance at the frame/patch level, as in QTSplus, QuoTA, EXPLORE-THEN-SELECT, and MCAT (Li et al., 14 Nov 2025, Luo et al., 11 Mar 2025, Shi et al., 30 Apr 2025, Mishra et al., 8 Apr 2025).
• Super-Resolution (SISR): SSCAN first routes query windows to most similar key-value windows via region-level dot-product similarity, reducing both compute and DRAM cost (Kim et al., 9 Apr 2025).
• Autonomous driving VQA: LaVida Drive’s query-aware selection ensures that only detail- and question-relevant image tokens are passed to the LLM, achieving both high compression (up to 168×) and state-of-the-art QA metrics (Jiao et al., 20 Nov 2024).

In most architectures, the selection layer sits between the frozen encoder (ViT, ResNet, or Transformer) and the task head, and is plug-and-play or minimally adapted, often incurring negligible additional parameter overhead.

5. Empirical Performance and Efficiency

Empirical studies consistently demonstrate that query-aware token selection yields substantial improvements in both efficiency (latency, memory, and FLOP reduction) and in retrieval/QA accuracy—often outperforming naive or heuristic selection:

• TokenSelect (Wu et al., 5 Nov 2024): Up to 23.84× speedup in attention compute, and 2.28× overall inference speedup, with near parity to full attention on InfiniteBench and other long-context benchmarks.
• Quest (Tang et al., 16 Jun 2024): Attention speedups of 7.03× and total decode reductions of 2.23×; accuracy loss <0.01 PPL for strong compression.
• QTSplus (Li et al., 14 Nov 2025): Up to 89% vision-token compression, 28% end-to-end latency reduction, with negligible accuracy loss and significant gains on TempCompass (+20.5 and +5.6 points direction/order).
• CTS (Yuan et al., 23 May 2025): Achieves up to 75.8% reduction in reasoning tokens with only a 5% drop in solution accuracy; in moderate compression regimes (>85% retention), accuracy often improves owing to removal of spurious "overthinking" steps.
• MCAT (Mishra et al., 8 Apr 2025): Reduces temporal tokens by ∼96% (from 150 to 9) while improving mIoU by 3.1 absolute points in ultrasound video localization.

Ablation studies indicate the necessity of query-contextual scoring (versus unconditional or heuristic metrics) and the retention of a minimum budget for robust semantic coverage. Training-free (plug-and-play) versions prevail in video and multimodal settings due to data scarcity and modularity needs.

6. Comparative Analysis: Query-Aware vs. Query-Agnostic Selection

Query-agnostic approaches (uniform sampling, sliding window, fixed chunking, or low-response token drop) fail to reliably preserve evidence for localized, complex, or retrieval-oriented queries. In contrast, query-aware selection:

• Adapts dynamically to the semantic and localization demands of each query (Li et al., 3 Dec 2025, Li et al., 14 Nov 2025, Luo et al., 11 Mar 2025).
• Enables fine-grained allocation of the token budget, critical in “needle-in-haystack” tasks or multi-hop reasoning.
• Outperforms post-hoc pruning methods that operate only after initial attention/cross-modal fusion, by front-loading the selection and minimizing "dead" token propagation (Luo et al., 11 Mar 2025).
• Preserves positional and temporal order when necessary (e.g., QTSplus’s re-encoder, LaVida Drive’s per-token position tracking).

However, selection granularity and score accuracy trade-off with precompute bandwidth, memory fragmentation, and the potential for selection cache staleness on rapidly-shifting queries.

7. Limitations, Design Challenges, and Future Directions

Design and deployment of query-aware token selection face several challenges:

• Score Proxy Accuracy: Dot-product or cross-attention proxies may underestimate the true conditional relevance, potentially missing critical tokens. Task- and domain-adaptive scoring remains an area of investigation (Yuan et al., 23 May 2025, Wang et al., 20 Feb 2025).
• Overhead of Selection Layers: While selection is far cheaper than full attention, its benefit depends on context length and query-token similarity; special kernels are usually required for maximum savings (Wu et al., 5 Nov 2024, Tang et al., 16 Jun 2024).
• Adaptive Compression Schedules: Budget prediction based on query content (as in QTSplus) is key for dynamic workloads but relies on robust meta-feature extraction.
• Generalization Across Modalities: While core mechanisms are analogous (token-level in text, patch/frame-level in vision), alignment and scoring can differ; multi-tier, multi-modal fusion approaches such as MCAT’s class-token routing demonstrate high compression but may require architectural changes (Mishra et al., 8 Apr 2025).
• Training vs. Plug-and-Play: Most selection modules are training-free; research into end-to-end, jointly-optimized selectors (possibly with learned metrics) is ongoing (Li et al., 14 Nov 2025, Shi et al., 30 Apr 2025).

Active research explores:

• CMI-based or gradient-based scores for more principled selection.
• Adaptive query routers for query-type-specific routing (e.g., DIG’s divide-then-ground (Li et al., 3 Dec 2025)).
• Integration with retrieval-augmented architectures and knowledge-tuned selectors.

---

References:

• "TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection" (Wu et al., 5 Nov 2024)
• "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference" (Tang et al., 16 Jun 2024)
• "Efficient Selective Attention (ESA): Unshackling Context Length..." (Wang et al., 20 Feb 2025)
• "Query-aware Token Selector (QTSplus)" (Li et al., 14 Nov 2025)
• "QuoTA: Query-oriented Token Assignment..." (Luo et al., 11 Mar 2025)
• "MCAT: Visual Query-Based Localization..." (Mishra et al., 8 Apr 2025)
• "Not All Tokens Are What You Need In Thinking (Conditional Token Selection)" (Yuan et al., 23 May 2025)
• "LaVida Drive: Vision-Text Interaction VLM..." (Jiao et al., 20 Nov 2024)
• "Divide, then Ground (DIG): Adapting Frame Selection..." (Li et al., 3 Dec 2025)
• "Static or Dynamic: Towards Query-Adaptive Token Selection..." (Shi et al., 30 Apr 2025)