QA-ViT: Query-Aware Vision Transformer

• QA-ViT is a vision transformer that conditions visual features on query context by integrating question embeddings directly into its upper layers.
• The method injects question tokens via per-layer MLPs and gated residual projections, resulting in dynamic feature extraction with measurable performance boosts.
• Empirical results show consistent improvements across VQA, captioning, and scene-text tasks, validating the benefit of early query-adaptive fusion.

Question Aware Vision Transformer (QA-ViT) directly integrates question awareness into the vision encoder component of vision-language (VL) architectures, enabling dynamic visual features that are conditioned on the textual input—typically, user questions about an image. Distinct from preceding approaches where the vision encoder processes the image in isolation and only later fuses with text representations, QA-ViT enhances the vision backbone itself so that image features attended to by the model are tailored to the query of interest. This method introduces minimal architectural changes, is model-agnostic, and consistently improves performance across a spectrum of vision-language reasoning tasks and datasets (Ganz et al., 2024).

1. Architectural Placement and Core Pipeline

Modern VL models are generally structured with three main components: a vision encoder (e.g., CLIP ViT or EVA-CLIP), a projection module (an MLP or cross-attention-based “Q-Former”), and a LLM such as Flan-T5 or Vicuna. The canonical forward flow for these systems is

$$I \xrightarrow{\text{Vision Encoder}} F_V \xrightarrow{\text{Projection}} \tilde F_V \xrightarrow{\text{LLM}(\text{w/}~Q)} \text{Answer},$$

where $I$ is the image, $F_V$ are visual tokens, and $Q$ is the question.

QA-ViT augments this pipeline by embedding question awareness directly inside the vision encoder. The visual feature extraction changes from

$V(I) \to F_V$

to

$V(I~|~Q) \to F_{VQ},$

so that before cross-modal projection or language modeling, the image representation is already conditioned on the question. This design ensures the features input into later multimodal modules are already relevant to the query context (Ganz et al., 2024).

2. Question Injection into Vision Transformer Layers

QA-ViT injects the question embedding into self-attention layers at the uppermost $L$ of $N$ total layers in the vision transformer. The injection process involves several key steps:

• Question Encoding: The question $Q$ is passed through an encoder $\mathcal{E}$—typically the LLM’s own text encoder—to yield $I$0, where $I$1 is the token count and $I$2 is the embedding dim.
• Layer-specific Projection: At each of the top $I$3 layers, $I$4 is projected to the vision feature space via a per-layer MLP, producing $I$5 for layer $I$6.
• Concatenated Self-Attention: At each chosen layer, the visual sequence $I$7 and question tokens $I$8 are concatenated: $I$9. This augmented sequence passes through the (frozen) ViT self-attention mechanism without modification:

$F_V$0

The output is sliced to retain only the first $F_V$1 visual positions, yielding $F_V$2.

• Gated Residual Projection: The conditioned features are processed as:

$F_V$3

where $F_V$4 is the frozen projection, $F_V$5 is learned, and $F_V$6 is a learnable scalar initialized at zero.

A summary of the sequence injection and update procedure is displayed in the following table:

Operation	Input	Output Dimension
Encode Question	$F_V$7	$F_V$8
Project (per layer)	$F_V$9	$Q$0
Concatenate	$Q$1	$Q$2
Self-attention	$Q$3	$Q$4
Slice Visual	$Q$5	$Q$6
Gated Projection	$Q$7	$Q$8

3. Training Objectives and Optimization

QA-ViT does not introduce bespoke losses. Training proceeds end-to-end using conventional cross-entropy or maximum-likelihood objectives for visual question answering (VQA) and image captioning:

$Q$9

where $V(I) \to F_V$0 is the $V(I) \to F_V$1-th token in the ground-truth answer. In multi-task regimes (VQA, scene-text VQA, captioning), this loss is summed across relevant examples. The only additional learnable parameter is the scalar $V(I) \to F_V$2 for gated projection, which is implicitly regularized by its initialization at zero.

4. Empirical Evaluation and Performance Gains

QA-ViT demonstrates consistent performance improvements across diverse datasets and VL architectures. Key datasets include:

• General VQA: VQA v2, Visual Genome, COCO caption, TextCaps
• Scene-text VQA: VQA T, VQA ST, TextCaps, VizWiz (zero-shot)
• Document VQA: DocVQA, InfoVQA, ChartQA

Systems augmented with QA-ViT, such as ViT+T5, BLIP2, InstructBLIP, and LLaVA-1.5, report the following representative gains:

Model	Dataset	Metric	Baseline	QA-ViT	$V(I) \to F_V$3
ViT+T5-base	VQA v2	Accuracy	66.5	71.7	+5.2
ViT+T5-base	TextCaps	CIDEr	86.3	96.1	+9.8
BLIP2-xl	VQA v2	Accuracy	72.5	74.6	+2.1
LLaVA-1.5	General avg	Score	106.6	107.6	+1.0
LLaVA-1.5	Scene-text	Score	93.0	94.7	+1.7

Metrics include vqa-score, CIDEr, and ANLS.

5. Ablations and Qualitative Analysis

Ablation studies reveal:

• Fusion Location: Introducing question-token injection solely in the top $V(I) \to F_V$4 of $V(I) \to F_V$5 ViT layers is optimal. Early or all-layer injection is inferior.
• Question Encoder: Leveraging the LLM’s own encoder for $V(I) \to F_V$6 outperforms alternatives such as raw token embeddings, BERT, or the CLIP text encoder.
• Prompt-Tuning vs. QA-ViT: Deep prompt insertion into ViT layers produces only a marginal improvement of $V(I) \to F_V$7 on VQA v2, compared to $V(I) \to F_V$8 for QA-ViT.
• Training Data Mix: Augmenting with COCO captions and document VQA data benefits both captioning and VQA, supporting the utility of multi-task learning.

Qualitative analyses using Grad-CAM and attention head maps show that, while a standard frozen ViT attends to large, salient objects regardless of $V(I) \to F_V$9, QA-ViT redirects attention to image regions referenced in the question, adapting its focus dynamically as $V(I~|~Q) \to F_{VQ},$0 changes. This substantiates genuine early-stage fusion between vision and language spaces.

6. Integration and Applicability

QA-ViT acts as a lightweight, plug-in module for ViT-based VL pipelines. It supports compatibility with any ViT variant and downstream architecture with negligible computational overhead since most weights remain frozen. The approach allows for dynamic, query-adaptive vision features without introducing specialized, task-specific losses or major retraining costs. Empirical evidence establishes large, robust gains not only on standard VQA and captioning tasks but also on zero-shot and scene-text VQA benchmarks (Ganz et al., 2024).

7. Broader Implications and Limitations

The integration of query awareness directly in the vision encoder demonstrates that conditioning visual feature extraction on question context yields both improved alignment between modalities and enhanced task performance. This suggests a general principle for multimodal systems: deferring cross-modal interaction until after visual encoding underutilizes the available information. A plausible implication is that further improvements in vision-language reasoning may result from earlier and deeper cross-modal fusions within model architectures.

No major controversies are reported in the source material; however, the method has not been shown to benefit from early or pervasive question injection across all layers, indicating limits to when/where query-awareness is beneficial. The use of frozen weights for almost all components minimizes overfitting but may, in some domains, restrict flexibility. Empirical results do not indicate failures in generalization; rather, consistent improvement across varied tasks is observed.