Cross-modal attention lets VLA models share information across vision, language, and state — revealing how robots ground language in perception and plan actions.
Layer
Click “Run Inference” to see ranked attention pairs.
How information flows between vision, language, and state inside a VLA transformer.
Attention is the mechanism by which a transformer decides which parts of its input to focus on when computing each output token. In single-modality transformers (e.g., a language model), this is self-attention: every token attends to every other token in the same sequence.
Cross-modal attention extends this to multiple input sequences from different modalities. In cross-attention, the query (Q) comes from one modality while the key (K) and value (V) come from another. The result is a weighted sum of values from the source modality, where the weights are determined by the similarity between the target query and the source keys.
In self-attention, all of Q, K, V come from the same sequence. A language token attending to other language tokens is self-attention — it helps the model understand syntax and coreference within the instruction.
In cross-attention, Q comes from sequence A while K and V come from sequence B. An action token querying the image token sequence is cross-attention — the action is asking the image: “what visual features are relevant to what I need to do next?” This is how language instructions get “grounded” in visual perception.
Both types appear in VLA models. pi0.5 uses the LLM’s native attention to mix vision and language tokens. UnifoLM-VLA’s DiT explicitly alternates: even layers = cross-attention (action → vision/text), odd layers = self-attention (within state/future/action tokens).
The Diffusion Transformer (DiT) in UnifoLM-VLA is the action expert network. It receives the current noisy action tokens and denoises them conditioned on the visual and language context from the language model backbone.
Layer structure:
VLAExplain reveals distinct patterns at different layers for both models:
These patterns are what the interactive attention matrix above lets you explore. The “Layer” slider simulates moving from early (vision-dominant) to late (action-refinement) layers.
Both pi0.5 and UnifoLM-VLA also expose cross-modal attention within the language model backbone — the part that processes vision and language together before the action expert takes over.
VLAExplain analyses: text → image (which image patches does a text token attend to?), image → text (which text tokens does an image patch attend to?), and in pi0.5 additionally: text → state, image → state, state → text, state → image. Patterns here reflect how the language model “grounds” the instruction in the visual context.
It is important to distinguish these two analysis types in VLAExplain:
Together they provide complementary views: ActionCode answers “what drove this specific action?” while cross-modal attention answers “how does information flow throughout the model?”
Understanding cross-modal attention patterns has practical consequences for robot safety and debugging:
The VLAExplain Gradio app supports the following controls for cross-modal attention analysis in both pi0.5 and UnifoLM-VLA: step selector (choose which action step to inspect), layer index (examine attention at any transformer layer), head number (inspect individual attention heads), normalization method, interpolation method, color map selection, and opacity control.
For pi0.5, all modality combinations (text → image, image → text, text → state, etc.) are supported. For UnifoLM-VLA, text → image and image → text are supported in the language model; in the DiT action expert, cross-attention (action → vision/text) and self-attention (state, future tokens, actions) are both available.
Say the answer out loud before you reveal it — recalling it is what makes it stick, and rereading it is not.
What is the key difference between self-attention and cross-attention?
How does the DiT in UnifoLM-VLA alternate between attention types across layers?
What is the difference between ActionCode attribution and cross-modal attention in VLAExplain?
Why do action tokens attend more to vision in early layers and more to language in later layers?
Self-attention: Q,K,V from same modality. Cross-attention: Q from A, K/V from B. UnifoLM-VLA DiT: even layers = cross-attention (action → vision/text), odd layers = self-attention (state/future/action). Early layers: vision-dominant. Later layers: language-conditioned. ActionCode = per-action view. Cross-modal = full model flow.