Virtual Task Prompting for Multi-Task Scene Understanding

Zedong Wang*,     Yucheng Wang*,     Dan Xu†
The Hong Kong University of Science and Technology
*Equal contribution    †Corresponding author

NeurIPS 2026

Figure 1: Overview of Virtual Task Prompting (VTP). (a) Encoder of TaskPrompter, (b) Encoder of VTP, (c) Decoder of VTP.

Overview of Virtual Task Prompting (VTP). (a) TaskPrompter modifies the ViT encoder to compensate for limited prompt capacity. (b) VTP instead maintains \(T{+}V\) virtual states alongside the encoder: before each block it reads \(T\) task vectors to interact with image tokens, and afterwards writes the results back to the \(T{+}V\) bypass, enriching task representations while keeping a vanilla ViT block. (c) VTP keeps the prompt-guided decoding pipeline but unifies its spatial and channel branches into one.

TL;DR VTP recasts multi-task prompting as learnable read–write routing over a small virtual task state that lives outside the ViT. The encoder stays vanilla, so off-the-shelf vision foundation models plug in unchanged, achieving competitive results on Pascal-Context, NYUD-v2, and Taskonomy, with fewer parameters and higher throughput than baselines.

Abstract

Multi-task scene understanding models typically require modifications to the vision encoder or complex decoders for cross-task interaction, which limits their flexibility and compatibility with off-the-shelf vision foundation models (VFMs). This paper introduces Virtual Task Prompting (VTP), which frames multi-task prompting as retrieving task-adaptive subsets from a shared virtual state, where both the state and accessing process itself are optimizable in an end-to-end manner. Concretely, VTP maintains a set of virtual tasks alongside the model, and at each layer, learnable routing matrices read an actual task composition as input for interaction with image tokens, then write the processed results back, forming explicit learning pathways that capture the intricacies among tasks before and after the computation of each encoder block. Since VTP keeps the backbone clean, it integrates seamlessly with recent VFMs like DINOv3, advancing the state of the art on Pascal-Context and NYUD-v2. Stress test on Taskonomy with 13 heterogeneous tasks further shows VTP's robustness in extreme settings. Ablations and analysis confirm that VTP's gains stem mainly from its prompting design rather than merely increased capacity.

Virtual Task Prompting

Inspired by hyper-connections, VTP keeps \(T{+}V\) learnable vectors—\(T\) task slots plus \(V\) virtual tasks—as a bypass outside the encoder. Before each ViT block it reads \(T\) task-adaptive compositions from this state; after the block it writes the processed vectors back. Each task is no longer a fixed vector but a learned composition of a richer state, where tasks may share slots or access distinct ones.

Swipe sideways to see the full diagram →

One VTP read-write cycle at layer l The virtual state of T task slots and V virtual slots is read by a routing matrix Phi into T task vectors, which are concatenated with image tokens and processed by an unchanged ViT block. The processed task vectors are written back by Psi and combined with the previous state carried by Gamma to form the updated state. T V p1 p2 p3 v1 v2 Pl−1 virtual state Read Φl T × (T+V) P̃l N image tokens Xl−1 x1 x2 x3 ⋮ ViT Block Fl Multi-Head Self-Attn. Feed-Forward LayerNorm · Residual T N T N attention over T+N tokens P̂l Xl x1 x2 x3 ⋮ to layer l+1 ✓ unchanged · loads VFM weights as-is Write Ψl (T+V) × T Carry Γl  (T+V) × (T+V) p1 p2 p3 v1 v2 Pl updated state to layer l+1
task slot virtual slot read / processed task vector image token routing weight 1 / 0

One VTP read–write cycle at layer \(l\) with \(T{=}3\) tasks and \(V{=}2\) virtual tasks. Routing matrices are drawn at their initialization: each task reads from and writes back to its own slot, which reduces to standard prompting; the routing structure then emerges during end-to-end training. Hover over a step below to highlight it.

1State ReadEq. 5
\[\tilde{\mathbf{P}}^{l} = \Phi^{l}\,\mathbf{P}^{l-1} \in \mathbb{R}^{T\times C}\]

A learnable routing matrix \(\Phi^{l}\in\mathbb{R}^{T\times(T+V)}\) reads exactly \(T\) task-specific compositions from the \(T{+}V\) virtual state. Tasks may draw on shared slots for common structure or on distinct ones for specialization.

2Vanilla InteractionEq. 6
\[[\hat{\mathbf{P}}^{l};\mathbf{X}^{l}] = \mathcal{F}_{l}([\tilde{\mathbf{P}}^{l};\mathbf{X}^{l-1}])\]

The read vectors are concatenated with the image tokens and processed by an unchanged ViT block. Because the input keeps the standard \((T{+}N)\times C\) shape, self-attention, FFN and residuals—and their pretrained weights—stay untouched.

3State WriteEq. 7
\[\mathbf{P}^{l} = \Psi^{l}\,\hat{\mathbf{P}}^{l} + \Gamma^{l}\,\mathbf{P}^{l-1}\]

\(\Psi^{l}\) writes the processed task vectors into selected virtual slots, while \(\Gamma^{l}\) carries forward what to retain from the previous state, so discriminative task information accumulates across layers.

Streamlined prompt-guided decoding

Because VTP's task vectors are not split into spatial and channel parts, TaskPrompter's dual-branch decoder collapses into a single pass. For each task \(t\), the task-specific affinity \(\mathbf{M}^{L}_{t,:}\) routes shared features to the task, and the cross-task affinity \(\mathbf{R}^{L}_{t,:}\) lets it retrieve complementary information from the others:

\[\mathbf{M}^{L}_{t,:} = \operatorname{softmax}\!\Big(\frac{\mathbf{X}^{L}{\mathbf{p}^{L}_{t}}^{\!\top}}{\sqrt{C}}\Big)\in\mathbb{R}^{N}, \qquad \hat{\mathbf{Z}}_{t} = \sum_{k=1}^{T}\mathbf{R}^{L}_{t,k}\,\big(\underbrace{\mathbf{M}^{L}_{k,:}\odot\mathbf{X}^{L}}_{\text{task-specific}}\big)\in\mathbb{R}^{N\times C}\]

A lightweight head (a \(3{\times}3\) convolution followed by a linear projection) then produces each task's prediction, applied to features from \(K{=}4\) evenly spaced encoder layers. All task losses are summed, \(\mathcal{L}_{\text{MTL}}=\sum_{t}\lambda_t\mathcal{L}_t\), so the routing matrices \(\Phi^l,\Psi^l,\Gamma^l\) receive gradients from every task simultaneously.

Design properties

Vanilla backbone, exact weight shapes

No attention heads are inserted and no residual streams are altered, so pretrained checkpoints (timm ViT-L, DINOv3, ViT-5) load as-is, with no remapping, no code changes and no VTP re-tuning.

Negligible routing overhead

Routing touches only the task stream at \(\mathcal{O}((T{+}V)^2C)\), independent of the number of image tokens \(N\). Since \(T,V\ll N\), its cost vanishes next to the encoder's \(\mathcal{O}((T{+}N)^2C)\).

Starts as standard prompting

\(\Phi^l\) and \(\Gamma^l\) are initialized as block identities and \(\Psi^l\) cyclically (\(\Psi^l_{ij}{=}1\) iff \(i\equiv j \bmod T\)); task vectors start from \(\mathcal{N}(1,1)\). Training begins from a stable pass-through and learns its routing.

Scales with the task count

A new task adds one prompt, one row/column in the routing matrices and one light head. The same design runs 4, 5 and 13 tasks, with a fixed \(V{=}2\).

  PyTorch-style sketch of the VTP encoder, transcribed from Eqs. 5–7 (illustrative, not the official implementation)
# T tasks, V virtual tasks (default 2), C channels; vit: any off-the-shelf ViT
L = len(vit.blocks)
cyclic = lambda n, T: torch.eye(T)[torch.arange(n) % T]    # Psi[i, j] = 1 iff i mod T == j

P0    = nn.Parameter(torch.randn(T + V, C) + 1.0)          # virtual state ~ N(1, 1)
Phi   = nn.Parameter(torch.eye(T, T + V).repeat(L, 1, 1))  # read : (L, T, T+V)
Psi   = nn.Parameter(cyclic(T + V, T).repeat(L, 1, 1))     # write: (L, T+V, T)
Gamma = nn.Parameter(torch.eye(T + V).repeat(L, 1, 1))     # carry: (L, T+V, T+V)

def vtp_encode(x):                                # x: (B, N, C) patch tokens
    P = P0.expand(x.shape[0], -1, -1)             # (B, T+V, C)
    feats = []
    for l, block in enumerate(vit.blocks):        # vanilla, pretrained ViT blocks
        P_read = Phi[l] @ P                       # Eq. 5  read          -> (B, T, C)
        z = block(torch.cat([P_read, x], dim=1))  # Eq. 6  standard (T+N) x C input
        P_hat, x = z[:, :T], z[:, T:]
        P = Psi[l] @ P_hat + Gamma[l] @ P         # Eq. 7  write + carry -> (B, T+V, C)
        if l in out_layers:                       # K = 4 evenly spaced blocks
            feats.append((P_hat, x))              # -> prompt-guided decoding (Eq. 8)
    return feats

What Do the Task Vectors Learn?

Pairwise cosine similarity among task vectors across the 24 ViT-L blocks on Pascal-Context, and a 3D UMAP of VTP's \(T{=}5\) read task vectors at blocks 6, 12, 18 and 24 over 500 test images.

Figure 4: (a) layer-wise correlation of TaskPrompter's task vectors, (b) UMAP of VTP's layer-wise task vectors, (c) layer-wise correlation of VTP's task vectors.

Task correlation and UMAP analysis. (a) TaskPrompter's task vectors collapse into uniformly high correlation with negligible variation in deep layers. (b) VTP's task vectors form clear layer-wise clusters, an internally cohesive latent space that encodes layer-wise discriminative information. (c) VTP exhibits structured task affinities that capture adaptive inter-task differentiation across depth.

BibTeX

@inproceedings{neurips2026vtp,
  title={Virtual Task Prompting for Multi-Task Scene Understanding},
  author={Wang, Zedong and Wang, Yucheng and Xu, Dan},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2026}
}