ZULU: A Flow-Matching Action Model Architecture with DINOv2 Dense Representations

Abstract

Currently, many vision-language-action (VLA) models rely on heavy vision-language backbones to interpret environmental observations. While pre-trained Vision Transformers provide robust out-of-the-box semantic and spatial representations. In this technical report, I explore an alternative approach. Rather than relying on a heavy vision-language (VL) backbone, I investigate driving an action model using a vision transformer that natively provides dense patch-level representations, clear object boundary detection, and consistent geometric alignment. To explore this idea, I present ZULU, an exploratory VLA architecture inspired by the joint-prediction paradigm of the DreamZero project. ZULU directly feeds the vision transformer features directly into the Diffusion Transformer (DiT) backbone. I introduce a block causal attention mechanism to manage an interleaved sequence of images, states, and continuous action tokens, and apply a flow-matching objective to model physical dynamics. This technical report models my architectural concept and implementation, sharing exploratory work as an open-source project.

I. Introduction

A. Architectural Trends in Vision-Language-Action Systems

The pursuit of generalist robotics control has increasingly relied on the development of Vision-Language-Action (VLA) models. Currently, the dominant trend in the field involves adapting massive vision-language backbones—often originally designed for internet-scale multi-modal reasoning—to interpret environmental observations and output physical actions. While these large-scale architectures demonstrate remarkable semantic understanding, they inherently require immense computational resources to train and deploy. Relying on heavy, parameter-dense layers to process visual inputs introduces significant memory constraints and computational overhead, making it difficult to explore continuous robotic control without access to large-scale computing infrastructure.

B. From Video Diffusion to World Action Models

A compelling alternative to standard VLA design is the World Action Model paradigm, as demonstrated by DreamZero, which utilizes an auto-regressive video diffusion model for the joint prediction of video frames and continuous actions. At its core, DreamZero relies on a highly structured attention mechanism to process multi-modal inputs. The architecture interleaves visual, proprioceptive state, and continuous action tokens into a unified sequence. It applies a block-causal self-attention mask across this sequence, allowing the model to strictly regulate temporal dependencies—ensuring that action and state tokens attend to the correct historical visual context without leaking future information. Language instructions are subsequently injected into the network via dedicated cross-attention layers, conditioning the physical trajectories on high-level semantic goals. While this token-routing architecture is elegant, the original DreamZero model operates at a massive 14-billion-parameter scale, largely because the backbone is tasked with learning the physical dynamics of the environment by predicting future world states and actions using video as representation. This presents an intriguing opportunity for exploration: what if we adapt this exact attention mechanism but replace the heavy visual processing burden with a pre-trained Vision Transformer? By feeding the attention blocks with features from a vision transformer that natively provides dense patch-level representations, clear object boundary detection, and consistent geometric alignment, it becomes possible to explore this advanced token-routing paradigm without requiring a massive parameter footprint.

C. Proposed Architecture of ZULU

To explore this alternative approach, I present ZULU, an exploratory World Action Model architecture. ZULU is designed to test whether the structural token-routing of a world action model can be successfully driven by DINOv2, a pre-trained Vision Transformer rather than a massive generative backbone. To implement this idea, the architecture is built upon three primary mechanical components:

- Direct Vision-to-DiT Integration: Instead of relying on a parameter-heavy video diffusion backbone to learn environmental physics from scratch, ZULU directly feeds spatial features from a pre-trained Vision Transformer into a Diffusion Transformer (DiT) backbone.
- Block-Causal Attention Masking: To process the multi-modal inputs, I implemented a block-causal attention mechanism that dictates a strict spatiotemporal routing. Visual (DINOv2) tokens remain clean; the first frame serves as a ground-truth anchor attending only to itself, while subsequent frames attend to the first frame and recent frame blocks. Noisy action tokens attend to the visual context, their own action block, and the corresponding state block, enabling the visual representations to guide action denoising. And all tokens uniformly process language instructions via cross-attention. This ensures a unidirectional flow of information (vision to action) preventing future leakage.
- Flow-Matching Objective: Action generation and physical dynamics modeling are formulated through a flow-matching objective rather than standard regression, allowing the DiT to smoothly predict continuous action trajectories.

Fig. 1. Overview of the ZULU architecture. Multi-modal inputs—including visual patches, proprioceptive states, and language instructions—are projected into a unified sequence and processed via block-causal attention layers. During training, the model optimizes a flow-matching objective to denoise continuous action trajectories, enabling it to iteratively generate motor commands from pure Gaussian noise during inference.

D. Implementation and Scope

To investigate the viability of this architectural concept, I implemented the ZULU framework and conducted initial training runs. The complete codebase, including the model architecture, attention mechanisms, and the training pipeline, is publicly available at github. The remainder of this paper details the specifics of this work. Section 2 breaks down the core ZULU architecture, focusing on token routing and flow-matching formulation. Section 3 outlines the implementation details and training setup. Finally, Section 4 discusses the current state of the model, the challenges encountered during training, and potential directions for future optimization and scaling.

II. ZULU Architecture

A. Input Tokenization and Feature extraction

To process heterogeneous multi-modal inputs within a single Diffusion Transformer backbone, ZULU projects visual frames, physical states, continuous actions, and natural language instructions into a unified continuous embedding space.

1) Visual Feature Extraction

Visual observations are processed using a pre-trained Vision Transformer DINOv2. For each frame in the input sequence, patch tokens are extracted by taking the final hidden states of the vision backbone while discarding the standard CLS tokens.

- Patch Extraction: The raw spatial frames are converted into a grid of dense patch embeddings that capture localized spatial features, object boundaries, and fine-grained geometric alignments.
- Layer Normalization and Projection: The extracted patch features are normalized via a LayerNorm layer and linearly projected to match the internal hidden dimension $d$ of the transformer backbone:
$\mathbf{H}_{vis} = \mathbf{W}_{proj} \cdot \text{LayerNorm}(\mathbf{F}_{patch}) + \mathbf{b}_{proj}$
- Clean Context: Crucially, unlike image generation diffusion models, visual patch tokens in ZULU remain un-noised throughout the model. They serve as deterministic, clean context tokens with fixed timestep embeddings set to zero.

2) State and Action Embedding

Physical state observations and continuous action trajectories are tokenized through specialized embodiment-aware mapping networks:

- Proprioceptive State Encoding: Robot joint states and proprioceptive vectors $\mathbf{s}$ are embedded using a category-specific Multi-Layer Perceptron ( $\text{MLP}_{state}$ ). This maps varying state representations into the shared embedding dimension $d$ :
$\mathbf{E}_{state} = \text{MLP}_{state}(\mathbf{s}, \text{id}_{emb})$
- Continuous Action Encoding: Action trajectories $\mathbf{a}$ are processed through a multi-embodiment action encoder. To inform the model of the noise level during the flow-matching process, the flow timestep $t_{action}$ is mapped using a sinusoidal positional encoding $\text{PE}(t_{action})$ and combined with the action representation via embodiment-specific linear transformations and non-linear activations:
$\mathbf{E}_{action} = \text{MLP}_{action}(\mathbf{a}, \text{PE}(t_{action}), \text{id}_{emb})$
- Training vs. Inference Disparity:
- During Training: The model receives ground-truth action trajectories perturbed with Gaussian noise according to sampled flow-matching timesteps $t_{action}$ . The resulting noisy action embeddings and state embeddings are concatenated to form an action-state register appended to the visual sequence.
- During Inference: Ground-truth action inputs are unavailable. The action sequence is initialized entirely from standard Gaussian noise $\mathcal{N}(\mathbf{0}, \mathbf{I})$ . The model iteratively denoises these noisy action tokens over multiple sampling steps using predicted vector fields, while the state tokens remain static physical conditioning context.

3) Language Conditioning: High-level textual instructions are processed through a frozen, pre-trained T5 text encoder:

Input text strings are tokenized and padded or truncated to a fixed maximum sequence length.
The text tokens are passed through the T5 encoder model to produce dense contextual embeddings.
To ensure no attention leakage from padding, attention masks zero out unused sequence positions.
The clean textual embeddings are linearly projected into the transformer hidden dimension $d$ :
$\mathbf{C}_{text} = \text{MLP}_{text}(\text{T5Encoder}(\mathbf{text}))$

These language context embeddings $\mathbf{C}_{text}$ are subsequently fed into the cross-attention blocks of the Diffusion Transformer to steer action generation according to the task instruction.

Fig. 2. Block-causal self-attention mask. This matrix illustrates the block-wise routing logic dictating how tokens attend to one another within the multi-modal sequence. The layout is partitioned into a ground-truth anchor frame ( $F_0$ ), followed by temporally corresponding vision blocks ( $V_0 \dots V_3$ ), action blocks ( $A_0 \dots A_3$ ), and state blocks ( $S_0 \dots S_3$ ). This masking strategy ensures that information flows causally from the clean visual context to guide action denoising without leaking future states.

B. Spatiotemporal Token Routing (The Attention Mask)

The core architecture of ZULU directly adapts the routing strategy introduced by DreamZero. Rather than utilizing computationally expensive, fully bidirectional attention across long sequences, ZULU implements highly structured, block-causal attention mechanism. This design rigidly controls the flow of information between visual inputs, proprioceptive states, and action trajectories, ensuring that the model learns physical dynamics without leaking future information during action prediction.

1) Sequence Layout and Partitioning:

Before entering the attention layers, the unified sequence of tokens is logically partitioned into discrete spatiotemporal blocks. The layout consists of four distinct token groups:

Initial Anchor Frame ( $F_0$ ): The visual tokens comprising the very first frame of the input sequence.
Visual Blocks ( $V_0 \dots V_n$ ): The subsequent visual tokens are grouped into discrete blocks, with each block containing the patches for a parameterized number of consecutive frames.
Action Blocks ( $A_0 \dots A_n$ ): The continuous action tokens are grouped into blocks that temporally correspond to the visual blocks.
State Blocks ( $S_0 \dots S_n$ ): The proprioceptive state tokens are similarly grouped to correspond with the temporal blocks of the actions and visuals.

2) Block-Causal Self-Attention Rules:

To process these interleaved modalities, the CausalAttentionBlock defines a strict masking matrix governing which tokens act as queries and which serve as valid keys and values. The routing dictates that information flows predominantly from the clean visual context to the noisy action tokens. Based on the implemented mask, the self-attention routing strictly adheres to the following rules:

First Frame Anchor Isolation: Queries originating from the initial frame $F_0$ attend exclusively to keys within $F_0$ . This anchor establishes a static, uncorrupted ground-truth visual context for the rest of the sequence.
Visual Context Routing:
- Queries originating from a visual block $V_i$ can attend to the anchor $F_0$ , all historical and current visual blocks up to $V_i$ , their corresponding temporal action block $A_i$ , and their corresponding state block $S_i$ .
- Visual queries are strictly prohibited from attending to future visual blocks (e.g., $V_{i+1}$ ) or any future action and state blocks, preserving causal temporal dynamics.
Action Denoising Context:
- Action blocks consist of the noisy tokens that the flow-matching objective seeks to denoise. Queries from an action block $A_i$ attend to the anchor $F_0$ , the full visual history up to the corresponding visual block $V_i$ , their own current action block $A_i$ , and the corresponding state block $S_i$ .
- Action tokens are masked from attending to past action blocks (e.g., $A_{i-1}$ ) or any future information. This ensures the denoising of an action chunk is guided exclusively by the visual and physical state context of its specific temporal window.
State Token Isolation: Proprioceptive state tokens represent pure physical conditioning. A query from state block $S_i$ is completely isolated, attending exclusively to its own keys within $S_i$ .

3) Language Integration via Cross-Attention:

Following the block-causal self-attention, the sequence undergoes a cross-attention phase to integrate the language instructions.

Bidirectional Semantic Conditioning: Unlike the strict causal masking of the self-attention mechanism, the cross-attention block applies a fully non-causal mask.
Routing Logic: The entire updated multi-modal sequence (visual, action, and state tokens) acts as the queries, while the dense embeddings extracted from the text encoder serve as the keys and values. This allows every token in the spatiotemporal sequence to freely attend to the entire natural language instruction, ensuring that visual tracking and action generation are uniformly conditioned on the semantic goal.

4) Timestep Modulation (AdaLN):

Throughout both the self-attention and cross-attention operations, the flow-matching timestep is injected into the network via Adaptive Layer Normalization (AdaLN). The timestep embedding is transformed into six distinct modulation chunks ( $e_{chunk}$ ). These chunks apply scale and shift transformations to the sequence immediately before and after the attention and feed-forward layers, modulating the network's behavior based on the current noise level of the action tokens.

C. Flow-Matching Action Generation

Rather than treating action generation as a discrete classification or standard regression task, ZULU formulates continuous trajectory prediction through a flow-matching objective. Within this framework, visual and state tokens serve purely as uncorrupted conditioning context, while the continuous action tokens are the exclusive variables subjected to the forward noising and reverse denoising process.

1) The Forward Process:

During the training phase, the model is provided with ground-truth continuous action trajectories, denoted as $\mathbf{a}$ . To construct the flow-matching objective, these clean actions are perturbed by sampling a random noise tensor $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ of the exact same shape.

A specific timestep $t$ is sampled from the scheduler, which corresponds to a distinct noise level $\sigma_t$ . The forward process connects the empirical data distribution (the clean actions) to a standard Gaussian noise distribution via an optimal transport interpolation path. The noisy action tokens $x_t$ at a given timestep are generated using the following linear interpolation:

x_t = (1 - \sigma_t) \mathbf{a} + \sigma_t \epsilon

Once constructed, these noisy action tokens $x_t$ are encoded and injected into the Diffusion Transformer (DiT) alongside the clean visual and state context.

2) The Flow-Matching Objective:

The core objective of the DiT backbone is to predict the vector field that drives the transformation from the prior noise distribution back toward the clean data distribution. In this rectified flow formulation, the true vector field (the training target) is explicitly defined as the difference between the sampled noise and the original clean action trajectory:

v_t = \epsilon - \mathbf{a}

The DiT processes the multi-modal sequence and outputs a predicted vector field $\hat{v}_t$ for the action tokens. This prediction is dynamically conditioned on the clean visual patches, the physical state tokens, the embedded language instructions, and the encoded flow timestep $t$ .

3) Loss Computation and Masking

The optimization metric for action generation is calculated using a Mean Squared Error (MSE) loss between the network's predicted vector field $\hat{v}_t$ and the ground-truth target $v_t$ . To balance learning across various stages of the noise schedule, this per-sample MSE loss is multiplied by a timestep-dependent weighting array provided by the flow scheduler.

4) Inference and Vector Field Integration

During inference, ground-truth action data is naturally unavailable. Action generation begins by initializing the entire action sequence horizon as pure Gaussian noise.

To integrate the network's predictions, ZULU relies on a specialized multi-step solver, Flow UniPC (Unified Predictor-Corrector) Multistep Scheduler. Prior to sampling, the scheduler calculates a discrete sequence of inference timesteps by interpolating noise levels (sigmas) and applying a continuous shifting transformation—dynamically warping the schedule to allocate integration steps optimally across the flow trajectory. At each discrete step, the DiT predicts the vector field conditioned on the current noisy action states and the streaming, clean visual context. Rather than a simple linear integration, the solver utilizes a multi-order predictor-corrector mechanism. It maintains an internal buffer of historical vector field predictions and employs a warmup phase, gradually increasing the solver's integration order as more historical data becomes available. At each timestep, the scheduler first executes a corrector phase—refining the current trajectory sample using the newly predicted vector field and the previous sample—before executing a predictor phase to propagate the trajectory to the next noise level. This iteratively pulls the noise towards the true action distribution to form a smooth, continuous physical trajectory.

III. Training Setup

ZULU is trained using a specialized subset of robotic manipulation data sourced from the DreamZero-DROID-Data-2000 dataset formatted under the LeRobot v2.1 structure. This dataset comprises 2,000 demonstration episodes focused on real-world household chores and pick-and-place tasks.

A. Dataset and Physical Embodiment

1) Multi-View Visual System

The dataset provides synchronized multi-camera RGB video captured at 15 FPS with an original resolution of $320 \times 180$ pixels. Observations are collected across three distinct camera perspectives:

- Exterior Left 1: Side-angle contextual view of the workstation.
- Exterior Left 2: Secondary side-angle view providing alternative spatial depth.
- Wrist Left 1: Hand-centric camera attached to the arm for close-up manipulation tracking.

2) Robot Embodiment

The physical demonstrations are executed on a Franka Emika Panda robotic arm equipped with a parallel-jaw gripper:

- State Vector: Captures 7-DoF joint positions combined with the 1-DoF gripper state ( $S_{horizon}$ ).
- Action Trajectory: Represents joint position targets and gripper commands chunked over a forward trajectory window ( $A_{horizon}$ )

B. Data Preprocessing and Transformations

To prepare raw observations for the multi-modal Diffusion Transformer, data samples undergo a standardized processing pipeline:

- Visual Augmentations: Camera frames are randomly rescaled with a factor of 0.95 and resized to a target tensor dimension of $160 \times 320$ pixels. Mild photometric color jittering (adjusting brightness, contrast, saturation, and hue) is applied across all camera streams simultaneously to improve generalization against lighting shifts.
- Quantile Normalization: Both state and action vectors are normalized using $q_{99}$ quantile normalization, mapping physical joint values into a bounded range while mitigating the impact of extreme mechanical velocity spikes.
- Relative Actions: Action joint targets are formulated relative to the robot's current joint positions, ensuring the network models differential joint trajectories rather than absolute spatial coordinates.
- Multi-modal Stacking: Visual views are concatenated along the spatial sequence dimension, while states, actions, and text annotations are collated into unified batched features.

C. Training Pipeline and Optimization

1) Optimization Hyperparameters

- Optimizer: AdamW with modified momentum parameters ( $\beta_1 = 0.95$ , $\beta_2 = 0.95$ ), weight decay set to $0.01$ , and numerical stability epsilon $\epsilon = 1 \times 10^{-8}$ .
- Learning Rate Schedule: Controlled via a cosine decay learning rate scheduler with an initial warmup ratio of $0.05$ .
- Compute Precision: Executed natively in bfloat16 to maximize throughput on modern GPU accelerators while maintaining numerical stability for flow-matching timesteps.
- Distributed Batching: Configured with a fixed global batch size of 32, dynamically distributed using gradient accumulation steps across physical training devices.

2) Loss Formulation

The model is optimized end-to-end using a joint objective function:

\mathcal{L}_{total} = \lambda_{action} \cdot \mathcal{L}_{action} + \lambda_{dynamics} \cdot \mathcal{L}_{dynamics}

Where $\mathcal{L}_{action}$ is the flow-matching Mean Squared Error computed over the predicted action vector fields, and $\mathcal{L}_{dynamics}$ is the residual next-frame MSE loss evaluated over predicted DINOv2 visual patch deltas.

D. Current Status, Challenges, and Future Directions

As this technical report presents the first exploratory iteration of the ZULU framework, the primary goal was to validate the structural token-routing and flow-matching mechanics. This section details the current physical footprint of the architecture, the engineering hurdles encountered during the initial training runs, and the roadmap for future scaling.

Fig. 3.1. ZULU action trajectory predictions on the training dataset. True Model MSE: 0.00348; True Model L1: 0.04326.

Fig. 3.2. ZULU action trajectory predictions on the test dataset. True Model MSE: 0.00639; True Model L1: 0.05655.

1) Model Scale

The first iteration of the ZULU architecture is deliberately constrained to serve as a lightweight proof-of-concept. The Diffusion Transformer (DiT) backbone is constructed with the following structural dimensions:

- Layers: 12
- Attention Heads: 16
- Hidden Dimension: 1024
- Feed-Forward Dimension: 4096
- Input/Output Dimension: 768
- Frequency Dimension: 256

This configuration results in exactly 216 Million trainable parameters. When accounting for the frozen spatial priors (the DINOv2 base) and the frozen semantic conditioning (the FLAN-T5 text encoder), the total system footprint reaches 412.21 Million parameters. By heavily leveraging these pre-trained frozen representations, ZULU successfully demonstrates that complex spatiotemporal token routing and flow-matching can be executed within a highly accessible parameter budget, avoiding the massive hardware overhead of learning visual physics from scratch.

2) Future Directions

As an open-source experiment, ZULU architecture lays the groundwork for several advanced optimizations and scaling opportunities.

- Training Optimization using DeepSpeed: To move beyond the limitations of standard PyTorch DDP, future iterations will benefit heavily from integrating Microsoft's DeepSpeed specifically the Zero Redundancy Optimizer (ZeRO). By partitioning optimizer states, gradients, and model parameters across GPUs (ZeRO Stages 1 through 3), the VRAM footprint per device could be drastically reduced. This would allow researchers to significantly increase the global batch size, expand the context window, or scale the DiT parameter count (e.g., to 1B+ parameters) without requiring extreme-memory accelerator hardware.
- Optimization of Attention and Visual Context Mechanisms: Currently, ZULU passes all extracted DINOv2 patches directly into the transformer, resulting in a dense and computationally heavy sequence. Future iterations should explore visual context optimization strategies to reduce this burden.
- Scaling for Long-Horizon Tasks and Reasoning: The current architecture operates on a short, action horizon which is ideal for immediate physical dynamics but insufficient for complex, multi-stage reasoning. Future expansions of the ZULU paradigm will focus on long-horizon task execution. Expansion of temporal context window and integrating history-aware memory tokens would further enable the model to handle tasks requiring prolonged physical reasoning and multi-step planning.

E. Conclusion and Discussion

In this technical report, I explored an alternative approach to Vision-Language-Action modeling by introducing ZULU, a 216M-parameter architecture. Rather than relying on massive generative vision-language backbones, ZULU successfully demonstrates how to ground a flow-matching Diffusion Transformer directly in pre-trained spatial priors like DINOv2. By engineering a custom block-causal attention mask, the model strictly manages the temporal routing of interleaved images, continuous actions, and proprioceptive states. While this report presents an early-stage exploration rather than a fully converged foundation model, I open-surce the complete architectural concept and implementation. My hope is that this structural blueprint provides a valuable starting point for further researches, looking to explore generative, flow-matching robotic control.

Currently, the model remains severely undertrained due to slow training throughput, requiring significantly more convergence time to reach full optimization. Nevertheless, qualitative closed-loop evaluations of an early checkpoint at 16,000 training steps reveal promising indicators for future exploration. Notably, during closed-loop execution, the robotic arm exhibits physically stable continuous trajectories without dangerous vibrations, rapid oscillation, or continuous jerking—a common failure mode in under-converged policy networks. While these initial robot movements remain small in magnitude and are not yet sufficient to complete full task, they demonstrate meaningful, non-random directional control aligned with visual context.

To explore domain adaptation, I conducted a fine-tuning experiment on this 16,000-step checkpoint using Low-Rank Adaptation (LoRA). The model was fine-tuned for 5,000 steps on a custom dataset of 56 pick-and-place demonstration episodes collected via Robosuite teleoperation. Although task execution performance remains far from optimal, the fine-tuned model consistently preserved non-erratic trajectory output.

Due to severe compute and hardware resource constraints, this preliminary report cannot provide exhaustive benchmark evaluations, formal quantitative metrics, or extensive ablation studies. However, the emergence of relatively smooth, physically coherent motor steps despite severe undertraining suggests that driving a block-causal, flow-matching Diffusion Transformer with pre-trained vision representations is an architectural direction well worth pursuing with scaled compute.