A long-standing goal of robotics is to develop versatile systems that perform effectively across diverse settings, applications, and hardware configurations. This makes deployment easier, improves model robustness, and enables transfer learning between domains. We’ve had tremendous success achieving this with foundation models for language. LLMs can be applied to many different problems and easily adapted with limited task-specific data, benefiting from learning across domains. We seek to achieve this same type of foundation model for embodied AI.

This blog post focuses on how foundational LLMs are being adapted for robotics applications. We will make frequent analogies to show how robotics research aligns with LLM research, cover the unique challenges of robotics and the adaptations needed to address them, and explore how the fields have become tightly coupled so that progress in one field benefits the other.

Looking toward the future, we should consider what form AGI will take. Many researchers believe intelligence can be built in the digital world through LLMs. This means agents for coding, research, and knowledge work. Many believe real-world tasks like robotics will be easily solved once AGI is achieved in text. There’s an alternate view that true intelligence requires understanding the physical world, and that AI shouldn’t be constrained to operating through digital chat interfaces. In building embodied AI, vision-language-action models represent another avenue that complements rather than competes with LLM research.

Data

The prevailing narrative is that data scarcity is the primary bottleneck holding robotics back from achieving the success we’ve seen in language applications. While hardware and algorithmic advances remain important, limited data availability represents the current constraint on robotics progress. To understand the data constraints of robotics, we need to examine the type and scale of data needed to train effective robotics models in comparison to LLMs.

LLM Training Scale

Let’s first establish the data scale used for LLM/VLM training so we can then understand the data gap of robotics.

LLM pretraining operates on trillions of tokens. Qwen2.5 trains its base model on 18 trillion tokens. The training follows a clear hierarchy: pretraining uses internet-scale data, while post-training stages use progressively less but higher-quality data. Qwen2.5 uses approximately 1 million examples (32 billion tokens with 32k sequence length) across all post-training phases, demonstrating how each successive stage requires orders of magnitude less data but with better task alignment.

VLMs follow similar patterns but with less transparency about exact data volumes. Kimi-VL reports using trillions of vision tokens, fewer than text tokens but still at massive scale. While they don’t specify image counts or video durations, the scale likely reaches billions of images.

Token counts provide only rough data comparisons since they depend heavily on tokenizer vocabulary size and encoding schemes. A tokenizer with 150k vocabulary packs more information per token than one with 30k vocabulary. For vision data, token counts depend on ViT patch sizes, input resolution, and video frame sampling rates, making direct comparisons even more problematic.

Fine-tuning data varies widely in scale. Pretraining makes these models sample efficient, so most use cases require only billions of tokens or less. Some applications, however, use substantially more data. For example, Qwen2.5-Coder constructs a code-heavy dataset of 5.5 trillion tokens for continued pretraining on top of the Qwen2.5 base model. The general principle is that models are trained on as much data as possible. Despite this variation, billions of tokens is typically sufficient to enable strong LLM performance in text domains.

Robot Data Scale

The data for training robots mainly consists of videos accompanied by robot actions. Actions are typically floating-point values to map to low-level commands to control the hardware. These robot “episodes” can also contain additional sensor data or multiple camera views. This data is diverse due to different robot embodiments, each with varying sensors, action spaces, and capabilities. The goal of a foundation model is to train a single model that can be deployed across many robots, requiring a large dataset that contains data from a diverse set of robots.

One of the most significant efforts in building this kind of dataset is Open X-Embodiment. This collective effort from 21 academic and industrial institutions amalgamates 72 different datasets from 27 different robots, covering 527 skills across 160,266 tasks. The dataset standardizes data from diverse robot types (single arm, double arm, wheeled robots) with varying sensors and action spaces into a unified format.

If we want to compare the training scale directly with VLMs, we have to rely on token count. It’s difficult to quantify the scale of videos that VLMs train on since they only report total token counts. These total token counts are not useful without knowing the frame sampling frequency and tokens per image.

VLMs are trained on trillion token internet-scale data. It is possible to get competitive in token count with robotics data. Open X-Embodiment contains 2.4 million episodes. If we assume 30 seconds per episode, 30 Hz frame sampling, and ~512 vision tokens per frame, we can get over a trillion tokens. This token count also doesn’t account for additional sensor and camera views, and the action tokens. However, this token count shouldn’t be compared with the trillions of tokens that LLMs train on. The parameters used in this calculation are estimates only. Also, there is a lot of redundancy in these video tokens. We should expect a trillion video tokens to have less signal than a trillion text tokens.

Rather than token counts, we should focus more on the quantity and diversity of video. It is hard to say how many videos the largest VLMs are trained on, but we do know that billions of general videos are available, far more than robot videos. This is because there are far more videos available through internet-scale pretraining than there are robot videos. The way to address this is to utilize non-robot videos to take advantage of transfer learning. The diversity of robot data is limited. While the 27 different embodiments represent unprecedented variety in robotics research, this is still minuscule compared to the diversity found in internet-scale datasets. One interesting avenue to address limited embodiment diversity is training in simulation. Skild AI reports training on 100,000 different robots generated in simulation.

These datasets are also limited in the diversity and coverage of tasks. For example, simple tasks are overrepresented while highly dexterous tasks are less covered.

Robotics datasets are large and are growing. However, they are limited in the diversity of embodiments, environments, and tasks. This is a fundamental challenge in building a general-purpose foundation model for robotics.

Transfer Learning

LLM applications utilize three types of data in a hierarchical structure, with each level decreasing in availability but increasing in quality and task alignment: internet-scale text data for pretraining → instruction tuning data → task-specific data. Robotics follows a similar data hierarchy: internet-scale video data → ego-centric video data → robot data with actions → specific robot data with actions. There are some efforts to collect large amounts of ego-centric data (Ego4D), as this kind of data is easier to collect than robot data (can use smart glasses). LLM researchers developed these training recipes due to the limited amount of high-quality instruction-following data that is available. We can follow this pattern for robotics.

These larger datasets used to train LLMs/VLMs are meant for models to learn an entire data modality (text, images). Robotics data uses the same modalities of data. The only new modality is actions, which is why we need to fine-tune to create VLAs. This is also why we want to leverage language and vision models that are trained on massive amounts of image/text data. Robotics data can be viewed as a type of fine-tuning data.

However, maybe we should consider interacting with the physical world to be a completely new domain. This means it requires a lot of data to learn, and transfer learning from digital tasks is limited. Then it makes sense to compare the scale of action data with the scale of image or text data. However, this is largely a philosophical question. The empirical reality is that robotics is far more data hungry than a typical LLM application.

Scaling Laws

If we assume that scaling laws remain consistent across domains, model performance should improve with increased data volume. This highlights the need to train on the largest possible datasets. Qwen2.5-Coder demonstrates this principle by training on 5.5 trillion tokens, significantly more than typical fine-tuning applications, to achieve strong coding performance. Language model scaling laws are logarithmic, meaning performance scales linearly with the logarithm of training tokens rather than the raw token count.

This work by Lin et al. 2025 argues that scaling laws in robotics depend more on the number of environments and objects rather than the number of demonstrations. A large amount of robotics data is collected in constrained lab environments. This suggests that large-scale deployment and data collection would unlock significant gains. It would also be interesting to study the scaling laws for general-purpose models on the number of embodiments.

Robotics is considered data-constrained because the amount of robotics data readily available is orders of magnitude less than what’s needed to reliably deploy general-purpose robots. This constraint exists because robot data collection is expensive and time-consuming, there is limited robotics data currently available compared to other domains, and robot intelligence appears to require substantial amounts of data to achieve reliable performance.

Data Collection

The fundamental challenge with robotics data collection is that it requires large-scale deployment of robots to generate sufficient training data. However, robots need to be effective and affordable enough to enable broad deployment in the first place. Many companies are betting that leveraging pretrained VLMs will make robots capable enough to reach this large-scale deployment threshold.

Large-scale deployment alone may not be sufficient to solve robotics. Tesla has had large-scale deployments of cars collecting massive amounts of data for years, yet has only recently started deploying autonomous vehicles successfully.

Despite these challenges, many companies are investing heavily in data collection to bridge the gap toward large-scale deployment. Figure is partnering with a real estate company to conduct large-scale humanoid robot data collection in thousands of different homes. This approach is expensive but may be necessary to achieve functional general-purpose robots.

LLMs for Robotics

We first consider how robotics can be treated as an application of LLMs. We will focus on Gemini Robotics, which applies a SOTA LLM for robotics without significant changes in architecture or training.

First, let’s consider what properties we need from an LLM in order for it to be useful for robotics. If we want to power robots with AI models, we need two capabilities. Embodied reasoning is the ability of models to understand the physical world and how to interact with it. We also need models capable of embodied action. This means that the model can use its embodied reasoning to interface with hardware and perform complex actions in the real-world.

The benefits of applying LLMs are that these models come pretrained with immense world knowledge. They also display common high-level intelligent capabilities such as reasoning and long-term planning. The goal can then be simplified to just adapting LLMs for robotics, which we expect to be simpler and more data efficient than training from scratch.

Embodied Reasoning

In the Gemini Robotics paper, the authors define embodied reasoning as follows.

“The set of world knowledge that encompasses the fundamental concepts which are critical for operating and acting in an inherently physically embodied world.” - Gemini Robotics

This is an ability of vision language models (VLMs) and is not necessarily tied to robotics. VLMs refer to multimodal language models that can process visual and text inputs. Testing embodied reasoning simply involves prompting VLMs about images. Classical computer vision tasks like object detection and multi-view correspondence fall under embodied reasoning. These tasks are all expressed as language prompts.

Embodied reasoning can also be tested through visual question answering. These questions test the understanding required to interact with the environment.

In addition to general physical reasoning, we can take advantage of world knowledge to make decisions. For example, we could ask a robot to pick up a healthy snack in the kitchen. The world knowledge in the VLM would be used to figure out how to execute this ambiguous command.

Applying VLMs means that robotics can take advantage of advances in LLM capabilities. For example, we can take advantage of the impressive recent progress made in LLM reasoning and test-time compute. We can prompt the VLM to output a reasoning trace which allows it to better solve an embodied reasoning problem.

Gemini Robotics introduces an Embodied Reasoning QA (ERQA) benchmark that measures the embodied reasoning capabilities of LLMs. This benchmark requires no fine-tuning and works with any VLM. For models intended for embodied AI applications, performance results on this or other embodied reasoning benchmarks (EmbodiedBench, EAR-Bench) should be reported.

Embodied Action

Embodied reasoning enables models to understand the physical world and reason about objects, spatial relationships, and physical interactions. For robotics applications, we want to leverage this understanding to enable meaningful actions in the real-world. This means translating high-level understanding into precise control commands through the robot’s hardware APIs. Each robot has a different interface, and the knowledge of how the robot is controlled is not present in the VLMs.

When applying LLMs to new domains, we typically follow a progression of increasingly complex methods: zero-shot prompting, few-shot prompting/prompt engineering, fine-tuning, and reinforcement learning. The application of LLMs to robotics follows this same pattern.

Zero-Shot Control

One way to express robot control is by specifying an API with code. The context to the VLM will include the documented API. The robot can reason based on this API and produce code that would execute a command.

Levine et al. 2025 explore this in more depth. Reasoning improves performance on robotics tasks as well as the interpretability of the actions.

Few-Shot Control

Zero-shot control works well when the API is high-level. For high-level APIs, we may have another model generating the low-level actions. However, if we want the VLM to output low-level actions, we need to give some examples. A few-shot prompt that includes examples of how to use low-level commands to complete actions would give the model the necessary information. It would then extrapolate based on these examples on how to execute novel commands.

Fine-tuning

Similar to traditional LLM applications, we can enhance VLMs for robotics by fine-tuning them on action data, creating what are known as vision-language-action (VLA) models. This approach bridges the gap between understanding the physical world and executing precise actions.

The fine-tuning process follows a supervised learning paradigm where human operators manually demonstrate tasks with robots. These demonstrations generate datasets of visual observations paired with corresponding low-level action trajectories. This approach is known as behavior cloning or imitation learning in robotics, but can also be considered a form of supervised fine-tuning (SFT) process as used in LLM training.

Modern implementations like Gemini Robotics employ a distributed architecture that balances capability with operational efficiency. A large VLA model runs in the cloud for high-level reasoning and planning, while a compact action decoder resides on the robot itself, translating model outputs into low-latency control signals.

Given the current research focus on fine-tuning foundation models for robotic applications, the remainder of this blog will explore VLAs in greater detail.

Reinforcement Learning

A natural extension of this approach is training VLAs with reinforcement learning. While RL for LLMs has become a very active research area recently, we see limited adoption of these methods for VLAs. This limitation exists because robots must physically perform actions to generate training data. In contrast, methods like RLHF/RLVR allow LLMs to generate responses in various contexts that are automatically verified by another model or algorithm. This online learning process scales easily by running multiple LLMs in parallel.

However, scaling online learning becomes significantly more challenging when it requires physical actions in the real-world. Researchers have attempted to address this by creating environments where multiple robots can interact simultaneously. For example, Google built a robot arm farm in 2016. Yet in the context of modern deep learning, a scale of 14 robots is minuscule. We would need millions. We may be able to achieve this once robots are able to be mass-deployed. Then there will be a new set of challenges if we want to do distributed online reinforcement learning.

There is a rich body of literature to work around the challenge of scaling online learning. The robot policy can be learned in simulation or in the real-world (which is more expensive and slower).

We have seen a lot of progress in applying RL on LLMs to improve performance on verifiable tasks such as math and coding. The advancements here may translate into robotics. The main gap is that we can’t quickly execute a trajectory in the real-world.

With the SFT approach, the robot policy is limited by the data it is fine-tuned on, since it is optimized to imitate it. With RL, the policy is optimized directly on task performance and can learn trajectories that are better than what is in the datasets. We don’t want robot performance to be limited by how well humans perform with teleoperation. In order for RL to be applicable for real-world actions, this likely requires a strong base policy trained through SFT along with sample efficient RL methods.

Synthetic Data

LLM training is increasingly relying on synthetic data generation. World models can be thought of as the source of synthetic data for the real-world. World models are another big area of research relevant for robotics. These models can be thought as the source of synthetic data for robotics that can be used both for training/eval data generation, and RL environment creation.

World models generate synthetic training data by predicting future video frames given robot actions. Sim-to-real transfer is a huge area of research in robotics, where the goal is to train policies in simulation and deploy them in the real-world. World models enable generating realistic simulations of different real-world environments, which are becoming indistinguishable from real data.

1X uses world models for evaluation in simulation. However, current world models suffer from high computational latency and accumulate errors over longer prediction horizons. There are some efforts like V-JEPA 2 that aim to use world models in the policy itself, but this is limited by the latency of the world model.

Language in Robotics

Do we need robots to be language promptable? A general-purpose robot needs to understand language. Without language the robot would be restricted to a subset of possible actions. For example, you can have a button for washing the dishes and one for doing the laundry. But eventually you’ll want the robot to do a task for which there is no button.

Language also serves as a source of world knowledge. When you ask a robot to “prepare a healthy snack,” the robot can leverage knowledge about nutrition and food categories from its pretrained language model without explicit programming for every food item.

I would argue that language understanding is a necessary component for a general-purpose robotics foundation model. However, it is an open question on whether the LLM has to be the main component in the model. For example, we may one day build powerful vision models for robotics, and add a small language component. Whereas with VLAs, we do the opposite.

The necessity of language for general-purpose robotics is one reason why many are betting on VLAs.

VLA Architectures

It is possible to fine-tune an LLM to predict robot actions directly. This approach treats actions as discrete tokens and formats input data to match the model’s requirements. However, specific architectural changes can significantly improve these models’ effectiveness.

As we explored in my VLM blog post, enabling LLMs to understand images requires architectural adaptations. Similarly, getting these models to understand and perform actions demands additional modifications. We will now examine the common architectural changes used to build vision-language-action models.

Discrete Actions

LLMs are built to predict discrete tokens. The simplest implementations of VLAs discretize actions so the model can predict them as tokens in its vocabulary. The actions are trained and decoded autoregressively, with discretized tokens added to the model’s vocabulary. During inference, the output is detokenized, meaning discrete tokens are mapped back to continuous actions for robot execution.

RT-2 and OpenVLA are two works that implement discretized autoregressive action training. Although it’s possible to represent actions as text using the normal text tokenizer (treating this as standard SFT), it’s more effective to treat actions as distinct tokens.

The autoregressive loss is typically applied only to the actions, with visual and text tokens treated as a prefix. This approach constrains the model to predict only actions rather than text, enabling a smaller output layer and more efficient cross-entropy calculation. However, it would be possible to simultaneously fine-tune on the text prompt by running SFT next-token prediction on both the prompt and actions. When fine-tuning only on actions, the output computation is smaller since cross-entropy operates over the action space rather than the full text vocabulary. In practice, it seems like it is more common to run SFT on vision-language data separately from action fine-tuning.

Continuous Actions

Although discretization is the most natural way to get VLMs to predict actions, this comes at a cost. Discrete tokens can’t effectively model smooth continuous trajectories so the resulting robot behavior may seem unnatural. It also limits the granularity in which we can control the robot, imposing an artificial limitation on the robot’s abilities. If we want higher frequency control for dexterous tasks, discrete tokens will become highly repetitive and lacking in the ability to represent subtle movements.

Discretization also introduces a tradeoff between output vocabulary size and action granularity. We could increase the number of discrete buckets, but this can also make it more difficult to train with limited data. Since actions are inherently continuous, it would be better to have model architectures that can predict continuous actions directly. Continuous action prediction enables finer grained actions that allow for high frequency dexterous tasks.

A naive solution is to fine-tune a VLM to predict continuous actions using a linear layer on top of the output embeddings with MSE loss. However, this approach is not effective when outputs have multiple modes. An MSE loss would train the model to average these modes, and this average is not a good action.

Some works, such as Liang et al. 2025, train VLAs with regression loss to predict continuous actions. However, this approach is not as effective on its own compared to discretization-based methods. They need to co-train with discrete action prediction for this to be effective.

Flow Matching / Diffusion

Flow matching and diffusion are the two most effective ways to train transformers to generate continuous outputs. These methods were first researched for image generation, but are well applied to action generation as well. We see some VLA architectures use flow matching while others use diffusion. Flow matching and diffusion are equivalent methods so we will not concern ourselves with the distinction.

The \(\pi_0\) model uses flow matching. Flow matching training involves sampling noise and time, and training a vector field to predict the direction to update the noise to get to the target. Inference involves taking multiple steps to integrate on the time to go from noise to a predicted final target. Diffusion is similar but is framed as removing a step of noise from an image.

def flow_matching_training(
        observation, # contains text prompt and latest images
        true_actions # next chunk of continuous actions
    ):
    # Sample random noise and timestep
    noise = random_normal(shape=true_actions.shape)
    t = random_uniform(0, 1)

    # Create noisy interpolation between true_actions and noise
    # t=0: pure true_actions, t=1: pure noise
    x_t = t * noise + (1 - t) * true_actions

    target_velocity = noise - true_actions  # Points FROM actions TO noise

    # Forward pass through vision-language-action model
    vision_lang_tokens = encode_vision_language(observation)
    action_tokens = encode_actions_with_time(x_t, t)
    predicted_velocity = transformer([vision_lang_tokens, action_tokens])

    # Train to predict the velocity field
    loss = mse(predicted_velocity, target_velocity)

def flow_matching_inference(observation, num_steps=10):
    x = random_normal(action_shape)  # Start from noise
    dt = 1.0 / num_steps

    # Iteratively follow the flow: t goes from 1 → 0
    for step in range(num_steps):
        t = 1.0 - step * dt

        # Get velocity prediction at current state and time
        vision_lang_tokens = encode_vision_language(observation)
        action_tokens = encode_actions_with_time(x, t)
        velocity = transformer([vision_lang_tokens, action_tokens])

        x = x - dt * velocity  # Follow flow BACKWARD (noise → actions)

    return x  # Continuous actions

Flow matching is more effective at generating continuous actions. The main drawback is that it is slower to train than autoregressive architectures. Flow matching can result in faster inference than autoregressive prediction since we can predict a whole chunk at once. To generate a chunk of 50 actions, we may need to only run 10 inferences for flow matching. This is faster than decoding 50 steps autoregressively. We can parallelize inference with autoregressive models. This is only in the case that the number of flow steps is less than the number of steps in the chunk.

Action Expert

Although it’s possible to fine-tune a VLM directly to predict actions, VLAs empirically perform better when using separate weights for action learning. This approach gives the model additional capacity to learn action outputs without interfering with the VLM’s existing capabilities.

These specialized weights can be added on top of the transformer, such as a multilayer MLP instead of a single output projection layer. However, incorporating these extra parameters into the transformer layers themselves proves more effective.

\(\pi_0\) adds a 300 million parameter action expert to a PaliGemma VLM. The model’s inputs are divided into a prefix and suffix. The prefix contains vision tokens for recent image frames and text prompts, which can be processed directly by the pretrained VLM. The suffix adds robotics inputs: \(q_t\) representing the proprioceptive state (current position of robot), and noise (interpolation between noise and actions for flow matching).

The action expert is a transformer model with the same number of layers but smaller embedding dimensions and MLP widths. The attention heads and per-head embedding dimension must match the main model to allow prefix tokens in the attention mechanism. When processing, suffix tokens go through the action expert transformer while incorporating the KV embeddings from the prefix (these are computed once then cached). The prefix uses full self-attention, while the suffix uses causal attention. This makes sense since we don’t need to predict tokens in the prefix. This smaller action expert architecture enables faster inference while still leveraging the capabilities of the larger VLM.

Various modifications to this architecture are possible. For instance, SmolVLA alternates between self-attention among action tokens and cross-attention with prefix tokens. This approach further reduces computational requirements compared to the semi-causal full attention mechanism used in \(\pi_0\).

Action Chunking

We see in the \(\pi_0\) and many other VLA architectures that we predict multiple tokens at the same time. This is called action chunking and was introduced by Zhao et al. 2023. \(\pi_0\) implements this by running flow matching on multiple action tokens in parallel. Since flow matching is an iterative process at inference time, we get an efficiency gain from this parallelization. This is similar to multi-token prediction used in DeepSeek-V3.

At inference time, the robot is continuously executing actions at 50 Hz. The model operates at 50 Hz with action chunks of size 50 to predict a full second of action. There are two inference strategies. In the synchronous strategy, we cycle through a process of generating a chunk, executing the chunk, and collecting observations for the next chunk. This means that during inference we introduce a delay where the robot is not performing any action.

We can also implement a real-time strategy. To avoid inference delays, we have to introduce lag in the observations. For example, every 0.5 seconds, while the robot continues executing the current action chunk, the model predicts the next 25 actions based on a new observation. The inference time is only 73ms for on-board inference (or 86ms for off-board inference), which is well within the 0.5 second window. However, this means the observations used for planning are always 0.5 seconds stale. By the time the robot starts executing actions 26-50, those actions were planned based on what the robot observed 0.5 seconds ago, before it executed actions 1-25.

Chunking is effective since actions have a lot of redundancy. This means that these consecutive tokens are highly correlated so they can usually be predicted in parallel. If the actions varied drastically between time-steps, action chunking likely wouldn’t be stable.

The original work (Zhao et al. 2023) also evaluates temporal ensembling, which is later found not to be effective in \(\pi_0\). This involves introducing an overlap between chunks and taking the average of the predicted actions. However, incorporating temporal ensembling could be useful if we want to increase the frequency of observations.

Real-Time Action Chunking

The main problem with action chunking is that we either have to have inference delays or introduce lag in the observations. If we use real-time with observation lag, the model doesn’t see the result of the second half of the latest chunk while planning the next chunk, since those actions haven’t been executed yet. This can cause discontinuities in the action that lead to jerky, unnatural movement at chunk boundaries. This delay can also be exacerbated by latency in retrieving the observations. It is possible to address this issue by adding the unexecuted actions to the input of the model, but the \(\pi_0\) model doesn’t seem to do this.

Black et al. 2025 proposes an inference time solution that does not require training changes. The main goal is to improve consistency between chunks as simply smoothing over the discontinuities wasn’t working well enough. They treat generating new action chunks like an image inpainting problem. This is another area where VLA research has taken advantage of image generation.

We can first introduce some terminology to define the inference process. \(H\) (prediction horizon) represents how many actions are predicted in each chunk. \(s\) (execution horizon) represents how many actions we execute from each chunk before starting inference for the next chunk. \(d\) (inference delay) represents how many controller timesteps it takes to generate a chunk. If we have 73ms delay at 50Hz, \(d\) is 3 timesteps.

With \(H=50\) and \(s=25\), if we start inference after executing \(s\) actions, then by the time Chunk 1 is ready (\(d\)=3 timesteps later), we’ll have executed \(s+d=28\) actions from Chunk 0. This means when we eventually switch to Chunk 1, the first d=3 actions need to be frozen to match actions 25-27 from Chunk 0. However, actions 28-49 can be edited while we predict Chunk 1. We can consider 3 regions:

1. Frozen region (first \(d\) actions): Weight = 1. These actions from the previous chunk will have already executed by the time inference finishes, so they must be frozen to match what actually happened.
2. Intermediate region (actions \(d\) to \(H-s\)): Exponentially decaying weights from 1 to 0. These actions belong to the current chunk. The exponential decay encourages smoothness between chunks.
3. Fresh region (last \(s\) actions): Weight = 0. These actions are beyond the end of the previous chunk and need to be freshly generated with no constraints.

These weights are used in pseudoinverse guidance (ΠGDM) adapted for flow matching. At each denoising step during action generation, a gradient-based guidance term encourages the final generation to match the frozen and intermediate values from the previous chunk. We use the weights on this term to determine how close we want the action to be to the previously generated action.

Hierarchical

We want these robot systems to be able to handle complex and even ambiguous commands. For example, we may want to prompt a robot with “organize my kitchen”. Such commands require longer term planning and breaking it down into lower level tasks. The goal of hierarchical systems is to have one system propose low-level instructions that another model executes.

The high-level policy takes a command and outputs both a response to the user and instructions for the low-level policy. For example, if you prompt the model to make a sandwich, the high-level policy will generate a sequence of actions like “1. Pick up bread 2. Place bread on plate…” The low-level policy then executes these commands through precise physical actions. This allows the high-level policy to focus on planning complex, long-horizon tasks.

\(\pi_0\) does an evaluation where they break apart complex high-level commands to lower level simpler text commands, using both expert humans and another VLM. They find that \(\pi_0\)-HL and \(\pi_0\)-Human both improve the performance on complex tasks, compared with giving \(\pi_0\) the complex prompt directly. This motivates further research into hierarchical systems where one model processes complex text prompts

In the Hi Robot work, the authors build this kind of hierarchical system with a high-level policy generating commands to a low-level policy. The high-level policy is just a VLM that outputs text, while the low-level policy generates actions.

This hierarchical approach enhances the system’s ability to handle longer-term and more complex tasks, while also improving interpretability by making the high-level policy outputs readable. The approach more than doubles performance on certain open-ended tasks such as table bussing and sandwich making, compared to \(\pi_0\).

One gap with breaking part a high-level task into low-level commands is that the low-level policy loses global context. Sometimes different subtasks can be partially parallelized. For example if a robot is going to retrieve a tool for a subtask, it may save time by also picking up a tool for the next subtask. An alternative architecture could be to have a single model that is trained to first output text (similar to a reasoning trace in LLMs) and then generate actions. This can be considered a combined VLM/VLA architecture. However, this means that the high-level policy VLM has to be the same as the low-level policy’s, whereas we would prefer to use a larger model for the high-level policy.

\(\pi_{0.5}\) implements this kind of combined model. It is trained to predict a single subtask at a time based on the observation. It runs an inference loop where it generates a subtask based on the observation, executes the subtask to completion, and then generates a new subtask based on the latest observation.

Although this approach is effective in executing complex multi-step tasks, it is limited in the fact that it greedily only generates one subtask at a time. Some tasks may benefit from generating a complete plan of subtasks before executing. However, we may also want to update the plan as we complete subtasks. In a follow-up work, they introduce a stop grad between the action expert and the VLM backbone. The idea is that the VLM parameters should only be trained to output text. Gradients from predicting continuous actions can interfere with the VLM’s abilities. This is considered knowledge insulation. However, we still want the VLM to learn from the actions, so it is co-trained to predict discretized actions.

Training effective VLAs requires balancing different capabilities: continuous action prediction, discrete action understanding, and the semantic comprehension of the VLM.

Model Size

GR00T N1 and \(\pi_0\) both use 2B parameter LLMs in their VLA architectures. Small models are required to enable on-device inference and real time latency. Some hierarchical approaches (Figure Helix uses a 7B VLM) can use larger models for the lower frequency component. Gemini Robotics uses a lower latency LLM hosted in the cloud, which likely has an even larger parameter count. However, we are generally limited to using the smallest LLMs for robotics.

Training Recipes

The \(\pi_0\) model also emphasizes developing effective training recipes. These methods are inspired from LLM training. They take a VLM that is pretrained on internet-scale data. They train the action expert using the Open-X Embodiment dataset. They then fine-tune the full system with a dataset that they manually collect. They argue that post training on high-quality data is important as the lower quality data contains mistakes from teleoperation that we don’t want robots to imitate.

Conclusion

Currently, reinforcement learning is one of the most active areas in AI research. Many researchers are trying to scale up RL training for LLMs in number of examples, environments, and tasks.

While hierarchical setups are useful to break apart complex tasks into subtasks, there are many alternative architectures. For example we could use Coconut (Chain of Continuous Thought) to simulate planning without forcing the steps to be human readable. We may see techniques from LLM reasoning research applied here.

Typically, VLAs operate on very short amounts of video. \(\pi_0\) only sees 2 seconds of video. This is mainly due to the latency constraint. Although long-context LLMs can process hours of video, this comes at a latency cost. Videos are very expensive in terms of token counts. Currently it is infeasible to train an effective VLA to use long visual context. However further research into token compression methods in LLMs/VLMs could benefit robotics. It is an open question as to how much long visual context is needed for effective robotics.

Compared to other tasks, robotics is especially latency-sensitive. Digital tasks like coding don’t have a time limit. We would like digital tasks to run fast, but in robotics because the world is changing, being slow can affect your ability to accomplish the task. Advances in the inference latency of LLMs could enable us to deploy larger and more performant models to robots.

It will be exciting to see how LLM research and VLA robotics research influence each other going forward.

Resources

Papers

Google DeepMind

Google’s line of work set up VLAs to be the very active area of research it is today. We will now go through architectural improvements to make VLAs more effective.

• SayCan
• RT-1
• RT-2
• RT-X
• Gemini Robotics

Physical Intelligence

• \(\pi_0\): A Vision-Language-Action Flow Model for General Robot Control
  • action expert trained with flow matching
• FAST: Efficient Action Tokenization for Vision-Language-Action Models
  • Autoregressive training is bad with repetitive tokens (which is common for robot actions). \(\pi_0\)-FAST does action compression to address this. Frequency sparse action sequence tokenization. 10% compression
  • Discrete Cosine Transform (DCT) used to compress action tokens.
• Hi Robot: Open-Ended Instruction Following with Hierarchical Vision-Language-Action Models
• \(\pi_{0.5}\): A Vision-Language-Action Model with Open-World Generalization
• Knowledge Insulating Vision-Language-Action Models: Train Fast, Run Fast, Generalize Better
• Real-Time Execution of Action Chunking Flow Policies

Other VLA

• Octo (May 2024)
• OpenVLA (June 2024)
  • open source
• TinyVLA (September 2024)
• CogAct (November 2024)
  • VLM + diffusion
• Magma (February 2025)
• HybridVLA (March 2025)
  • Cotraining with diffusion and autoregressive training
• GR00T N1 (March 2025)
  • Nvidia
  • VLM for reasoning + Diffusion transformer for actions
• SmolVLA (June 2025)
  • Hugging Face
  • flow matching action expert (pi0 architecture)
• MolmoAct (August 2025)
  • open source (Allen Institute for AI)
  • action reasoning model (ARM)
  • Completely autoregressive
  • Pre/Mid/Post training phases
• Many many more

Startups

Although their approaches aren’t always very public (except PI), it’s worth noting their approaches.

• DYNA
  • DYNA-1 (June 2025) https://www.dyna.co/dyna-1/research
    • First scalable foundation reward model for robotics
    • Folding laundry can be improved with a reward model. Humans are imperfect at folding. Behavior cloning is limited from teleoperated folding is imperfect. It is more powerful to use a strong reward model to judge the robot’s folding and use that signal to iterate and improve.
• Figure
  • Helix (February 2025) https://www.figure.ai/news/helix
  • S2: VLM built off of an open source 7B VLM
    • Passes a latent vector to the S1 model
  • S1: Higher frequency 80M model
• Skild
  • (July 2025) https://www.skild.ai/blogs/building-the-general-purpose-robotic-brain
  • Contains a low frequency and high frequency model

Talks

• Stanford CS25: V2 I Robotics and Imitation Learning
• Stanford Seminar - Connecting Robotics and Foundation Models, Brian Ichter of Google DeepMind
• Stanford CS25: V3 I Low-level Embodied Intelligence w/ Foundation Models
• Stanford CS25: V3 I Low-level Embodied Intelligence w/ Foundation Models
• U of T Robotics Institute Seminar: Sergey Levine (UC Berkeley)
• Princeton Robotics Seminar - Brian Ichter

The best way to specify few-shot tasks is through language. For example, you can input any image to ChatGPT or Gemini and ask questions like “What species of bird is in this photo?”. Without language you would need to collect training datasets of bird species and frame this as a classification problem. Alternatively, you could use few-shot learning with support sets by showing the model a few labeled examples of each bird species, but this approach is fundamentally brittle. You must pre-define every possible category and cannot handle abstract concepts like migration patterns or behavioral traits. Zero-shot learning is essentially impossible without language, as there’s no way to specify novel tasks the model is not trained to predict. Connecting vision to language allows us to specify arbitrary tasks without having to train any model. Some tasks also inherently require language generation, like image captioning and visual question answering.

In this blog, we will explore methods that connect computer vision with language and build an understanding of Vision-Language Models (VLMs). The distinction between VLMs and LLMs has become blurred, as most state-of-the-art LLMs now include vision capabilities. Nevertheless, some labs do release separate checkpoints for VLMs and text-only LLMs.

We will first clarify some of the terminology:

• Large Language Model (LLM): Large model primarily trained on language, however can be used to refer to multimodal models.
• Vision Language Model (VLM): An LLM with a vision encoder, allowing it to understand images.
• Multimodal Language Model (MLLM): A more general type of model that is built to process multiple modalities: images, video, audio.

Open VLMs

• PaliGemma, PaliGemma 2 (2024)
• DeepSeek-VL, DeepSeek-VL2 (2024)
• Qwen-VL (2023), Qwen2.5-VL (2025)
• Kimi-VL (2025)
• GLM (2025)

In this blog we will cover the vision encoders used in VLMs, different types of VLM architectures, and the training recipes used to train SOTA models.

Vision Encoders

In order to build a VLM, we need a vision encoder to map images/videos into embeddings or sequences of embeddings. While we could use a pure vision encoder model (as explained in my previous SSL blog posts), we achieve better performance when training vision encoders alongside text. Let’s explore some methods for pretraining vision encoders using image-text pairs.

CLIP

Introduced by Radford et al. 2021, CLIP is a simple way to learn image representations. They use a dataset of 400 million image - text pairs scraped from the internet.

An image encoder extracts an embedding for each image in the batch. The text encoder does the same thing for the corresponding texts. They then use a contrastive loss to push matching image and text embeddings closer to each other, and repel pairs from different examples. The outputs from the encoders are first linearly projected into “multimodal embeddings”. We then take the cosine similarity of each pair of these image and text multimodal embeddings. Cross entropy loss is computed for each row (image-to-text) and column (text-to-image) of the similarity matrix. For each image we want to be able to predict the corresponding text, and vice versa.

The CLIP loss is a sum of two cross entropy losses. $\tau$ is the temperature, which is the exponentiation of a learned scalar $t$: $\tau = e^t$. The first cross entropy term is the image-to-text loss, which is minimized when the text embedding corresponding to an image has a high dot product with that image. The second term inverts this relationship to form a text-to-image loss.

CLIP uses separate image and text encoders to map both into a shared embedding space for similarity comparison. Although the main goal of CLIP is representation learning, the resulting encoders can be used for zero shot tasks. You can specify many computer vision tasks as selecting a text among a group. For example, this can be applied to image classification. You can input a batch of images to the model. Rather than their corresponding captions, you input the classes as texts. The CLIP model can then match images to these arbitrary classes without any finetuning. Because it is trained on a large amount of images and text, this approach generalizes well.

The accuracy of the zero shot tasks can depend on how the texts are generated. Prompt engineering the class texts can improve performance. Multiple text prompts can be ensembled by averaging their embeddings.”

CLIP can also be used in a few shot setting by training a linear probe. This is just a single linear layer on top of the image encoder.

One clever trick is that the text embeddings can be used to warmstart a linear probe. Comparison with a set of text embeddings is mathematically the same as multiplying with a matrix containing the same number of rows.

Jia et al. 2021 released ALIGN shortly after which implements a similar approach. This method uses a different model architecture but uses the same contrastive training method. The dataset is much larger but noisier.

SigLIP

SigLIP, introduced by Zhai et al. 2023, replaces the softmax-based contrastive learning approach used in CLIP with a pairwise sigmoid loss. Instead of computing softmax across all pairs in a batch, SigLIP treats each image-text pair as an independent binary classification problem. This makes the training easier to parallelize across the batch.

The loss is computed for each image and text pair independently. Since most of the $N^2$ pairs are not matching, a learnable bias $b$ is used to correct this.

The advantage of this compared to CLIP is that we do not need to calculate a global normalization factor for the cross entropy. The SigLIP loss can be better parallelized across devices.

Each device gets a chunk of the images. In step 1, they compute the loss with respect to the matching chunk of text. Each device then iterates through the other chunks of text embeddings (containing only negatives) and accumulates the loss.

CapPa

An alternative to the contrastive training methods of CLIP/ALIGN and SigLIP is image captioning. The pretraining task is to predict the text from the image. This is also closer to how the vision encoder will be used in a VLM.

Captioning is an alternative training method. A transformer decoder (LLM) is trained to predict the text, while attending to the image through cross attention. This contrasts with contrastive approaches in that we are generating the text, rather than text embeddings. This requires a decoder architecture for the text encoder.

In addition to captioning (Cap), this work also incorporates parallel decoding to get CapPa. Instead of auto-regressive token prediction, the model predicts all tokens simultaneously. This is done by masking out all of the input tokens to the model and replacing causal attention with bidirectional self attention. The model switches between these autoregressive and parallel decoding modes. The authors find that this mixture performs better than using either mode alone.

CoCa

This paper introduces Contrastive Captioner (CoCa), which is a method that unifies the contrastive and captioning approaches to learning from image text pairs.

CoCa combines both approaches by training with a multi-task loss function that incorporates both contrastive learning and captioning objectives. Unlike CapPa, CoCa’s text decoder is split into unimodal and multimodal components.

The loss is a weighted combination of contrastive and captioning losses:

Attention pooling aggregates the vision embeddings into a single embedding, which is paired with the text decoder’s CLS token embedding for the contrastive loss. This implements a CLIP-style loss. A separate attention pooling layer creates embeddings for the multimodal text decoder. The decoder then uses standard next-token prediction loss for caption generation.

The contrastive loss can suffer from the hard negatives problem, where random sampling doesn’t yield the hard negatives that the model needs to learn fine-grained distinctions. The challenge with captioning is that it forces the model to predict very specific details that may be noisy or inconsistent across large datasets, making it overly sensitive to fine-grained variations that parallel the hard negatives issue. Combining these losses is a promising approach to get the best of both worlds.

Dynamic Resolution

So far, we’ve discussed encoders that process fixed-resolution images, such as 224×224 or 384×384 pixels. However, this approach has fundamental limitations. High-resolution images lose important details when downscaled, while low-resolution images waste computational resources when upscaled with padding. Unlike CNNs, transformers can handle variable-sized inputs—a capability we should leverage for vision tasks.

Modern VLMs have shifted toward dynamic resolution approaches that can handle variable input sizes more effectively. Rather than forcing all images into the same resolution, these methods adapt the processing to match the image’s native aspect ratio and information density. Implementation approaches include Native Resolution ViT and dynamic tiling, though we won’t cover these in detail here.

VLM Architectures

In multimodal ML and data processing, there are different ways to combine modalities:

• Early fusion: Combine modalities at the input level, processing them together through the entire model
• Late fusion: Process each modality separately and combine their outputs at the end
• Middle fusion: Combine the modalities at an intermediate representation

For modern VLMs, the general pattern is to use vision encoders for additional processing of visual data. LLMs prioritize language, with most of the training focused on language. Visual data is processed and adapted specifically for use by the language model.

Late fusion wouldn’t be effective for VLMs since the LLM needs access to visual data earlier in the pipeline to generate appropriate text outputs.

Intermediate fusion, where vision embeddings are added at a middle layer, could potentially lower computational costs. Visual data adds many embeddings and increases context length significantly. By skipping earlier layers, we could reduce computation and shrink the KV cache size. However, this approach isn’t commonly implemented in state-of-the-art VLMs, as model quality is generally prioritized over efficiency.

Early fusion VLM architectures use a vision encoder to output a sequence of vision embeddings. These embeddings are processed by an adapter before being appended to the LLM’s input. Typically, the LLM applies full attention to these visual tokens while using causal attention for the language tokens. This approach works because the model isn’t generating the image tokens but using them to generate text. Models that perform autoregressive image generation are an exception to this pattern (it’s rumored that GPT-4o uses this approach). Chameleon from Meta implements this. Currently, it is more popular to have separate models for image generation, so we’ll focus more on images as inputs rather than outputs.

Transformers have a quadratic computational complexity with respect to token length. Images introduce a large number of tokens, particularly at high resolutions, significantly increasing the computational demands on the LLM. When designing VLM architectures, token compression becomes essential. Several techniques exist to reduce the number of visual tokens that the language model must process.

We will now examine the architectural patterns used in VLMs. For early fusion, there are two main approaches: inserting vision tokens into the input sequence with an MLP adapter, or using cross attention with learned query embeddings.

Cross Attention Adapter

Cross attention adapters, also called attention poolers, compress sets of embeddings into a fixed number of output embeddings using learned query embeddings. These learned query embeddings attend to the vision tokens through cross attention, where keys and values are derived from the vision encoder’s output.

This is implemented in Qwen-VL as a “position-aware vision-language adapter”. The number of query embeddings dictates how many embeddings we choose to represent the image. There is a trade-off between capturing more information and reducing the sequence length that the LLM sees.

This cross attention layer is added in between the ViT and the LLM. The existing transformer blocks of the vision encoder and language model are not modified.

This architecture can gracefully handle different resolutions of images, but they train with 224x224 and post-train with 448x448. For the higher resolution post training, they increase the number of query embeddings from 256 to 1024.

MLP Adapter

In the LLaVA work by Liu et al. 2023, the authors train a simple VLM architecture by combining a pretrained LLM (Vicuna model, open-source instruction-tuned LLaMA) and a pretrained vision encoder (CLIP ViT-L). The vision tokens are linearly projected and then concatenated to the input of the LLM decoder.

This approach is currently more popular than cross attention. Qwen switched to MLP for Qwen-VL2. LLaMA 3.2-Vision, PaliGemma 2, and DeepSeek-VL also use MLP adapters.

With MLPs, we have less control over the token lengths. The simplest way to compress the token length is to project multiple vision encoder embeddings into the same input embedding. This is used in Qwen2.5-VL. Groups of 4 vision tokens are concatenated and passed to a two-layer MLP. This projects each group to a single embedding of the LLM’s dimension.

BLIP

BLIP-2 (Li et al. 2023) uses a Querying Transformer or Q-Former as an adapter. These architectures are similar to Flamingo in that they keep the vision encoder and text decoder frozen. However, BLIP is actually an early fusion architecture.

The Q-Former is similar to the Perceiver in that it processes images with a fixed number of learned query embeddings. However, it also takes in the text tokens. The text input shares the self attention layer but skips the vision cross attention and learns a new feed forward layer.

This is trained in two stages. In stage 1, three losses are optimized: Image-Text Matching (ITM, binary classification of whether image-text pairs match), Image-Grounded Text Generation (ITG, next token prediction loss conditioned on image), and Image-Text Contrastive Learning (ITC, CLIP loss). These different losses are set up by changing the attention masks.

In Stage 2, they connect the Q-Former to a frozen LLM, by concatenating the visual outputs to the LLM’s inputs. The text component of the Q-Former is discarded. This is finetuned to generate text.

Cross Attention Mid Fusion

The cross attention adapter uses attention to map embeddings to the input of the language model. An alternative design would be to add cross attention to the transformer blocks of the language model. These methods aren’t as popular as early fusion but are worth exploring.

Flamingo

Alayrac et al. 2022 introduce the Flamingo VLM architecture. It keeps the vision encoder and LLM parameters frozen but adds perceiver resamplers and gated cross-attention dense blocks. Since the LLM parameters are frozen, this method has no effect on the text only performance of the model and we can easily recover the text only LLM by deleting these additional modules.

The perceiver resampler takes a variable number of input embeddings and outputs a fixed number of output embeddings. It is essentially an attention pooler with a feed forward layer that is repeated by num_layers.

The gated cross attention dense layers are similar to the cross attention layers in encoder-decoder transformers. It is gated by layer specific learnable alpha parameters. These are initialized to 0, so the added vision modules have no impact at the start of training. The tanh activation ensures non-zero gradients for the alpha parameters after initialization.

While architecturally interesting, Flamingo’s approach has been superseded by simpler early fusion methods in modern VLMs. Generally, we don’t want to separate the vision and language parameters of our models. We want foundation models that are natively multimodal.

LLaMA

LLaMA 3.2 uses a ViT encoder trained with a contrastive objective for visual understanding. However, they discovered that the contrastive objective doesn’t preserve fine-grained localization information. To address this, they extract features from multiple layers (4th, 8th, 16th, 24th, and 31st) of the encoder.

They enhance the encoder by adding 8 additional transformer layers. The original ViT encoder, trained with CLIP, has 32 layers. These 8 extra layers increase the total to 40 and are designed to learn features more relevant to language generation. Both the attention and feed-forward layers in these additional layers use tanh gating. Unlike Flamingo, the gating parameters are initialized to $\pi/4$ instead of 0 (code), which effectively reduces the contribution of these added layers by about 40%.

The language decoder incorporates cross attention for images in every fourth transformer block. Unlike Flamingo, LLaMA doesn’t gate this cross attention.

LLaMA 3’s architecture is distinctive for using a mid fusion approach, setting it apart from most other models. This approach was necessary because LLaMA 3 was originally developed as a text-only model, and mid fusion is often required when adapting text-only models to handle images. However, with LLaMA 4, Meta transitioned to an early fusion MLP adapter approach, which they market as “native multimodality.” For early fusion to work effectively, the model must be trained end-to-end from early in the training process.

Training Recipes

LLMs undergo training through a progressive multistep pipeline. For text-only LLMs, we usually have the following stages: pretraining → supervised fine-tuning (SFT) → reinforcement learning (RL). Let’s examine how visual capabilities are integrated into these training pipelines.

LLaVA

The training recipe of LLaVA has two stages. The first stage involves pretraining the model on 595k image-text pairs filtered from the CC3M image pair dataset. The goal is to align visual and language representations. Only the MLP projector is unfrozen in this stage. The model is prompted to generate a brief description of the image, with a randomly sampled question (ex: “Summarize the visual content of the image.”, Appendix E). This prompt is generic and doesn’t consider the content of the image-text pair, however at this stage we only need to train the adapter.

In stage 2, the LLM and projector are unfrozen while the vision encoder stays frozen. The model is finetuned on higher quality data and more complex prompts. At the time of release, multimodal instruction following data was scarce. The authors addressed this gap with a synthetic data strategy. They leveraged a language only GPT-4 model to generate diverse questions related to the image. Note that the image itself is not used in the generation of these prompts. Instead, they input the caption of the image as well as bounding boxes if available.

To guide GPT-4’s generation, the authors manually designed seed examples for three data types. Conversations create multi-turn dialogues with simple questions and answers about details in the image such as object types, counting, locations, and spatial relationships. Detailed descriptions generate comprehensive scene descriptions that describe each object in the image. Complex reasoning involves more complex questions that require multiple steps to arrive at an answer. These seed examples were the only human annotations needed and enabled GPT-4 to generate thousands of similar examples through in-context learning.

Qwen-VL

In the pretraining stage, both the vision encoder and the adapter are unfrozen. Qwen trains at a much larger scale than LLaVA, using 5 billion image-text pairs. They initially pretrain at a lower resolution and then switch to higher resolution images for subsequent stages.

The multi-task training phase uses higher quality data covering various vision-language tasks, including captioning, visual question answering, grounding, and OCR. They also incorporate text-only data to maintain the model’s text generation abilities. All components remain unfrozen during this stage to develop general visual understanding, using 76.8 million examples.

In the final stage, the model is finetuned specifically for improved instruction following and dialogue capabilities. Since this stage doesn’t target better image understanding, the ViT is frozen. This phase uses a small but high-quality dataset of only 350k examples focused on multimodal instruction following and captioning. This progression demonstrates how each stage uses progressively less data of increasing quality.

Qwen2-VL uses the same training recipe but significantly scales up the datasets. It also uses the Native Resolution ViT architecture for dynamic resolution with the vision encoder. Unlike Qwen-VL, which used fixed image resolutions, Qwen2-VL also adds video support to both its architecture and training data.

Qwen2.5-VL further improves this training recipe by scaling up sequence length, which is crucial for handling longer videos. The first two training stages maintain the same sequence length. However, Qwen2.5-VL replaces stage 3 with “Long-Context Pretraining,” where the model is trained end-to-end (with both ViT and LLM unfrozen) on long-context data such as extended videos or images with substantial text. On top of these three pretraining stages, they add SFT and RLHF (DPO) post training stages.

They also show that this VLM training slightly reduces the text only performance of the LLM. This may be why VLMs are often released separately from the base text models.

DeepSeek-VL

The same pattern is used in DeepSeek-VL, except the vision encoder is not frozen during the SFT stage.

In DeepSeek-VL2, they introduce a dynamic tiling strategy to enable dynamic resolution. The same training stages are applied but with scaled up and improved datasets.

Kimi-VL

Kimi-VL follows a strategy similar to Qwen2.5-VL. They pretrain MoonVIT, a dynamic resolution vision encoder. Both the Qwen2.5-VL and Kimi-VL vision encoders are trained using SigLIP and captioning losses, similar to CoCa. Unlike approaches that freeze certain components, Kimi-VL trains all parameters during each stage.

During joint pretraining, the VLM trains on both image-text pairs and text-only data. This balanced approach is crucial at scale to prevent degradation of text generation capabilities. The cool down stage uses higher quality data, while the long context stage incorporates both long-context text and multimodal examples. A notable feature of Kimi’s approach is the inclusion of both text and multimodal inputs throughout all training stages. They continue training with SFT and RL post training stages to develop a multimodal reasoning model.

Despite variations in training pipelines, these models share common patterns. They all use multi-stage pretraining, with each progressive stage processing fewer examples but incorporating higher image resolution, longer context length, and improved data quality.

Video

We have focused on image inputs but VLMs can also be developed to understand video inputs. Rather than extracting tokens for a single image, you can process each frame separately and concatenate the tokens. Audio tokens can be processed separately and then interleaved/concatenated with the frame tokens.

One important thing to note is that video is far more expensive to process than images. A single image already requires hundreds to thousands of tokens, but videos add a high multiplier. For example, a 1-minute video at 30fps means processing 1,800 frames. Depending on the tokenizer, each frame could be 256 tokens. This means the one minute video would be 460k tokens. If we use the rule of thumb that a token is equal to 3/4 of a word, this means this is equal to 345K words. If we assume there are 500 words on a page, this is approximately equivalent to a 700 page book. However, a minute of video has far less useful information than a long book. This redundancy creates opportunities for compression, which is critical for processing longer videos.

The simplest way to reduce token count is to choose a low frequency to sample the video frames in order to reduce processing and token count. If we use 1 FPS, we reduce the token count for 1 minute of video to 15K tokens. It is impractical to train a VLM with videos at full frequency.

There are some interesting works on making this more efficient, leveraging the fact that there is a lot of redundant information between nearby video frames. One simple method is to project patches from multiple frames to a single token. Qwen2.5-VL does this through a strided Conv3D. Typically in a vision transformer encoder, the image is split into 2D patches which get projected to a token embedding (equivalent to a strided Conv2D). Qwen uses 3D patches. This aggregates information from multiple frames into individual tokens. This is like reducing the frame frequency while not dropping any frames entirely.

There are more sophisticated methods that allow for a more dynamic compression that takes into account the variable amount of redundancy. For example, Run-Length Tokenization (RLT) (Choudhury et al. 2024) identifies consecutive patches that are repeated across video frames and replaces them with a single token plus a duration encoding. Consecutive patches are considered repeated if the L1 distance between them is below a threshold $\tau$. This method is content-aware in that the number of tokens depends on how redundant the frames of the video are.

RLT is a simple heuristic-based method of video token compression. There are also model based approaches for this.

Conclusion

When I first started working on ML, it seemed like computer vision and natural language processing were parallel and equally important fields. There was some knowledge transfer, but they were fairly independent, both in training methods and model architectures. Transformers were invented for NLP but were soon applied to other data modalities including vision. More recently, LLMs have been dominating the field of AI. As these models have become more multimodal, it seems like other domains are getting absorbed into LLM research. While computer vision remains an active area of research, a large portion of the field is focused on improving the visual capabilities of LLMs.

This convergence and the rise of VLMs is leading to many interesting applications, such as action planning in robotics and visual accessibility tools. As these models continue to improve, we’re moving closer to AI systems that can understand the real world, which is very visual.

It is a common objective for machine learning models to predict future observations based on current and past observations. Language models, for instance, predict subsequent words from preceding ones. World models extend this concept by incorporating actions. These models learn how actions affect the environment, enabling the development of intelligent agents that can effectively plan and interact with their surroundings.

A world model is a function that maps the current state of the world \(x_t\) and an action \(a_t\) to a prediction of the next state \(x_{t+1}\). The definition of observation and action varies depending on the implementation.

World models enable many important capabilities for intelligent systems. Learning an effective world model requires understanding physics, causality, spatial intelligence, and more. These are especially important as we deploy AI systems in the real world.

World models have diverse applications across different domains. While significant research has focused on video game environments for training game-playing agents, their most crucial application lies in real-world embodied AI. Robotics applications require world models to understand and interact with their physical environment, predict the outcomes of their actions, and adapt to changing conditions. These models help robots perform complex tasks that require planning.

World models represent an exciting but still nascent area of AI research. The term “world model” remains somewhat loosely defined. In this blog post, we will discuss different implementations of world models.

Generative World Models

World models have immediate applications in physical AI, particularly in robotics and autonomous vehicles. These systems require interactive data that demonstrates how actions lead to specific outcomes, unlike traditional vision models that simply observe static data.

A significant challenge in robotics is the limited availability of training data compared to language models. While the most valuable data comes from real-world interactions with physical agents, collecting this data is both expensive and time-consuming. We also need a large amount and variety of environments so the agents can learn to generalize well. The goal of this class of world models is to generate realistic and diverse environments in order to train robust agents.

Genie

This work from Google is a “generative interactive environment” trained on internet data (Genie-2 Blog Post). It is capable of generating interactive environments from a text or image prompt. We will explain each component of the model.

Spatial-Temporal Transformer

Since it is trained on videos, Genie uses the Spatial-Temporal Transformer (ST-Transformer) architecture to implement most model components. The ST-Transformer block has three components, temporal attention, spatial attention, and a single feed forward block. The spatial attention is self attention among patches of a single frame. The temporal attention is for a single patch on \(T\) time steps. This uses causal attention. This architecture scales linearly rather than quadratically with the number of frames, which makes it efficient at processing video.

Latent Action Model (LAM)

In world models, we need to associate actions with transitions between observations. Since actions are not available on internet scale data, this model learns “latent actions”. This model uses an encoder and decoder. The input to the encoder is all prior frames \(x_{1:t}\) and the next frame \(x_{t+1}\). The encoder outputs latent actions \(\tilde{a}_{1:t}\). The decoder uses these latent actions and the prior frames to predict the next frame \(\hat{x}_{t+1}\). This is implemented using a VQ-VAE objective. A small code book of size 8 is used for controllability. The small codebook size also acts as a bottleneck. This forces the representation to only store information about the change rather than encode the observation itself.

When trained on 2d videos, the changes between frames can be described in a compact code. Think of moving left, right, jump etc. this may be more difficult in complex real world environments. Since the latent actions use a small codebook and the decoder has access to the prior frames, the model is forced to learn the differences between frames which can be represented as actions.

Video Tokenizer

Another VQ-VAE is trained to generate tokens from the video. The difference here is that the input and output are identical and just the video frames.

Dynamics Model

The dynamics model takes the prior video tokens and latent actions and predicts future video tokens. This also uses a ST-Transformer architecture. Since this is a causal transformer, at each timestep \(i\) the model uses the current latent embedding \(z_i\) to predict the next timestep’s embeddings \(z_{i+1}\). At training time, this is parallelized so \(z_{1:t-1}\) is mapped to \(z_{2:t}\).

The MaskGIT method enhances this model’s performance. Since we’re dealing with images/video frames, \(z_i\) represents a grid of latent embeddings. Traditional autoregressive prediction of these embeddings in one shot can be unstable. In this setting, if a single token is mispredicted, all subsequent predictions suffer. MaskGIT resolves this through an iterative process that enables prediction corrections. The tokens within the next frame can attend to each other bidirectionally.

During training, tokens are randomly masked according to a Bernoulli distribution with a sampled probability between 0.5 and 1, and the model predicts these masked tokens. This training approach enables iterative decoding at inference time. The next frame is initialized as zeros, and the dynamics model makes predictions based on the context. The dynamics model then processes this output as an input. Since the predictions are a softmax over discrete token values, we can determine which tokens are low in confidence. The model in subsequent iteration only updates these low confidence predictions. This is repeated across 25 steps, where output tokens are regenerated until reaching high confidence. This method is similar to diffusion, but uses discrete rather than continuous tokens. The masking approach makes the dynamics model’s predictions more robust and reduces the redundancy in video data for more efficient learning.

Training and Inference

The latent action model and video tokenizer are trained separately first. They are trained using the VQ-VAE objective. The LAM and video tokenizer are then used to train the dynamics model. There is a stop grad so that the latent action model is frozen.

At inference time, the LAM is discarded as we can define the actions directly. The model is prompted with a single initial frame which is tokenized. At each step, the user chooses an integer for the latent action. The initial frame tokens and the latent action are used to generate the next latent token. This generated token is appended to the input . The user can then select another action. At each step, the tokenizer decoder can map the latent embeddings to images which can be displayed to the user.

The video tokenizer and LAM are very similar. One significant difference is that LAMs use smaller codebook sizes. The focus is on having a controllable set of actions (8 codes). The video tokenizer uses a larger codebook size (1024) to more effectively reconstruct the video. The LAM is able to use a smaller codebook size since it is only trying to encode one frame at a time. Whereas the video tokenizer tokenizes the entire video and then reconstructs it all together.

Genie achieves a significant breakthrough by training on large-scale data without requiring explicit actions. Instead, it learns a finite set of latent actions that enable interaction with generated environments.

Through these latent actions, users can control newly generated games. While this approach is powerful because it enables training on extensive datasets without ground truth actions, the learned actions come with limitations. Since the actions are discovered rather than defined, there’s no way to enforce specific mappings. For example, with 4 latent actions, we cannot guarantee that action 1 means “move left” and action 2 means “move right.” The model might learn diagonal movements instead. This use of latent actions somewhat weakens Genie’s classification as a true world model, since traditional world models aim to learn how observations change based on well-defined, real-world actions.

Cosmos

Released by NVIDIA, Cosmos is a “world foundation model”. It is trained on a large dataset of 20 million hours of video. Unlike other works, Cosmos focuses on physical AI and real world data rather than simulated video game environments. It is built for physical AI applications like robotics and autonomous vehicles. We will cover the different components of this model.

Tokenizer

Effective world modeling requires efficiently tokenizing videos. They make available both discrete and continuous tokenizers. They are both trained in the same way, but the discrete tokenizer uses Finite-Scalar-Quantization (FSQ) to discretize the tokens.

The tokenizer models use a temporally causal encoder-decoder architecture to map videos from \(x_{0:T} \in \mathbb{R}^{(1+T) \times H \times W \times 3}\) to “token videos” \(z_{0:T'} \in \mathbb{R}^{(1+T') \times H' \times W' \times C}\) which are more compact spatially and temporally. To train this, they use a simple L1 reconstruction loss. In a second stage of training, they use an optical flow loss to improve the smoothness between frames of the video, and a Gram-matrix loss to improve the sharpness of the frames. An adversarial loss is used to further enhance this.

World Model

Like with tokenizers, the world model is also implemented for continuous and discrete tokens. However, the model architectures are more different. Diffusion is used for continuous tokens, while autoregressive models are learned with the discrete tokens.

Diffusion

They use a DiT (Diffusion Transformer) based architecture. The continuous tokens are perturbed with Gaussian noise and the DiT model denoises these tokens. The model is conditioned on text through cross attention. The text prompt provides additional information about the video that is used for denoising. At inference time, the text can be used to generate new videos. The model can also be conditioned with input image frames to generate video continuations.

Autoregressive

The autoregressive world model trains on the next token prediction objective using discrete video tokens. This training happens in multiple stages:

Stage 1: Predict tokens of future frames given a context of the first 17 frames.

Stage 1.1: Context length is increased to 34 frames. Note that each frame produces a large number of tokens.

Stage 2: Train with text conditioning.

A diffusion model is used as a decoder to achieve higher quality video generations. This decoder is trained separately.

Cosmos includes two other classifications of world models. Text2World maps text prompts to videos. Video2World models take text prompts and video as input. These models generate continuations of the video.

In Cosmos, actions are referred to as perturbations or \(c_t\). During WFM training, these perturbations are simply text prompts describing the videos. Text effectively encodes actions when continuing videos, allowing for generalizable world model training. During finetuning, actions can be encoded in more specific formats. For example, when finetuning on camera control, the model is conditioned with Plücker coordinates that represent the camera’s pose. The model learns to generate video matching these coordinates. Additional models can also be incorporated to condition the WFM on task-specific actions.

During world model pretraining, explicit actions aren’t used. The model functions essentially as a video generation model. However, the models are later finetuned to incorporate explicit actions for different use cases, transforming them into true world models.

Cosmos was released very recently. It will be interesting to see whether “World Foundation Models” become more widely developed and used in the coming months. No model has been trained yet on this scale of physical AI data.

World Models for Agents

The previous works focus on generating interactive environments that can be used to train intelligent autonomous agents. Now, we will explore approaches that use world models to develop the agents themselves. This is essential for leveraging all the generated environments. We want to create autonomous agents capable of performing well across diverse environments and tasks, especially with limited data. World models provide valuable solutions to this challenge.

World models are an important component of model-based reinforcement learning. In MBRL, a dynamics model predicts the next state from the current state and action. There are different methods to use this for planning. Typically, the dynamics model is learned from a dataset of trajectories for the task at hand, from the real world or simulation. The states that are being predicted are task specific representations.

The dynamics model can be implemented with a world model. This is when the model is trained on a large and diverse dataset of offline data, rather than task-specific trajectories. World models enable learning a single dynamics model, which can generalize to multiple tasks. This is important for robotics, since you want the robot to be able to do things that aren’t optimized for during training. For this generalization, it is important that the world model is trained on a variety of tasks so it doesn’t focus on task-specific information.

Ha et al. 2018

This relatively early paper shows the advantages of world models in the context of reinforcement learning. They train world models on different OpenAI Gym environments to learn effective policies.

V: The vision model is a VAE that learns a latent representation \(z\) of the observation (image) from the environment.

M: The memory RNN learns representations across time and makes predictions in the latent space. It is modeling \(P(z_{t+1}\mid a_t, z_t, h_t)\). The next latent representation is predicted from the current representation, action, and hidden state. This is a probability distribution using a mixture density model, so we can sample from it.

C: The controller predicts an action from the hidden state and current latent representation. This can be a lightweight linear model.

Steps to train

1. Generate rollouts under a random policy
2. Train \(V\)on the frames
3. Train \(M\) on the latent representations and the actions of the random policy
4. Optimize the controller
5. Restart to learn from rollouts from the new policy

We can also use the world model to generate new rollouts in the latent space. This is done by sampling latent representations in \(M\) and using it to train the controller. One issue with this approach is that the policy can exploit gaps in the world model. The stochasticity of \(M\) helps train a more robust controller.

This early work laid important foundations for subsequent research in world models, showcasing their potential for creating more sample-efficient and generalized learning systems.

Dreamer

The paper “Mastering Diverse Domains through World Models” by Hafner et al. 2023 introduces the DreamerV3 model. This line of work aims to train world models to use for model-based reinforcement learning. This approach utilizes two stages: world model learning and actor critic learning.

World Model

The world model is implemented as a Recurrent State-Space Model (RSSM), which is learned using trajectories with frames and actions.

Sequence model \(h_t=f_{\phi}(h_{t-1},z_{t-1},a_{t-1})\): This model predicts the next hidden state from the prior hidden state, latent representation, and action.

Encoder \(z_t \sim q_{\phi}(z_t \mid h_t, x_t)\): This model generates the latent representation of a timestep \(z_t\) from the observations (frames). This also utilizes the hidden state \(h_t\).

Dynamics predictor \(\hat{z}_t \sim p_{\phi}(\hat{z}_t\mid h_t)\): This predicts the latent representation from the hidden state.

Reward / Continue predictor \(\hat{r}_t \sim p_{\phi}(\hat{r}_t \mid h_t, z_t)\), \(\hat{c}_t \sim p_{\phi}(\hat{c}_t \mid h_t, z_t)\): These models predict the reward and continue signal. The RL environments are set up to have a scalar value for these at each timestep.

Decoder \(x_t \sim p_{\phi}(\hat{x}_t \mid h_t, z_t)\): This predicts the observation frame from the latent representation. The decoder is not necessary for planning but training with it ensures that \(z_t\) captures information from the observations. This model also enables visualizing planned trajectories. The encoder and decoder together form a discrete autoencoder.

With these modules, we can train the world model using offline trajectories. There are three different losses which train different modules.

The prediction loss \(L_{\text{pred}}\) trains the decoder, reward predictor, and continue predictor (corresponding to the three terms in the equation). \(L_{\text{dyn}}\) and \(L_{\text{rep}}\) train the encoder and dynamics model, similar to the VQ-VAE codebook and commitment losses (VAE blog). The dynamics loss trains the dynamics predictor to predict latent representations from hidden states that match the encoder’s output, while the representation loss trains the encoder to generate latent representations that match the dynamics predictor’s output. These losses are clipped so the learning focuses on the prediction loss when the other two are reasonably optimized. This also prevents the encoder from including irrelevant information in the representations, which makes control more difficult. The sequence model isn’t directly included in any of these losses, but receives gradients as we generate hidden states through the trajectories.

These losses are weighted and calculated for each timestep of the trajectories to train all the modules of the world model at once.

Actor Critic Learning

Once we train the world model using offline trajectories, we can use actor-critic learning on simulated online trajectories. This is a popular RL method. The actor (\(a_t\sim \pi_{\theta}(a_t\mid s_t)\)) predicts the action to take given the current state. The critic (\(v_{\psi}(R_t\mid s_t)\)) predicts the expected sum of future rewards for a state.

Rather than interacting with real environments, the actor can be trained on imagined trajectories from the world model. The world model and actor simulate trajectories along with reward and continue signals. This stage is the dreaming. The critic is trained to learn the sum of rewards, while the actor learns a policy to maximize returns. This enables learning an effective policy with simulated interactions.

The actor and critic are trained concurrently on both imagined trajectories from the world model and real trajectories stored in a replay buffer. While the world model improves sample efficiency, some real interaction is still needed to fine-tune on the specific task. The ratio of imagined to real trajectories and the weighting of losses are tunable hyperparameters.

Dreamer uses a general purpose world model to be able to more effectively and efficiently learn policies for a variety of tasks and environments. One shortcoming of Ha et al. 2018 and Dreamer is that they only use trajectories with rewards for training. As we know, this kind of data is limited. We want to be able to train with videos that have no actions and rewards, and even with static unlabeled images.

DINO-WM

This method uses a pretrained DINOv2 SSL model (covered here) to bootstrap a world model. The DINO model is trained at a very large scale to produce image representations for visual understanding. This method incorporates this capability as a component of a world model, rather than trying to learn it through the world model itself.

The authors also opt for avoiding reconstruction of images and rather predict latent representations. This follows the JEPA framework.

The environment is modeled as a partially observed Markov Decision Process POMDP defined as \((\mathcal{O},\mathcal{A},p)\) (observations, actions, probability). Like the other methods, there are three components of this world model:

The observation model learns a latent representation. The transition model predicts future representations given the prior observations and actions. The decoder model reconstructs the image observation, which is optional (visualization/generation).

Observation Model

The observation model is trained to learn visual representations. This model is meant to generalize to many tasks and environments. For this they use a pre-trained DINOv2 model. This model is trained on internet scale image data, so it generalizes well. At time step \(t\) it encodes an image \(o_t\) into patch embeddings \(z_t\) of shape \((N, E)\), where \(N\) is the number of patches and \(E\) is the embedding dimension.

Transition Model

The transition model is trained as a decoder-only transformer model. Given past latents, the model predicts future latents. Through an attention mask, patch embeddings indexed by \(i\) only attend to embeddings of the same index. That is the latent state \(z_t^i\) attends to \(\{z_{t-H:t-1}^i)_{i=1}^N\). \(H\) represents the context length, which is the number of steps the model can look back. There is attention between the patches in the observation model, but the transition model looks at each patch independently.

To incorporate actions, a K-dimensional vector is concatenated to each patch embedding in the input. This is mapped from the original action representation, which is a continuous vector defining the action, using an MLP. The representation of an action changes based on the environment or task. For example, it could be a 2D vector like (-0.1, 2.1), describing how to push an object. The actions are included in the training datasets. This model is trained with teacher forcing using these trajectories.

Additional metadata, such as proprioceptive information, can also be mapped to embeddings and concatenated. The model is trained with a simple L2 loss:

Decoder

The output of the observation model is processed to reconstruct the input image using a simple reconstruction loss. The decoder loss is not backpropagated to the encoder/observation model. This approach proves empirically more effective and maintains the world model’s planning capabilities entirely in the latent space.

Test Time Optimization

While the model is trained with trajectories from the dataset, at test time we want to do planning in order to get the best possible sequence of actions.

Given the current observation \(o_0\) and a goal observation \(o_g\), which is an image of the desired outcome. Model predictive control is used to find a sequence of actions to reach the goal. The following cost is optimized: \(\mathcal{C} = \| \hat{z}_T - z_g^2 \|\) where \(\hat{z}_t = p(\hat{z}_{t-1},a_{t-1})\), \(\hat{z}_0 = enc(o_0)\), \(\hat{z}_g = enc(o_g)\).

This method defines goal states with images. This is only feasible for simulated environments. This method is not directly applicable to more complex real world applications.

Cross Entropy Method (CEM)

To implement MPC, this method was found to be most effective. The policy is represented as a tensor of Gaussian distributions of shape \((T, K)\), where \(T\) is the number of steps in the trajectory, and \(K\) is the size of the action embedding. The method samples \(N\) trajectories from this distribution. The world model processes each trajectory and computes the final latent state, which is used to calculate each trajectory’s cost. The top \(K\) trajectories (not to be confused with the action embedding size) are selected. Their mean and standard deviation are then used to redefine the Gaussian distributions. This iterative process continues until it finds a sequence of actions that minimizes the cost.

CEM is similar to importance sampling in that it samples a set of trajectories and uses the successful ones to update the distribution.

Gradient Descent

This method was found to be less effective but is still worth noting as a direction for future work. In this method, the actions are initialized as learning embeddings. These actions are initialized randomly. The cost is used to update these embeddings through gradient descent.

Other Methods

Video Generation Models

Video generation models like Veo and Sora occupy an ambiguous space in their classification as world models. While they can generate realistic videos from text prompts and create natural continuations of existing videos, they don’t model how the world responds to actions. Though one could conceptualize an action as a video segment and its continuation as a world model simulation, this approach falls short. We need models that can simulate the world without requiring actions as input and understand the consequences of multiple actions in sequence.

JEPA World Model

I-JEPA and V-JEPA (discussed in my JEPA blog post), and similar self-supervised learning methods can be interpreted as world models. These methods work by masking parts of the image or video. The masking itself functions as the action. The model has access to the location of the masks and predicts the latent representations of the masked regions.

This work introduces Image World Models (IWM). IWM extends I-JEPA by using various image distortions beyond simple masking—such as changes to brightness, contrast, and saturation. These distortions are commonly used in contrastive SSL methods. The action becomes a representation of the distortion applied to the source image to produce the target image. Using the source latent and the action, the IWM predicts the latent representation of the target image.

Since this work focuses solely on images, it lacks a time component. The world that the IWM models exists without a temporal dimension, limiting its application to real-world situations. Additionally, the action space isn’t particularly useful for practical applications. Nevertheless, this approach helps develop richer image representations than those from I-JEPA, which relies exclusively on masking. Distortions are often used to make the image representations invariant to them. However, in this use case, we don’t directly learn invariance but rather the relationships between distorted images. I would classify this method as a great improvement in self supervised representation learning than an advancement in world models.

It is worth noting that many other world model architectures, such as DINO-WM, can be considered Joint Embedding Predicative Architectures (JEPAs).

Conclusion

World models are an important area of research with significant potential. The works described here generally use simple actions, like movement in 2D video games. In the future, we need to model more complex actions at different granularities.

As an example, let’s consider a home robot that does chores. A high-level action would be selecting which chore to start with. This requires considering which task is most important now and how it would affect the household in the near future. Executing the chore requires lower levels of granularity. Say the robot decides to wash dishes, it has to select the order of dishes to wash to optimize storage. And in washing a dish, it has to determine the sequence of mechanical movements to do so effectively and efficiently. We may also need world models that plan for even longer timeframes, like days or months. Our internal world model is powerful, and these kinds of inferences and planning come very easy to us. We underestimate what it would take to replicate this in machines.

The vanilla autoencoder excels at compression and representation learning but faces challenges when used as a generative model. Due to the curse of dimensionality, we can’t effectively sample from the latent space of a vanilla autoencoder. The space of all possible latent representations is vast, with only a small portion occurring in the dataset. Consequently, most random points in this space would decode into noise.

While it might be possible to sample from a low-dimensional latent space, if the dimension is too small, the model loses expressiveness and doesn’t scale well. This type of autoencoder is valuable for representation learning but falls short as an effective generative model. Variational autoencoders seek to address this.

For the vanilla autoencoder, the loss is defined as \(L_{reconstruction}(x, g(f(x)))\). Where \(x\) is the input, \(f\) is the encoder, and \(g\) is the decoder. \(L\) is a reconstruction loss that is often set to mean squared error (MSE). The latent embedding can be represented as \(z = f(x)\). All autoencoders map observed data to latent embeddings where the latent dimension is a hyperparameter.

This architecture is simple to implement and is useful for compression and representation learning. Variational autoencoders were introduced to address different deficiencies of this architecture, which we will cover.

Variational Autoencoders (VAE)

The goal of variational autoencoders is to constrain the latent space of an autoencoder so that it can be sampled from. VAEs use variational inference to create a probabilistic latent space. VAE was introduced in Auto-Encoding Variational Bayes in 2013. Although the acronym VAE isn’t present in this paper.

A simple way to construct this would be to have the encoder output the parameters to a Gaussian distribution. A mean and standard deviation are predicted for each latent dimension. Using these parameters, we can sample from the distribution as such.

Reparameterization Trick

However, the architecture defined above is not trainable. The sample operation is not differentiable, so it would not be possible to train the encoder. A “reparameterization trick” is used to solve this.

This trick is a simple and clever implementation detail, since Gaussians are normally sampled in this way. We are sampling \(\epsilon\) and then treating the sampling of \(z\) as a deterministic process. This can be done only because \(\epsilon\) is independent of any parameters. The stochastic component is refactored to be a fixed distribution that has no dependencies for backpropagation.

Although Gaussians are commonly used, any distribution that can be reparameterized like this, can be used in a VAE.

Forward Pass

In order to understand how the VAE is trained, we need to introduce some mathematical notation.

Encoder network \(q(z\mid x)\): For each input \(x_i\), this outputs vectors for mean \(\mu_i\) and log variance \(logvar_i\). These parameters are used to generate a latent embedding.

Latent variable \(z\): This is computed using the reparameterization trick. \(z_i = μ_i + σ_i * ε\), where \(σ_i = \sqrt{\exp(logvar_i)}\) and \(ε ∼ N(0, I)\)

Decoder network \(p(x\mid z)\): Given the latent embedding this outputs the parameters of the output distribution. This distribution can be a Gaussian distribution for continuous data (like natural images) or a Bernoulli distribution for binary data.

This notation describes the forward pass of the VAE. We will now explain how this is trained.

Loss

During training, we sample from the conditional distribution\(q(z\mid x)\). At test time, however, we need to use the decoder to generate new samples by sampling from the prior \(p(z)\). For this to work, the model’s marginal posterior \(q(z)\) must approximate a standard Gaussian prior \(p(z) = N(0, I)\). To achieve this while sampling from \(q(z\mid x)\), we add a KL divergence term to the loss function: \(\text{KL}(q(z\mid x) \mid\mid p(z))\). We use the reverse KL divergence specifically to ensure the approximate posterior places probability mass only where the prior does. This KL term has a closed form for Gaussian distributions:

The reconstruction loss is defined differently for Gaussian and Bernoulli output distributions. We will focus on the Gaussian form.

In practice this is implemented as the MSE between the mean of the output \(p(x\mid z)\) and the input \(x\). We can expand the log probability as follows (\(\mu_q\) and \(\sigma_q\) represents the outputs of the decoder):

We can take the mean and drop the constants and optimize the MSE:

The optimization process combines both losses as follows:

1. Take a batch of input examples (images): \(x_i\).
2. Use the encoder to generate mean and log variance parameters: \(\mu_i\) , \(logvar_i\).
3. Sample \(\epsilon\) to generate latent embeddings \(z_i\) for each example using the reparameterization trick.
4. Pass the latent embeddings through the decoder and take the mean of the output distribution or reconstruction: \(\hat{x}_i\).
5. Calculate the reconstruction loss between the decoder’s output mean and the inputs \(x_i\). Calculate the KL loss using the mean and variance from the encoder. Sum these losses and backpropagate through both networks.

Generation

To generate examples with a VAE, we only use the decoder. First, we sample \(z\) from the prior \(p(z)\), which is set to the unit Gaussian. This sampled value passes through the decoder to produce the distribution \(p(x\mid z)\), which can be either Gaussian or Bernoulli. For images, we typically don’t sample from \(p(x\mid z)\) since it represents pixel-level probabilities and sampling from it would only add noise. By sampling different values of \(z\), we can generate diverse images.

ELBO

The loss we defined for VAE is known as the Expected Lower Bound (ELBO). This can be formulated as follows:

We take the negations of the losses for ELBO because it is a lower bound that we want to maximize through gradient ascent. However, in practice we may just minimize negative ELBO.

We have arrived at ELBO from a practical standpoint. This can also be derived through variational inference. We will now go deeper into the mathematical intuitions of ELBO and VAE.

We can start by considering training the decoder by maximizing the likelihood:

This is intractable because it requires taking an integral over the continuous latent variable. Instead of optimizing the intractable likelihood, we can formulate a lower bound. Maximizing this lower bound is a proxy to optimizing the likelihood. Variational inference is a mathematical framework that approximates complex probability distributions by optimizing a simpler distribution to be as close as possible to the target distribution.

Derivations

In unsupervised learning we want to increase the probability of the data, which we express as \(\mathrm{log}(p(x_i))\).

We can first express the log probability as an expectation with respect to the latent variable. This is done by expanding \(p(x)\) as an integral over \(z\). We then multiply the numerator and denominator by \(q(z)\) and rearrange these terms:

Jensen’s Inequality states that for any concave function \(f\): \(f(E[y]) \geq E[f(y)]\). This applies to the log likelihood because \(\log\) is a concave function. We can apply Jensen’s inequality to the expectation we have derived.

The quantity on the right represents the lower bound on the log probability. We can rearrange the terms on the right for a more interpretable formula for ELBO.

The first term is the reconstruction loss; the second is the KL term.

Alternative

It is also possible to derive ELBO from the definition of KL divergence. The first step to this derivation is the same. However, we continue by applying Bayes’ rule instead.

We start with the definition of KL divergence between the posterior and approximate posterior:

Using Bayes’ rule, we can express the posterior as:

Substituting this into the KL divergence:

We are able to remove \(p(x)\) from the expectation since it is the only term that doesn’t depend on \(z\). We can now rearrange the terms to isolate \(\log p(x)\). We want to place a bound on the quantity, like we did in the first derivation.

We drop the last term since KL is non-negative. We are able to get to the same lower bound.

Intuitions of ELBO

Now that we have derived ELBO, we will explain the intuition of its terms. We can first decompose the KL term.

The first term, \(\mathbb{E}_{z \sim q(z\mid x_i)} \left[ \log q(z\mid x_i) \right]\), is the negative entropy of the approximate posterior distribution. We want to increase the entropy of the decoder distribution so that we get diverse latent embeddings. This encourages the model to give probability to a wider range of latent embeddings and not overfit to precise values. This makes the VAE able to generalize better and make it easier to sample from.

The second term, \(\mathbb{E}_{z \sim q(z\mid x_i)} \left[ \log p(z) \right]\), is the cross-entropy between the approximate posterior and the prior. This term encourages the approximate posterior to be similar to the prior distribution. Since we typically use a unit Gaussian as the prior, this term pushes the approximate posterior toward a standard normal distribution. It is important that this distribution is close to the distribution we will be sampling from.

We use the reverse KL to strongly penalize \(q\) taking non zero values when \(p\) is zero. This encourages the supports of each distribution to be more aligned. We don’t want to generate latent embeddings during training that would be low probability during inference / generation.

Problems in VAE

Although it is a powerful architecture, there are several weaknesses to the VAE that follow up research seeks to address.

• Blurry images: VAEs tend to produce blurry output images due to their continuous latent embeddings.
• Posterior collapse: This occurs when the decoder ignores the latent embedding. If a decoder is too powerful, it can generate high-quality images without using information from the latent space. While this doesn’t harm image generation itself, it renders the latent embedding meaningless as a representation.
• Entanglement of the latent embedding: The latent dimensions in a VAE often become entangled, meaning that individual dimensions do not correspond to interpretable features. This makes it difficult to control specific attributes during generation or perform meaningful latent space manipulation.

β-VAE

This work by Higgins et al. 2017 explores a simple change to the VAE ELBO loss, by simply adding a weight \(\beta\) to the KL term of the loss:

In the paper, they set \(\beta > 1\) and observe that it encourages more disentangled representations. This is because the prior \(p(z)\) is a multivariate distribution where each latent variable is completely independent of others. This also constrains the information capacity of the embedding, which forces the encoder to learn more of the underlying factors of the data. This is built on the assumption that disentangled representations are more efficient than entangled ones.

The authors point out that a latent embedding with independent factors isn’t necessarily disentangled. Entangled representations can be expressed as independent factors through PCA. We want these factors to be semantically interpretable.

Vector-Quantized VAE (VQ-VAE)

The latent embedding of the VAE is continuous. However, for some applications, we want to have a discrete latent representation. This can be used as categorical features in a downstream model. This can also be used for autoregressive modeling if the discrete values can be used to create a vocabulary.

van den Oord et al. (2017) introduce VQ-VAE to address this problem.

Forward Pass

1. A CNN encoder maps the input image to a grid of latent embeddings. These are deterministic and do not require sampling through reparameterization.
2. For each embedding, we find the closest embedding in the codebook. We replace the embedding with this codebook embedding.
3. A CNN decoder reconstructs the image from the codebook embeddings.

The encoder generates an embedding \(z_e(x)\). This is matched to a codebook embedding \(e_j\). \(z_q(e)\) represents the quantized embedding. The encoder and quantization can be interpreted as a categorical probability distribution:

The decoder reconstructs the input from the quantized embedding: \(p(x\mid z_q(x))\)

Training

The VQ-VAE is trained with three loss terms:

1. The first term is the reconstruction loss, which is used to train the decoder.
2. The second term is the codebook loss which pushes the codebook embedding closer to the encoder’s output embedding.
3. The third term is the commitment loss that pushes the encoder’s output embedding closer to the quantized embedding.

This loss term is meant to optimize reconstruction while also improving the quality of the quantization. There is no KL term as in the standard VAE. The second and third terms replace the KL term as it forces the encoder to match the prior of a categorical distribution.

We also want the quantization to be close to the prior, which is a uniform categorical distribution of the latent codes. We can consider the KL divergence of the quantized encoder with this prior:

From the prior we get \(p(z=k) = \frac{1}{K}\). From the quantization operation we get \(q(z=k^*\mid x) = 1\) for some code \(k^*\), and 0 for all other values. We can use this to remove the summation.

We arrive at a constant value which we can ignore in the optimization. However, we still have a potential issue with VQ-VAE in code imbalance, where certain codes are used more often than others. This could lead to posterior collapse.

One important detail to note is that the quantization operation is non-differentiable. The authors use a straight through estimator to address this. This simply involves copying the gradients from the codebook embedding \(z_q(x)\) to the unquantized embedding \(z_e(x)\). As the model trains, these two embeddings become closer together and the gradients become more accurate.

Generation

In VQ-VAE we don’t sample from the categorical distribution of codes directly. The encoder maps each image to a fixed number of codes (ex: 128) with an order. We can use these codes to train an autoregressive model like PixelCNN. We can sample latent codes from this model through ancestral sampling. These codes are then fed into the VQ-VAE decoder to generate an image. VQ-VAE is powerful in that it enables usage of powerful autoregressive models. More recently it is common to use transforms to model in the latent space.

Discrete VAE

The discrete variational autoencoder (dVAE), used in OpenAI’s image generation model DALLE, provides an alternative to vector quantization. In VQ-VAE the encoder maps the images to a grid of vectors which are then quantized. In dVAE, the encoder maps the inputs directly to the discrete encodings. Instead of a grid of vectors, we get a grid of scalars.

The encoder first maps inputs to grids of vectors, as done in VQ-VAE. However, these are treated as one-hot encoding from a categorical distribution. The argmax is then taken to discretize this encoding into a scalar. These scalars can be auto-regressively modeled and then sampled from.

Gumbel-Softmax

In order to make this trainable, we use the Gumbel-Softmax relaxation. This relaxation for VAEs is introduced concurrently in the papers Jang et al. 2016 and Maddison et al. 2016 (Eric Jang blog post). The Gumbel-Softmax can be defined as follows:

\(i\) is the index in the one-hot encoding, which corresponds to the number of discrete codes. \(\pi_i\) is the probability of code \(i\) in the encoding.

\(g_i\) is a sample from the Gumbel distribution, also referred to as Gumbel noise. It is sampled as follows: \(g = -\log(-\log(u))\) where \(u \sim \text{Uniform}(0,1)\). This is applying the reparameterization trick to sampling from a categorical distribution. \(g_i\) is analogous to \(\sigma \epsilon\) in VAE. This noise also makes sure that codes that aren’t the max are selected and trained on.

This argmax can be approximated as a softmax:

During training the temperature \(\tau\) is gradually reduced. As it approaches 0, the softmax approaches the argmax. We can then use an argmax for inference / generation. The temperature is annealed during training and not used during inference. At inference, we also do not need to add the Gumbel noise.

dVAE and VQ-VAE are similar architectures with a key difference in how they discretize the data.

Conclusion

This blog post explained the evolution of autoencoders to variational autoencoders to vector quantized variational autoencoders. This just covers the fundamental mathematical intuitions. The VAE and vector quantization are malleable frameworks that can be used in a variety of settings and within many variations of architecture. The VQGAN builds upon VQ-VAE by adding a discriminator to improve the image generation. VQ-VAE-2 improves upon the original by using a hierarchical approach with multiple levels of latent codes to capture both fine and coarse details. Although these models are most known for image generation, they are also very useful for discrete representation learning. This blog post focuses on the VAE and the encoding of images. There are also methods of additional conditions to the image generation, such as text prompts.

This post will focus on changes to the Transformer architecture that have been popular in the last couple of years (as of 2024). Time will tell which of these will be relevant in the future. For a deep dive on the original transformer architecture from Attention is All You Need (2017), see part 1 of this blog post. ML researchers and engineers use a lot of jargon when discussing transformers. This blog post seeks to elucidate it.

The original transformer architecture is very powerful and can handle a very wide range of applications. However, there are some limitations that more recent research has sought to address:

• Long context
  • Transformers are limited by their \(N^2\) computational complexity. This limits the sequence length a transformer can process, especially under resource and latency constraints.
• Multimodality
  • While transformers were originally developed for language processing, which remains their primary application, they have since expanded to handle nearly all data modalities and power large multimodal models.
• Efficiency
  • Scaling is attributed to the recent successes of LLMs. To enable this scaling, the model architecture needs to be designed efficiently.
  • There have also been some research into more efficient computation. The same operations can be rewritten to run faster on hardware.

This blog post will cover the modifications to the architecture that are popular in the ML community and are applied in cutting edge industry applications such as LLMs (GPT, Claude, Gemini). This will give the prerequisite knowledge to understand these SOTA models. Many of these optimizations are only popularly used in LLMs. I find that transformers in other modalities, such as vision transformers, tend to stick to more vanilla architectures and training procedures.

While these architectural improvements are significant, they aren’t the main driver of AI’s remarkable progress in recent years. The real paradigm shift has come from post-training techniques. Specifically, fine-tuning and instruction tuning methods have enabled us to leverage language models more effectively for a variety of use cases. I initially planned to cover these techniques here, but as this post grew longer, I realized they deserve a part 3 to do them justice. For now, we’ll focus purely on the architectural changes to the model itself.

Many of the topics covered in this blog are large and active areas of research. These may warrant their own blog posts to give the research justice. However, in this post, we will take a myopic view and not cover this breadth of research. This will allow us to focus on more aspects of modern transformer applications and how they relate to each other.

To understand the current state of the art, we will look at technical reports for foundational LLMs: GPT, Gemini, LLaMA, Claude. We will also look at recent papers in other modalities to see how the transformer is adapted or evolved for these use cases. However, we find that there are generally fewer architectural changes in other domains.

Unfortunately many of the most used transformer models are opaque in that the public has access to very limited details of their architecture. Even open source models usually only make the inference code open, while keeping the training code closed. However, the following models can still help us understand the current state of the art:

• Open Source LLMs
  • LLaMA (2023): LLaMA, LLaMA 2, LLaMA 3
  • Gemma (2024): Gemma, Gemma 2
  • Mistral (2023): Mistral 7B, Mixtral of Experts
  • Qwen (2023): Qwen
  • Phi (2024): Phi-3, Phi-4
  • DeepSeek (2024): DeepSeek LLM DeepSeek V2 DeepSeek V3
  • OLMo (2024): OLMo, OLMo 2
    • Open training code!
• Text encoder LLMs
  • ModernBERT (2024)
• Open Source Vision Language Models
  • PaliGemma (2024): PaliGemma, PaliGemma 2
  • DeepSeek-VL2 (2024)
• Computer Vision Models
  • ViT-22B (2023)
• Many more!

Through looking at these technical reports, you find that they share many of the same architectural improvements. This blog seeks to explain these improvements. I believe it is likely that the closed source transformers (GPT 3+, Gemini, Claude) aren’t that different.

Normalization

In part 1, we covered Pre-LayerNorm. We also see in architectures such as GPT that layer norm is added before and after the feed-forward and attention layers. This change was introduced in GPT-2. Generally, adding more normalization is inexpensive and has little downside, but can significantly improve and speed up training.

There are also new types of normalization being used. RMSNorm is currently a popular variant, used in LLaMA.

To refresh, layer normalization is defined as follows:

RMSNorm is defined as:

The notation we use for RMSNorm is changed to be more aligned to the LayerNorm notation. \(x\) and \(\gamma\) are both vectors, and the division by \(\mathrm{RMS}(x)\) is broadcasted to each element of \(x\). \(\gamma\) is a vector or learned scaling parameters for each value of \(x\). It is written as \(g\) in the RMSNorm paper.

The difference is that RMSNorm rescales the features, while LayerNorm rescales and recenters the features. RMSNorm is designed to be shift invariant and not depend on the mean or recenter the values in any way. The mean subtraction is dropped from the numerator. \(\beta\) is also dropped to avoid a learned recentering. In LN, variance \(\sigma(x)^2\) depends on the mean (\(\frac{\sum_{i=1}^{N}(x_i - \mu)^2}{N}\)). RMS is used in place of variance to be completely independent of the mean. RMSNorm is completely shift invariant.

RMSNorm is more efficient than LayerNorm. The memory usage is halved because there is only one learned parameter \(\gamma\), since \(\beta\) is dropped. There is also less computation since there is no mean subtraction. These small efficiency gains become significant when applied to transformers with hundreds of billions of parameters. Another possible advantage is that it also constrains the model less, while still keeping the advantages of normalization.

Activation Functions

The original transformer used ReLU activation functions, which appear only between the two layers of the feed-forward block. Since then, two other activation functions have gained popularity: GeLU and SwiGLU.

Research into activation functions remains largely empirical, with limited theoretical understanding of why certain functions outperform others. The research community has spent years optimizing every component of the transformer architecture, including the ReLU activation function. These functions should be considered for any new transformer application, but one should only expect a modest improvement in performance. Nevertheless, their mathematical formulations are quite interesting and worth understanding.

“We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence.” - Noam Shazeer in “GLU Variants Improve Transformer”

Gaussian Error Linear Units (GELUs)

GeLU is formally defined as follows:

\(\Phi(x)\) or \(P(X \leq x)\) is the standard Gaussian cumulative distribution function (CDF). It is the probability of values less than \(x\) occurring in a standard normal distribution: \(\Phi(x)=P(X \leq x); X\sim\mathcal{N}(0,1)\). GeLU multiples the input by this probability.

The CDF is defined as follows:

The error function is defined as:

The integral in \(\mathrm{erf}\) needs to be approximated for practical use. There are two ways to approximate GeLU:

GeLU can be interpreted as a smooth version of ReLU with the added property of non-zero gradients for negative inputs.

The GPT-2 code uses the longer and more exact approximation. The shorter approximation is much more computationally efficient. Nowadays, this is more popular because the difference in model performance is negligible.

GeLU is used in BERT and GPT models. It is also popular for vision transformers such as ViT-22B.

SwiGLU

SwiGLU builds on the Swish activation function. This is also referred to as SiLU (Sigmoid Linear Unit) in the GeLU paper.

SwiGLU is a gated linear unit (GLU) version of this. GLUs were introduced in 2016. GLUs modify the feed-forward network by introducing a gating mechanism that controls information flow, allowing the network to selectively emphasize or suppress different parts of the input. This works by adding another linear transformation of the input \(Vx\) that acts as the gating function. The gating function performs element-wise multiplication with the output of the first feedforward layer and activation function. SwiGLU is a GLU that uses Swish as the activation function.

GLUs are more expensive since \(V\) is essentially a third linear layer in the feed-forward block, the same size as \(W_1\). To ameliorate this cost, the SwiGLU paper sets \(d_{ff}\) to \(\frac{2}{3}4d_{model}\), instead of \(4d_{model}\). \(\frac{2}{3}\)is chosen to keep the total number of parameters in the feed-forward block constant, since we have three weight matrices of the same parameter count instead of one. LLaMA also uses this ratio. However, this additional matrix multiplication is in parallel to the first linear layer. Depending on the hardware, it may have negligible training speed implications. The gating mechanism improves the performance of the model enough to warrant a reduction in the width of the feed-forward block.

GLU also increases the expressivity of the linear layer. The first linear layer is able to represent second order polynomials like \(x^2\), due to the element-wise multiplication. This may compensate for reduced width.

SwiGLU is used by many recent LLMs such as Mixtral/Mistral, LLaMA, and Qwen. Gemma uses GeGLU, which was introduced in the same paper as SwiGLU (GLU Variants Improve Transformer). GeGLU is a gated linear unit (GLU) variant of GeLU (GELU in place of Swish).

Position Embedding

In part 1, we covered two basic implementations of position embeddings tested in the “Attention Is All You Need” paper: learned position embeddings and sinusoidal position embeddings. These absolute position embeddings are functions of each token’s position in the sequence. While these methods are straightforward, they have notable limitations.

One key challenge is generalizing position embeddings to longer context lengths than those used in training. LLMs are typically pretrained with a fixed context length (e.g., 4096), but we often want to use them with longer sequences during inference. Learned position embeddings can’t generalize beyond their training sequence lengths since they lack embeddings for later positions. This is why sinusoidal position embeddings, which can naturally extend to any length, have become more popular.

Another limitation is that absolute position embeddings are permutation invariant, which isn’t ideal for many applications. This has led to research into relative position embeddings, which encode the positional relationship between pairs of tokens rather than their absolute positions in the sequence.

Relative Position Encoding

Relative position encoding (RPE) uses the distance between tokens to define an embedding. Absolute position embeddings are simple in that they only require adding an embedding to the input token embeddings of the model. Each token maps to one embedding. With relative position embeddings, there is an embedding for different pairs of tokens. This requires computing the embeddings in the attention operation.

For a sequence of length n, there are \(2n-1\) different pairwise distances ranging from \(-n/2\) to \(n/2\), including 0. The distance can also be clamped by a max distance to reduce the number of parameters. These distances are mapped to learned embeddings.

RPE was introduced by Shaw et al. 2018. This particular implementation can be referred to as the Shaw relative position embedding. Rather than using matrices, the paper describes how individual attention scores, for query token \(i\) and key token \(j\).

The following equation describes how relative position is encoded into attention:

Equation (1) describes how the attention scores are calculated. Equation (2) shows how relative position is encoded into the values. It essentially reweighs the values based on relative positions. \(\alpha\) is just \(e\) after softmax along the key dimension. The change from standard attention is the addition of \(a_{ij}^K\) and \(a_{ij}^V\). These are learned embeddings that bias the keys and values based on the relative position between \(i\) and \(j\). The bias on the key is different based on which query it will be in a dot product with. The embedding lookup is based on the difference between \(i\) and \(j\) and a max distance parameter \(k\). There are separate embedding tables for keys and values: \(w^K\), \(w^V\).

Relative position is encoded in the keys so that the attention matrix can capture positional relationships between tokens. The position can also be encoded in the values, so that positional information can be encoded in the embedding themselves. This makes sense since the attention mechanism is an information bottleneck since the scores are scalars, whereas values are embeddings. However, the paper reports that they get the best results from only adding relative position the keys. Relative position on the values does not add any additional performance but comes with an efficiency cost. We see subsequent works drop this term and focus on incorporating relative position in the attention matrix.

Now that we have defined how we can incorporate relative position within the attention operation, we must understand how to efficiently implement this. The notation above doesn’t use matrix multiplies and we cannot naively iterate over \(i\) and \(j\) with for loops. To implement the Shaw relative position embedding, \(a^K\) and \(a^V\) are computed as two matrices of size \((N, N, h, d_k)\) and \((N, N, h, d_v)\), result in \(O(N^2hd)\) memory utilization. The memory usage is reduced to \(O(N^2d)\) by sharing these embeddings between each attention head (this is also in line with absolute position embedding). The notation above disregards the attention heads dimension for simplicity. We can formulate attention with relative position as follows:

\(S^{rel}\) represents the relative attention bias which is calculated by initializing a matrix of size \([N,N,h,d_k]\). This contains embeddings of size \(d_k\) from each pair of \(N\) positions. This matrix adds significant memory complexity and contains repeated embeddings.

Huang et al. 2018 introduce in their Music Transformer paper another optimization to RPE to bring the memory complexity down to \(O(Nd)\). The key intuition is that there are only \(2n-1\) or \(O(N)\) possible relative positions possible. The Shaw RPE matrices include many duplicate embeddings because the same pairwise distances repeat many times in the model. We can restructure the computation so that these repeated embeddings aren’t stored in memory at the same time. This implementation does not include relative position in the values.

Music Transformer optimizes this by multiplying each query directly with the relative position embedding matrix: \(QE^{r\top}\). This produces a matrix of size \([N,N]\), matching the expected size of the attention matrix bias. \(E^r\) represents the embedding table as a matrix storing representations for every relative position (pairwise position difference). The matrix is then transformed into the correct attention bias \(S^{rel}\) through a series of operations which they refer to as a skewing procedure. While these transformations are confusing, the key point is that we avoid creating any \(O(N^2d)\) matrices. Instead, \(E^r\) only requires \(O(Nd)\) additional space, with all other matrices matching the attention matrix size of \(O(N^2)\).

1. Mask the top left triangle, this will be shifted to result in a causal mask
2. Pad an empty “dummy” column on the left side
3. Reshape the matrix from \([L,L+1]\) to \([L+1,L]\)
   1. Row-major ordering would cause this to result in the values to be shifted such that the relative positions are aligned to the diagonals.
4. Remove the empty first row through slicing

Since relative position is encoded into the attention layers, it possible to use it alongside absolute position embeddings. Relative position embeddings are only required to be used in the first attention layer, since the information can propagate to subsequent layers. However, it is more common to use relative position attention in all of the transformer blocks. This may ensure that the position information is not disregarded by the model. It is much less common to add absolute position embeddings in every block.

Relative position embeddings add an inductive bias of translation invariance to the models. Given the scale of the datasets modern transformers are trained on, it seems like this bias would limit the model’s ability to learn. However, as with many results in this post, relative position embedding has empirical advantages in language modeling.

Attention Biasing

Relative position embeddings initialize embeddings for each relative position. However this can be simplified by mapping each relative distance to a scalar bias to the attention scores.

ALiBi

Attention with Linear Biases (ALiBi) is another method of position embedding introduced by Press et al. 2021.

\(m\) is a hyperparameter that can be set per attention head. No absolute position embeddings are added. This method adds a strong inductive bias for locality. The attention scores decay with relative distance. Since there are no learned parameters for position this bias can generalize to longer sequence lengths during inference. The biases can be extended indefinitely.

T5

The T5 models introduced by Raffel et al. 2019, also utilize relative biases, but make these learned parameters that are shared among all the layers. Rather than setting a single hyperparameter \(m\), each relative distance maps to a learnable scalar bias value. This is more parameter efficient when compared to RoPE since the biases are shared across all layers, and they are scalars rather than embeddings.

RoPE

While relative position embeddings improve model performance, they require modifying the attention operation with position bias, adding computational complexity and memory overhead. Rotary Position Embedding (RoPE), introduced in the RoFormer paper by Su et al. 2021, offers an elegant solution by encoding relative position directly into the queries and keys. This approach keeps the attention operation unchanged by leveraging mathematical properties to achieve the same benefits more efficiently.

The attention matrix contains dot products between query and key embeddings: \(\langle q_m,k_n \rangle\) . The key property of relative position embeddings is that this dot product does not depend on the absolute positions \(m\) and \(n\), but rather the relative position \(m-n\). This is translation invariant, meaning that the output is the same if you shift \(m\) and \(n\) by any constant amount. RoPE exploits the fact that we do not care what information is in the queries and keys, we only care about the dot product.

RoPE groups the embedding dimensions into groups of two adjacent indices. These two values can be considered to be defining a complex number. The first number represents the real part, and the second is the imaginary dimension.

If the embeddings are of size \(d\), \(d/2\) of these size two blocks / complex numbers are present. The absolute position of the embedding in the input sequence is denoted \(m\).

Each of these groups representing a vector will be rotated. There is a different angle \(\theta\) of rotation for each position \(m\) from 0 to \(d/2\).

The value 10000 follows the sinusoidal position embedding in Attention Is All You Need and is equivalent to \(\theta_1\). Each group is rotated by the angle \(m\theta_i\). The rotation matrix is defined as:

Given that \(m\) is a variable index. The \(\theta\) values define frequencies. Each group’s vector is rotated at a different frequency. This helps the model learn short and long range relationships between tokens.

We want to multiply a vector containing \(d/2\) groups by different rotation matrices. To accomplish this, the different rotation matrices are placed on a diagonal. This matrix is sparse and the rotation can be efficiently computed.

We can multiply this matrix by the query \(Q\) and key \(K\) matrices to encode the positions. The standard attention formula can be applied to these rotated queries and keys. We have defined the rotation and the implementation of RoPE but now we’ll show how this actually encodes position.

Let’s take a arbitrary query embedding \(q_m\) and a key embedding \(k_m\). We can assume an embedding dimension of 2 so that the query and key map to only one complex number. The multiplication with the 2d rotation matrix is equivalent with multiplying with \(e^{im\theta}\) and \(e^{in\theta}\). Multiplying a complex number \(z = x + iy\) by \(e^{imθ}\) effectively rotates \(z\) counterclockwise by an angle \(m\theta\). \(x\) and \(y\) are the two components in the embedding group. Expanding the product:

This is equivalent to multiplying with the 2d rotation matrix, assuming the first value \(x\) in the group represents the real component and \(y\) is the second and imaginary component.

The query embedding is rotated by \(m\theta\) and the key is rotated by \(n\theta\).

In attention we are interested in the dot product of these rotated vectors. The dot product of the vectors is equivalent to the inner product of the complex numbers the 2d vector represents. The complex inner product is defined as \(Re[ab^*]\). We take the conjugate of the rotated key embedding:

We take the conjugate value of \(b\) and take the real part of the product. This can be applied to the query key product:

This shows that we have reached our desired property that the dot product depends on the relative position \(m-n\), but not the absolute positions \(m\) and \(n\). This is also translation invariant in that you can shift \(m\) and \(n\) by a constant amount and not change this dot product.

Unlike RPE, RoPE doesn’t require modifying the attention operation. It just requires transforming the input query and key matrices. This makes RoPE easier to efficiently implement. RoPE also has no learned embeddings. Attention is All You Need shows that sinusoidal position embeddings, with no learned parameters, performs just as well as learned position embeddings. RoPE can be thought of as an extension of sinusoidal embeddings that prefers relative position information.

We’ll now revisit the \(\theta\) parameters to understand how they affect long context inference. \(10000\) can be considered the base frequency \(b\). This is an important hyperparameter that determines the lowest and highest frequency rotations. Given the embedding dimension \(d\), the lowest frequency rotation is \(\theta_{d/2} = b^{2/d-1}\). With the settings \(b=10000\) and \(d=768\), this corresponds to a wavelength of about 9763 tokens. This means for long context LLM inference (>32k tokens), this rotation would become periodic. ****This directly affects the sequence lengths that RoPE can handle effectively. If \(b\) is set too low, RoPE becomes periodic during longer context inference and performance suffers.

RoPE is still periodic like the sinusoidal absolute position embedding in the original transformer. However the difference is that it is periodic in the relative positions, not the absolute positions. The translation invariance significantly improves performance, but there is still room for improvement in the periodicity of the relative positions. Several approaches can be used to help RoPE perform better at long context inference.

Positional Interpolation

Tian et al. 2023 introduces a simple method to improve RoPE’s extrapolation capabilities. For models trained with context length \(L\) that need to handle longer contexts \(L'\) during inference, RoPE can be rescaled using:

This rescaling reduces the rotation applied to each token embedding. Though this approach requires finetuning for the new context length, the authors show that just 200 finetuning steps yield strong results.

Another solution is to pretrain with scaled RoPE from the start. This approach was explored in Zhang et al. 2023 and implemented in Llama 3. By setting the base frequency \(b\) to 500,000, they achieved effective inference with context lengths of 32k. While this offers a simple solution, it may limit the model’s ability to learn proper processing of higher relative positions during training.

YaRN

Peng et al. 2023 improves on positional interpolation (PI) by scaling different RoPE frequencies differently. They hypothesize that PI hurts high frequency dimensions by changing the wavelengths by a significant portion. They adapt to this by scaling high frequencies less, and low frequencies more. As we previously discusses, the low frequency components are the most important to avoid periodicity.

This method in used in DeepSeek-V3 and Qwen2.

Additional Resources

• Self-Attention with Relative Position Representations – Paper explained - AI Coffee Break with Letitia
• Relative Positional Encoding - Jake Tae
• Understanding Music Transformer - gudgud96
• Rotary Embeddings: A Relative Revolution - EleutherAI
• You could have designed state of the art Positional Encoding - fleetwood.dev
• Rotary Positional Embeddings (RoPE) - labml.ai
• Relative position embeddings according to T5 paper - AliHaiderAhmad001
• Extending the RoPE - EleuterAI
• LLaMA explained: KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLU - Umar Jamil

Efficient Attention

The \(N^2\) complexity of self-attention has long been considered a major bottleneck. There has been a lot of research into sparse attention methods, such as BigBird and Sparse Transformers. While these methods enhance efficiency through sparse attention patterns, they come at the cost of model performance. Because of this, we don’t see these methods adopted in frontier LLMs.

Sliding Window Attention (SWA) is another sparse attention method introduced in Longformer. Instead of having a query token attend to all prior key tokens, it attends only to \(w\) (window length) prior tokens. This reduces attention complexity from \(N^2\) to \(Nw\). However, effective performance requires a substantial window length. For example, Mistral uses a window length of half the context length. This is implemented through an attention mask during training. During inference, the KV cache (explained later) size can be reduced from \(N\) to \(w\).

SWA resembles convolution in that tokens attend only to nearby tokens. Longformer employs smaller window sizes in earlier transformer blocks and larger ones in later blocks. This enables both local and global information flow, allowing learning at different scales—similar to CNN architecture.

A key limitation is that distant tokens cannot communicate directly. Information must flow through multiple transformer layers to connect distant tokens. To address this, models like Gemma 2 and ModernBERT alternate between SWA and global attention blocks. By maintaining global attention blocks throughout the network, distant token can still communicate quickly. We can adjust the number of SWA blocks and window size to manage this efficiency tradeoff.

SWA has emerged as a powerful and practical method for improving transformer efficiency. It has been gaining more traction in LLMs than other sparse attention approaches. SWA typically uses a large window size - usually half the global context length. This makes it more similar to global attention than sparse/local attention. However, it still offers substantial benefits. It halves both the KV cache size during inference. Additionally, it reduces the attention computation during both training and inference. SWA also avoids the issues of position embedding extrapolation by capping the maximum relative positions in the attention operation.

FlashAttention

FlashAttention (Dao et al. 2022) is an optimization to the computation of the attention operation. They are designed to be completely accurate in their results, but more efficiently run on GPUs. We will focus on FlashAttention-2, which has several significant improvements. FlashAttention-3 is another optimization, however, it is designed specifically for NVIDIA GPUs.

When designing algorithms like attention, we often only consider the number of floating point operations (FLOPs). However, we need to more carefully consider the hardware. Specifically, GPUs are very fast in computing highly parallelized floating point operations, but have slower memory access. We can speed up attention by reducing the amount of memory reads and writes, even if it comes at the cost of additional computation.

GPUs have two types of memory: slow but large HBM (high bandwidth memory) and small but fast SRAM (static random access memory). When working with large matrices, we load blocks (parts of the matrices) from HBM into SRAM. The computation takes place and then the output is written back to HBM.

We will review two prerequisite topics before covering FlashAttention.

Block Matrix Multiplication (BMM)

Block matrix multiplication enables matrix multiplications to be parallelized by splitting the inputs into blocks and concatenating output blocks. If we are trying to compute \(C=AB\), a certain output block requires a subset of rows from \(A\) and subset of columns from \(B\). We can just take the input blocks containing these rows and columns, rather than reading the entire matrix into memory.

Each output block \(C_{ij}\) is computed using only the required blocks: \(C_{11} = A_{11}B_{11} + A_{12}B_{21}\). The shapes of these blocks need to be carefully defined to be compatible for the matrix multiplication, and be efficient on hardware. The number of rows in the blocks of \(A\) must match the columns of blocks of \(B\).

Computing an output block doesn’t require reading all the required input blocks at once. They can be read iteratively since it is a sum. When viewing matrix multiplication as dot products between rows and columns, the summation works by splitting these vectors into blocks, then summing the dot products of these blocks.

Parallel GPU cores can compute different output blocks while only reading the required input blocks. Attention is just two matrix multiplications \(S = QK^\top\) and \(O = P V\) and a softmax. FlashAttention takes advantage of the BMM algorithm to compute this more efficiently.

Online Softmax

With BMM, we can compute matrix multiplications block by block without needing any global information. However, there is a challenge in integrating softmax into this framework since it requires normalizations which need to be calculated globally from all blocks.

FlashAttention uses an online softmax to minimize the memory access. Recall the softmax operation:

The denominator \(l\) is the normalization factor. Also, softmax function is numerically unstable because exponents can reach extremely high values when \(x_i\) is a large positive number. To achieve numerical stability, we subtract the maximum value before computing the exponentials:

This subtraction keeps each exponent between 0 and 1, ensuring numerical stability. The result remains unchanged since this operation is equivalent to multiplying both numerator and denominator by the same constant. However, this maximum value is another global value.

A naive implementation would require 3 passes, one to find the max value, one to compute the numerator and accumulate the normalization factor, and one to apply the normalization factor. This requires a lot of memory access. However, we can optimize by fusing the first two steps into a single \(O(N)\) pass using this algorithm:

1. Assume the first element is the max \(m_j=s_{0j}\) and initialize an accumulator for the normalization constant as \(l_j := e^{s_{0j}-m_j}\)
2. Iterate through the rest of the list
   1. If the value \(s_{ij}\) is less than or equal to the current max, add it to the accumulator: \(l_j := l_j+e^{s_{ij}-m_j}\).
   2. If the value \(s_{ij}\) is greater than the current max, adjust the normalization factor (which was computed using the wrong max) by multiplying: \(l_j := l_j * e^{(s_{ij}-m_j)}\). This correction factor ensures \(s_{ij}-m_j\) is added to each incorrect exponent in the sum. Then update the max value: \(m_j:=s_{ij}\).
   3. This entire update can be expressed without conditional logic: \(m_i := \max(m_{i-1}, s_{ij});\quad l_j := l_j * e^{(m_{i-1}-m_{i})} + e^{(s_{ij}-m_{i})}\)

Through this optimization, we’ve reduced the softmax computation from three linear passes to two. The correctness of this algorithm can be simply proved by induction. The method is defined as iterating over individual items in \(x\), however, it easily extends to iterating over blocks in \(x\). We compute the local softmax within the blocks and then accumulate a global normalization factor.

There are two main changes to attention implemented in FlashAttention: tiling and recomputation.

Tiling

We can reduce the number of memory reads and writes by fusing operations together. With two separate operations, the data will be read twice and written twice. If the operations can be fused together, there will be only one read and write. However, we have to consider that we are working with very large matrices and using highly parallelized computation. Block matrix multiplication along with online softmax enables us to compute attention in pass without writing the full matrix to memory.

With tiling, we restructure the attention operation to load data block by block from HBM to SRAM. The full attention matrix of size \((N,N)\) is never materialized at once or written to HBM. This is done by looping through the blocks multiple times to avoid the need to write to memory. This is increasing computation to reduce memory access.

With fused operations you are trading accessibility of the code for efficiency. Simple matrix operations for attention (like multiplications and softmax) are easy to understand and modify. However, when running attention at scale, implementing an optimized fused version makes practical sense. Currently, billions of dollars are being spent running attention. It is important that we make this operation as efficient as possible.

Recomputation

Typically, the attention matrix is stored in HBM during the forward pass for use in the backward pass. FlashAttention takes a different approach. It recomputes the attention matrix during the backward pass instead of storing it. While this requires computing the attention matrix twice, it eliminates the need to write the large matrix to memory.

The softmax normalization parameters from the forward pass are still stored, enabling the backward pass attention computation to be completed in a single pass.

logsumexp

For each example in the batch, we can save only one value for the softmax normalization instead of two. We need to save both the max value \(m_i\) and the normalization sum \(l_i\). We can store the logsumexp \(L_i = m_i+\log(l_i)\). We’ll show how we can compute a correct softmax with this.

This small trick reduces the memory writes from the forward pass by half.

Algorithm

Forward Pass

1. Iterate through query blocks to get \(\text{Q}_i\) (this step runs in parallel)
2. For each query block, iterate through key and value blocks to obtain \(\text{K}_j\) and \(\text{V}_j\)
   1. Compute the attention block \(\text{S}_i^{(j)}=\text{Q}_i\text{K}_j^T\)
   2. Update the maximum value \(m_i^{(j)} = \max(m_i^{j-1}, \text{rowmax}(\text{S}_i^{(j)}))\), where rowmax finds the maximum in the current block
   3. Compute the unnormalized attention weights for the current block using the current maximum: \(\tilde{\text{P}}_i^{(j)}=\exp(\text{S}_i^{(j)}-m_i^{(j)})\)
   4. Update the block output: \(\text{O}_i^{(j)} = \text{diag}(e^{m_i^{(j-1)}} - m_i^{(j)})^{-1} \text{O}_i^{(j-1)} + \tilde{\text{P}}_i^{(j)} \text{V}_j\)
      1. The first term adjusts the exponents from previous blocks
      2. This output remains unnormalized and contains only the softmax numerator
3. Normalize the output block and write to HBM: \(\text{O}_i=\text{diag}(l_i^{(T_c)})^{-1}\text{O}_i^{T_c}\)
   1. \(T_c\) represents the final block index after completing the loop
4. Compute logexpsum and save to HBM for the backward pass: \(L_i = m_i^{(T_c)} + \log(l_i ^{(T_c)})\)

Backward Pass

1. Initialize gradient buffers in HBM for \(\text{dQ}\), \(\text{dK}\), \(\text{dV}\)
   1. Initialize \(D=\text{rowsum}(\text{dO}\circ O)\) in HBM
2. Iterate through key and value blocks to obtain \(K_j\) and \(V_j\)
   1. Initialize \(\text{dK}_j\) and \(\text{dV}_j\) in SRAM
   2. Iterate through the output blocks and load \(\text{Q}_i\), \(\text{O}_i\), \(\text{dO}_i\), \(\text{dQ}_i\), \(L_i\), \(D_i\) to SRAM.
      1. Recompute the attention block \(\text{S}_i^{(j)}=\text{Q}_i\text{K}_j^T\)
      2. Normalize the attention using the saved logsumexp \(\text{P}_i^{(j)}=\exp(\text{S}_i^{(j)}-L_i)\)
      3. Calculate the following gradients
         1. \[\text{d}\text{V}_j \leftarrow \text{d}\text{V}_j + (\text{P}_i^{(j)})^\top \text{dO}_i\]
         2. \[\text{d}\text{P}_i^{(j)} = \text{dO}_i \text{V}_j^\top\]
         3. \[\text{d}\text{S}_i^{(j)} = \text{P}_i^{(j)} \circ (\text{d}\text{P}_i^{(j)} - D_i)\]
            1. This is derived from the Jacobian of the softmax
      4. Update \(\text{dQ}_i \leftarrow \text{dQ}_i + \text{dS}_i^{(j)}\text{K}_j\)
         1. This read from HBM and written back. This memory expensive step is necessary since we need to accumulate gradients for all query blocks before writing the final result
      5. Update in SRAM: \(\text{dK}_j \leftarrow \text{dK}_j + \text{dS}_i^{(j)T}\text{Q}_j\)
   3. Write \(\text{dK}_j\) and \(\text{dV}_j\) to HBM

FlashAttention is a kernel fusion that torch.compile cannot find because it is a rewrite of the computation. NxN attention matrix is not materialized. This also allows larger values of N which is important for long context. It is now implemented in most deep learning frameworks.

FlashAttention achieves its performance improvements through careful memory management and tiling of the attention computation. Rather than computing the full attention matrix at once, it processes smaller blocks of queries and keys, reducing the peak memory usage. This optimization is particularly important for training and inference with long sequences, where the quadratic memory scaling of attention would otherwise be prohibitive.

Additional Resources

• Flash Attention derived and coded from first principles with Triton (Python) - Umar Jamil
• FlashAttention - Tri Dao - Stanford MLSys #67
• Slides

Mixture of Experts (MOE)

Mixture of Experts (MoE) is transformer architectural change applied to all of the feed-forward blocks. Rather than having one feed-forward module in each transformer block, there are multiple parallel feed-forward blocks. A gating network selects among these blocks. These blocks are considered experts that specialize in different things. Multiple blocks are selected and the output embeddings are added together, creating a mixture.

MoEs are considered to be sparse models. MoE models can have very high numbers of trainable parameters, but have only a fraction used for a given example during training and inference.

MoE is generally structured as follows:

1. Create \(n\) experts, which are separate versions of the feed-forward MLP block
2. A gating network / router processes the token embedding \(x\) and outputs a score for each embedding. This is a vector of size \(n\).
   1. The gating network is typically implemented as a single learned MLP layer: \(W_gx\)
3. We select the top k experts and route the token embedding to these experts.
4. The output is a weighted sum between the outputs of the k selected experts
   1. A softmax is used on the scores of the top k experts to normalize the weights

The number of selected experts \(k\) is chosen to be small (typically 2) to achieve this sparsity. With this approach and parameter setting, we achieve the goal of a sparse transformer architecture.

We want all \(n\) experts to be utilized with roughly equal frequency. In the worst case, the same \(k\) experts are selected for every input, and the remaining experts are just a waste of space. While each input should activate only a few experts, we want to ensure that all experts are activated with roughly equal frequency over the entire training set. We want to encourage diversity in expert selection. Different MoE implements achieve this differently.

Sparsely Gated MoE (2017)

Shazeer et al. 2017 introduces mixture of experts for language modeling, however, using LSTMs instead of transformers.

In order to enforce diversity, they add tunable Gaussian noise to the expert score \((xW_g)_i\) of each expert in the function \(H()\).

The noise for each component is learned through \(W_{noise}\). A load balancing loss is optimizes the magnitude of the loss for different experts. Without a proper loss term, the noise would collapse to 0. This noise term is important for the load balancing loss.

Softplus is similar to ReLU but smooth and always non-negative: \(\text{Softplus}(x)=\log(1+e^x)\). This is just to make sure the noise added is non-negative.

After adding the noise, we apply the \(\text{KeepTopK}\) function. This sets the values to \(-\infty\) for experts outside of the top \(k\) to ignore them in the softmax. After the softmax, only two experts have weights, which add up to 1. The final weights on the experts is calculate as:

The softmax normalizes the weights of the two experts, which can be used to compute a weighted sum of the outputs of the selected experts.

For diversity, an additional importance loss term is used. Given an expert \(i\) and a training batch \(X\), we define importance as follows:

We then sum this across experts to define an auxiliary loss, which is minimized when all the experts are activated equally on the batch:

However, one shortcoming of the importance loss term is that it uses the weights. It is possible for an expert to have a high average weight, but never be selected in \(KeepTopK\). They define another loss \(L_{load}\) weighted by \(w_{load}\) to address this. We want a smooth estimation of the number of examples assigned to each expert. The definition of this is out of scope for this post, but it is designed to be smooth operation on a discrete operator.

A key challenge with MoE is that each expert only sees a fraction of the training batch, limiting its learning effectiveness. To address this, one solution is to use data parallelism with very large batch sizes.

MoE also requires careful handling in multi-device training setups. Through expert parallelism—a form of model parallelism—different devices manage different experts. After the gating function determines expert assignments, tokens are shuffled between GPUs in a ring pattern before returning to their original GPU. This allows us to use MoE to scale the number of parameters in a model, but there is an added communication cost with expert parallelism.

The auxiliary loss is calculated separately for each MoE layer in the transformer.

GShard

GShard (Lepikhin et al. 2020) applies MoE to transformer, also with k=2 experts. Each expert implements the same two-layer MLP architecture used in standard transformers. GShard uses a simpler auxiliary loss term to encourage expert diversity. It is simply \(\frac{c_e}{S}m_e\) where \(e\) is a particular expert. This loss is averaged across all experts. \(c_e/S\) is the fraction of the \(S\) tokens routed to expert \(e\). This is not differentiable because \(c_e\) is determined by the top k operator. This ratio acts as a weight on the mean weight of the expert \(m_e\) which is differentiable. The authors were able to train very large MoE transformer models while using this loss to balance tokens between experts.

The expert with the highest score is selected first. For the second expert, rather than choosing the next highest score, the system randomly samples from the remaining experts. This random selection promotes diversity, functioning similarly to the Gaussian noise term.

Switch Transformer

Fedus et al. 2022 applies MoE to transformers with some simplifications compared to the 2017 paper. They only use one expert by setting \(k\) to 1. Instead of having separate loading balancing and importance losses, they use one auxiliary loss term (same loss as GShard):

\(f_i\) is the fraction of tokens routed to expert \(i\) and \(P_i\) is the fraction of router probability routed to expert \(i\). These are among all tokens in a batch. \(f_i\) is analogous to the load balancing loss and \(P_i\) is analogous to the importance loss. \(P\) is differentiable, while \(f_i\) is not. The loss works because \(f_i\) is treated as a weight on a differentiable loss. The multiplication with \(N\) ensures that the optimal loss is the same with different numbers of experts.

Frontier LLMs

Mixtral

MoEs gained significant popularity in LLMs with the Mixtral paper from Mistral AI. The model uses 8 experts but selects only 2 for each computation, with SwiGLU serving as the expert function. The MoE block operates as follows, where \(x\) represents input token embeddings and \(y\) represents output token embeddings:

While the paper doesn’t detail the exact auxiliary loss, the Hugging Face implementation extends the Switch Transformer loss to work with any number of selected experts.

DeepSeek

DeepSeek shares more details in how they train their MoE models. They do not use an auxiliary loss. Instead, they add a bias term \(b_i\) to the scores of each expert. At the end of each training step, they update the biases by \(\gamma\), which they term as the bias update speed. If the expert is overloaded, the bias is subtracted. If it is underloaded, it is added. This is heuristic method that avoids the challenges of optimizing an auxiliary loss.

Although they don’t use an auxiliary loss for load balancing between experts, they do use one to encourage expert diversity within sequences. Their goal is for tokens within a single example to use different experts. This contrasts with other methods and challenges the intuition that experts should specialize in specific topics. They find empirically that this reduces expert specialization. They implement this similarly to the Switch Transformer loss but average the probabilities within a sequence. However, a recent paper by Wang et al. 2024 suggests this loss isn’t necessary for achieving expert diversity across sequences.

They also introduce the notion of shared experts. They have routed experts which are selected by a gating function. The shared experts are always used. These are meant to store more general knowledge, while the routed experts are more specialized. This requires having more experts activated if you want to have more than 1 routed expert. DeepSeek-V3 uses high numbers of experts (ex: 2 shared, 64 routed). To implement this efficiently, they implement each expert as a 2 layer MLP with a small hidden dimension.

MoE represents one of the most significant improvements to the transformer architecture. However, the research community has yet to converge on a standard implementation. Different papers approach expert diversity in varying ways. Understanding these documented methods provides a foundation for designing your own MoE model. The technique’s success has led to applications beyond language models into other domains, such as the Vision MoE (V-MoE) paper.

Additional Resources

• Mixtral AI Talk
• Mixture of Experts Explained - Hugging Face
• Training MoEs at Scale with PyTorch
• Why Llama 3 is not a MoE? - /r/LocalLLaMA
• Bug from the Hugging Face Mixtral code explaining the auxiliary loss
• Mixtral MoE inference code
• DeepSeek-V3 MoE inference code

KV Cache

Modern transformer architectures have introduced several innovations to improve decoding efficiency, especially for handling long sequences and reducing memory usage during inference. These improvements are essential when deploying large language models in production environments with limited computational resources.

Decoding with an LLM involves two main steps:

1. Prefill: Processing the input prompt and prior context
2. Decoding: Generating new tokens

Efficiently Scaling Transformer Inference (Pope et al., 2022) explores how these decoding steps can be optimized. In language model decoding, both the context and previously generated tokens are needed to predict the next token. However, reprocessing these tokens from scratch would be inefficient. Instead, we can save the intermediate outputs of prior tokens to avoid recomputation.

Only the keys and values from the attention layers need to be preserved. Since transformer decoders use causal attention, new tokens don’t influence the embeddings of previous tokens. We can safely discard the query and output embeddings for prior tokens as they’re no longer needed for computation. The key and value embeddings are kept in the KV cache, allowing newly decoded tokens to attend to previous ones.

During the prefill step, a KV cache is generated, containing the attention key and value tensors from each layer. This cache grows as each new token is decoded.

The quadratic growth of KV cache with context length presents a significant scaling constraint. When we say a model uses a KV cache, it means the keys and values from previous tokens are preserved during the decoding process. The KV cache can have significant memory utilization as the number of tokens increases. However, the KV caching mechanism means that for the same model and hardware, more tokens can be processed during inference than during training.

MQA/GQA

Multi Query Attention (MQA) and Group Query Attention (GQA) are architectural changes to the attention layers of the model. These methods slightly hurt the model’s performance, but enable more efficient decoding.

Multi Query Attention (MQA)

Multi-Query Attention was introduced in the paper Fast Transformer Decoding: One Write-Head is All You Need.

In standard multi-head attention, each attention head has its own set of query (Q), key (K), and value (V) projections. This means that for \(h\) heads, each token will have \(h\) query, key, and value embeddings.

MQA modifies this by using a single shared key and value head across all query heads. There are \(h\) query embeddings but only one key and value embedding. All of these embeddings are of size \(d_{model}/h\). This significantly reduces the memory footprint of the model. This is especially impactful when decoding with a long context length.

MQA trades a minor decrease in model performance for substantial memory efficiency. The size KV cache size by a factor equal to the number of attention heads. For example, with 32 attention heads, MQA requires only 1/32 of the KV cache memory compared to standard multi-head attention. This enables decoding with longer contexts.

Multi query attention maybe more aptly named single key and value attention, since MHA already has multiple queries. The difference is that MQA doesn’t have multiple keys and values.

Group Query Attention (GQA)

Group-Query Attention was introduced in the paper GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

GQA is a middle ground between standard multi-head attention and Multi-Query Attention (MQA). Instead of sharing a single key-value head across all query heads (as in MQA), GQA shares key-value heads among groups of query heads.

For example, if we have 8 query heads (\(h=8\)) and 2 key-value heads (\(g=2\)), each key-value head would be shared by 4 query heads. This provides a better balance between computational efficiency and model quality compared to MQA.

There is reduced computational cost in that \(g\) keys and values need to be computed instead of \(h\). However, the main motivation is memory saving. Pope et al., 2022 shows that the memory savings of GQA enable longer context lengths. The computation of MHA might be slower. But the memory requirements can make longer context lengths impossible, due to out of memory errors. With more heads, the resource savings are further multiplied.

GQA is currently more popular as it offers a tunable tradeoff between efficiency and quality. MQA is a special case of GQA where the number of groups is 1. Both MQA and GQA take advantage of the fact that in attention the number of keys/values and queries are not required to be the same.

Conclusion

This blog post is meant to bridge the gap between “Attention Is All You Need” and frontier LLMs. While the transformer architecture has evolved significantly over the years, the original design remains remarkably relevant and far from obsolete. Very few architectural changes have been universally adopted. Individual transformer applications make different design decisions, but the transformer itself has remained universal. There have also been challengers to the transformer architecture, such as Mamba . For the time being, I expect that the transformer is here to stay.

This blog post will be in two parts:

Part 1 will be a deep dive of the standard Transformer architecture. It is a highly modular architecture, so we will explain each component in detail and how they integrate. This will also cover how to design the components for different use cases. It was introduced by the famous paper Attention Is All You Need. There is no shortage of resources to learn about transformers, but I hope to offer some new perspectives.

Part 2 will cover recent advancements that have further advanced the capabilities of transformers. The original transformer architecture is robust and versatile and has led to many successful applications. However, in recent years with the surge in investment into transformers / LLMs, we have seen many useful advances. These impart new capabilities such as longer context length, faster training, and more efficient inference. This is a guide to designing modern transformer architectures for any use case.

Transformer Architecture

Attention Is All You Need introduced the transformer architecture in 2017 specifically for machine translation. Since then, architectures derived from this have been used not only for various NLP tasks but also for other modalities such as vision, audio, and time series. We’ll take a modality-agnostic approach. As we explore each component, we’ll focus on how to design it for different modalities and use cases. For instance, position embeddings might be designed differently for text than for images.

Transformers are so generalizable because they have relatively few inductive biases. Unlike CNNs, which require Euclidean input, Transformers do not enforce a data structure. It is up to the designer to incorporate domain-specific inductive biases into the model. These design decisions are important for the model to be effective at a given task.

The transformer architecture can be understood as 3 steps.

1. Input processing (Generation of a set of embeddings to input into the transformer)
2. Transformer blocks (Bulk of the computation)
3. Output processing (Using the output embeddings of the transformer to perform a task and train the model)

Input Processing

The input to the transformer is an unordered set of embeddings. These embeddings are high dimension vectors of floating point values that represent a part of the input. We refer to input processing as the steps taken to compute these embeddings. Input processing changes the most between modalities.

The general pattern for input processing is as follows

1. Split up the input into pieces
2. Map each piece to an embedding

The output of these two steps is a set of embeddings that represent the raw input in a way the transformer architecture can process.

Image Processing

Transformers for images were introduced in the ViT paper. The input image is processed in two steps:

1. The image is split into patches of 16x16.
2. The pixel values of each patch are flattened to vector and fed through a learned linear projection, resulting in patch embeddings.

The result is a set of embeddings that can be processed by the transformer.

Text Tokenizer

Text is represented as a sequence of characters or bytes. Unlike images, text isn’t inherently numerical data that can be directly transformed into embeddings. Text is processed by tokenization, which is mapping it to a sequence of discrete tokens. Tokenizers create a vocabulary, which is mapping of all possible tokens to vocab indices. These indices are used to retrieve a learned embedding from a table. Text input processing involves two steps: tokenization and embedding lookup.

Let’s consider two basic options for tokenization:

• Character-Level Tokenization: Every character in the text becomes a separate token. This creates a very long sequence but has a very small vocabulary.
• Word-Level Tokenization: Each word is a distinct token. This results in a more manageable sequence length but results in a huge vocabulary.

There is an obvious tradeoff between vocabulary size and input size. Character-level tokenization creates very long sequences, which can be inefficient for transformers. On the other hand, word-level tokenization can lead to a massive vocabulary size, requiring a large embedding table. This can be computationally expensive and struggle with unseen words or typos.

The ideal approach considers several factors:

• Sequence Length: Shorter sequences are generally more efficient for processing, but extremely short sequences may not capture enough context.
• Embedding Table Size: A larger vocabulary requires a bigger embedding table, increasing memory usage and training time.
• Rare Words: Very infrequent tokens may not be adequately learned during training, impacting model performance.
• Token Complexity: A single token shouldn’t represent too much information. Complex concepts might benefit from being broken down into smaller tokens for better processing by the model.

In practice, finding the right balance often involves a compromise between character-level and word-level tokenization. Techniques like subword tokenization (splitting words into smaller meaningful units) can offer a middle ground, achieving a balance between sequence length, vocabulary size, and capturing text information effectively.

Byte Pair Encoding (BPE)

The most common approach to implementing sub-word tokenization is Byte Pair Encoding (BPE).

It works by first starting with the individual characters (bytes) in the text as its initial vocabulary. This ensures that all text can be encoded, though not efficiently. BPE then iteratively identifies the most frequently occurring pair of characters and merges them into a single new token. This process continues until a predefined maximum number of tokens is reached, preventing the vocabulary from becoming too large.

One interesting feature of this approach is that the tokenizer uses a small and separate dataset for BPE. This dataset can be engineered to achieve certain properties in the tokenizer. For example, it is beneficial for this data to be balanced between different languages. For example, if the amount of data for Japanese is significantly lower than that for English. Rare pairs in English would be prioritized over common pairs in Japanese. This would be unfair to Japanese, and Japanese text would require far more tokens. To address this, the tokenizer dataset can be balanced between different languages.

See platform.openai.com/tokenizer for an interactive demo on how text is tokenized and mapped in token indices.

BPE also involves some hardcoded rules. Some bytes, such as punctuation can be ignored in the tokenizer merging. GPT tokenizers use different regex patterns to split the string prior to tokenization to prevent certain bytes from merging.

BPE is expensive to run and to encode/decode text since it is an iterative process. This is optimized by OpenAI (tiktoken) by implementing it in Rust. SentencePiece by Google is another popular tokenizer. SentencePiece runs BPE on Unicode code points (Unicode characters). It falls back to bytes for rare code points. Unicode has nearly 150k code points, a large number of which are very rare. Most tokenizers use less than 100k tokens. Having 150k tokens before adding more through BPE is not practical.

Once we have a trained tokenizer, we use it to map input text to token indices. These token indices are mapped to learned embeddings. Transformer models often include embedding tables, which store learned embeddings for each item in the model’s vocabulary.

See this video by Andrej Karpathy for a deep dive into text tokenizers.

Audio and Other Modalities

Like images, audio is a continuous data modality. A popular method of tokenizing audio is to generate a spectrogram using a Fourier Transform. This creates an image that can be tokenized in the same way as images in ViT. The AST: Audio Spectrogram Transformer paper does exactly this.

This paper uses a 2D position embedding so they can warmstart from a ViT model. If it were to train from audio only, a 1D position embedding could be used, as in OpenAI’s Whisper.

We have covered the basic methods for tokenizing text and continuous data domains, however, there is a lot of research covering alternative methods. This includes vector quantization which generates discrete tokens from continuous data modalities.

Position Embedding

The transformer takes a set of tokens as input. However, many inputs are better represented as sequences, such as text and audio Where a token occurs in the sequence is a crucial piece of information. Position embeddings can be added to the token embeddings to encode the position of the token in the sequence. Although the input is still a set, we are not losing the information of the order of the tokens within the input sequence. Position embeddings implicitly turn the transformer from a set processing architecture to a sequence processing one.

The original transformer paper evaluates two methods of configuring the position embedding. These have equivalent results.

• Learned: A separate learned embedding is used for every position in the sequence up to the maximum sequence length.
• Fixed: The embedding values aren’t learned but are configured as a function of the position. \(i\) is the index in the embedding. \(d_{model}\) values have to be generated so the position embedding can be the same dimension as the token embedding.

These embeddings are added to the input token embeddings. This assumes that the addition of the positional encoding doesn’t cause conflicts in the embedding space (which is typically in a high dimension). However, it is also possible to concatenate position embedding values.

For language, 1D position encodings are used. These embeddings should be designed to fit the data. For example, in images, the position embedding is 2 dimensional. For videos, an additional time dimension can be added. The embedding can be designed in any way to encode the structure of the data.

Transformer Blocks

The transformer blocks are where the bulk of the computation takes place. We will first go through the components needed to build these blocks, and then put them together.

Attention

The core component of the Transformer architecture, as highlighted in the title “Attention Is All You Need,” is Attention. This concept predates transformers in Natural Language Processing (NLP). The fundamental equation for attention is:

Attention comes in two forms: self-attention and cross-attention. We’ll begin with self-attention. It can be conceptualized as a set-to-set mapping of embeddings where information is shared among all embeddings. Here’s how it works:

\(Q\), \(K\), and \(V\) represent queries, keys, and values. These represent different linear projections of the input embeddings that are used in the attention operation. You can think of attention as tokens communicating information with each other. Each token’s query determines which other tokens it wants to read from, while its key determines which tokens will read from it. When a token’s query matches well with another token’s key, it receives more of that token’s value. The value represents the information that the token shares with others . Through these learned projections, the model determines how information flows between tokens.

Attention is implemented by first generating these three matrices. Let’s say the input embeddings are stored in a matrix \(X\). Learned weight matrices \(W^Q\) , \(W^K\), and \(W^V\). Are used to project the input embeddings: \(Q = W^QX, K = W^KX, V = W^VX\).

At this stage, no information has been transferred between tokens. Given the variable number of tokens, we can’t use a single large MLP layer. To mix the information, we set each embedding to be a weighted sum of all value embeddings.

To compute this weighted sum, we first compute the attention matrix \(QK^T\). This matrix is of shape \((N, N)\). This is the source of the \(N^2\) complexity of transformers. The attention matrix contains scores for every combination of token query and key embeddings: \(q*k\), which is scaled by a factor \(\frac{1}{\sqrt{d_k}}\). This scaling is applied to normalize the gradients, such that the magnitude of the dot product isn’t dependent on the embedding dimension. This specific attention formulation is called scaled dot-product attention.

A softmax is applied to each row or column, creating a weight vector for each token. This is multiplied by the value matrix to generate a weighted sum for each token. Because all matrices are learned, each token can determine which tokens to attend to, which tokens should attend to it, and what information to broadcast. This function is highly flexible. A token could even learn to nullify its own information and instead read from other tokens.

Self-attention is like a fully connected neural network layer in that information from all tokens can propagate to other tokens. However, self-attention has the benefit of supporting variable length input.

Cross-Attention

Self-attention is a mechanism where queries, keys, and values all derive from the same set of input embeddings. In contrast, cross-attention operates on two distinct sets of embeddings, which can have different lengths. The queries come from one set, while the keys and values come from another.

In cross-attention, the attention matrix \(QK^T\) is of shape \((N_Q, N_K)\) . \(N_Q\) is the sequence length of the queries, and \(N_K\) is the sequence length of keys. The softmax is taken on the rows, so each query token has a probability distribution with respect to keys.

Cross-attention is particularly relevant for machine translation. In this context, the keys and values come from the source language text, while the queries come from the target language. As the model generates the translation, it attends to the set of tokens from the source language, allowing it to draw information from the original text throughout the translation process.

Masked Self-Attention

The transformer decoder uses causal masking. The decoder is trained to predict the next token. This task becomes trivial if the next token and all future tokens are visible, as in full self-attention. Causal masking constrains the attention operation to only look at tokens to the left, making the decoder auto-regressive (meaning each output depends only on previous outputs). This one-way flow of information is essential for generating sequences one token at a time.

Masking is applied on the \(QK^T\) matrix. Masked indices are set to \(-\infty\), this causes the softmax function to assign zero weight to these tokens. In many implementations of attention, the mask can be customized by passing in a Boolean matrix.

Multi-Head Attention

Multi-head attention (MHA) is a way to increase the expressivity of the attention operator. It is essentially running multiple attention operations in parallel and concatenating the output. This improves expressivity because each head is free to attend to different tokens.

Multi-head attention has two scaling parameters. Feature dimension for each head \(d_v\), and number of heads \(h\).

Each head projects the input embeddings into queries, keys, and values of size \(d_v\). This means that the weight matrices \(W^Q\) , \(W^K\), and \(W^V\) are of size \((d_{model}, d_v)\). Attention is applied to the set of queries, keys, and values independently. This results in \(h\) sets of output embeddings of size \(d_v\). The output embeddings of each head are concatenated resulting in embeddings of size \(d_v*h\). The output needs to be the same dimension as the input, so there is linear projection back to size \(d_{model}\).

Typically the embedding dimension to each head is \(d_v = d_{model}/h\). In this case, the concatenated output is the same dimension as the input token embeddings. However, it is possible to set \(d_v\) to be higher or lower.

The output projection, which is a linear layer \(W^O\), learns to combine the outputs of the different heads. The output size of this layer is the same as the input token embedding size. This layer allows the model to give different importance to different attention heads. When \(d_v\) is set to \(d_{model}/h\), this projection is not required for dimensionality matching. However, it is beneficial in that the information from different heads can be mixed before the residual connection.

Multi-head attention (MHA) effectively divides the softmax operation into separate parts. Each head has a fixed amount of attention weight to distribute among different value functions. This multi-headed approach allows for more complex token interactions. One of the advantages of the attention mechanism is its interpretability. For each head, it’s possible to examine which tokens are attending to which other tokens, providing insight into the model’s internal workings.

Normalization

In transformer architectures, layer normalization is typically used. Unlike batch normalization, the values are independent of other items in the batch. This is because batch-wide statistics aren’t used. Layer Normalization was introduced just a year prior to transformers.

For each set \(x\) in the batch, the mean \(\mu(x)\) and variance \(\sigma(x)^2\) of the embedding values (across all \(d_{model}\) values of each embedding) are calculated. LayerNorm operates on each embedding in the input completely independently. These values are used to normalize each embedding value:

\(\gamma\) and \(\beta\) are learned scalar parameters. \(\epsilon\) is a small constant used for numerical stability.

Layer norm is effective for multiple reasons. Since batch statistics aren’t used, it makes data parallelism more efficient. This is because the batch statistics do not have to be communicated between GPUs. LayerNorm is also not affected by the size of the batch, which means different batch sizes can be used at different times.

In Attention Is All You Need, the layer normalization occurs after each attention and feed-forward layer (Post-LN architecture). However, it is now more popular to put the layer normalization before these layers (Pre-LN architecture). This paper from 2020 shows that the Pre-LN architecture generally performs better. This is the only fundamental change to the original transformer architecture.

Feed-Forward

After each attention layer, a small feed forward neural network processes each token embedding. This is a position-wise operation. In the original paper, this is a two layer network with a ReLU activation after the first layer. The first layer outputs a dimension \(d_{ff}=2048\). The second layer projects these embeddings back to the original token embedding dimension \(d_{model}=512\). The first layer is set to the 4x the size of the token embedding. This multiplier is arbitrary but is considered to be an effective value given the efficiency tradeoff.

The attention layer has \(4*d_{model}*d_{model}\) parameters (accounting for query, key, value, and output projection matrices), which is \(1.0*10^6\) for the default embedding size. The feed forward layer has \(d_{model}*d_{ff} + d_{ff}*d_{model}\) which is over \(2.1*10^6\) parameters with the default configuration. When \(d_{ff} = 4*d_{model}\), this is equivalent to \(8*d_{model}*d_{model}\). The feed forward layer has roughly twice the number of parameters.

The attention layer computation scales quadratically with input length (default value is \(n=1024\)). The computational complexity of the feed forward layers is \(n*d_{model}*d_{ff}\). For attention, it is \(n^2d_{model}+d_{model}*d_{model}\). A recent trend is increasing the sequence length \(n\), which causes the attention layer to further dominate the computational cost. Due to the complexity of the attention operation and the different ways to implement it on hardware, we will skip calculating numerical values of the computation.

The feed forward layers contain the bulk of the transformer’s parameters, while the attention layers have the bulk of the computation. Attention is meant to learn the relationships between tokens, while the feed forward layers are meant to learn the individual token representations themselves. The attention operation is computationally intense in modeling the relationships between tokens, but it does not process individual token embeddings as much. The feed forward layers complement attention by enabling complex transformations of these embeddings.

This paper by Geva et al. argues that the feed forward layers act as key value memories. The high parameter counts of these layers enable the model to store rich information about the data they are trained on. Models like GPT-4 may not have their impressive world knowledge without the storage capacity of the feed forward layer. The transformer is a powerful architecture due to its balance of computational complexity and high parameterization.

Blocks

Now that we have covered each component, we can describe the transformer blocks. There are three main types of transformer blocks: encoder, decoder with cross-attention, and decoder without cross-attention. These blocks can be repeated any number of times.

Encoder Block

The encoder block maps a set of embeddings to another set of embeddings. It uses full self-attention, so each token can attend to all other tokens.

Decoder Block

The decoder block takes in a set of input embeddings but also attends to a set of embeddings from the encoder. The first attention layer processes input embeddings with causal attention. The second attention layer is cross-attention, where the keys and values come from the encoder output. This kind of block is only used in encoder-decoder architectures since it relies on the encoder output embeddings.

The cross-attention block is omitted in decoder-only transformers. This is because there are no encoder tokens to attend to. This block is identical to the encoder block, but the attention is masked.

Encoder-Only, Decoder-Only, and Encoder-Decoder Architectures

The original transformer paper introduced an encoder-decoder architecture. Since then, encoder-only and decoder-only architectures have gained significant popularity for various use cases. You can even have “encoder-heavy” or “decoder-heavy” architectures where one part of the transformer has more layers than the other. Let’s explore the different types of transformer models and their applications.

Encoder-Decoder

The encoder-decoder architecture can be viewed as two interconnected transformers. An encoder, which is a stack of encoder blocks, maps a sequence of input embeddings to output embeddings. A stack of decoder blocks with cross-attention then processes these embeddings. The decoder blocks read the final output embeddings from the encoder in the cross-attention layer.

The encoder and decoder can process different types of data. For instance, in speech recognition, the encoder might encode audio while the decoder translates it into text.

This architecture employs cross-attention, whereas encoder-only and decoder-only architectures rely solely on self-attention.

Encoder-Only

Encoder-only transformers, popularized by BERT, perform a one-to-one mapping of input embeddings to output embeddings. They can’t perform sequence-to-sequence modeling unless the input and output sequences have identical lengths.

These models excel at scalar prediction tasks, such as classification or regression, where the output is a single value rather than a set or sequence. Text sentiment analysis and ImageNet classification are prime examples of their application.

Encoder-only models are useful in tasks reducible to token classification. For instance, BERT’s evaluation on the Stanford Question Answering Dataset (SQuAD) doesn’t generate text answers but identifies the relevant span in the input text. The task becomes classifying which tokens mark the start and end of the answer span. Similarly, Vision Transformer (ViT), another encoder-only architecture, is trained for ImageNet classification.

Decoder-Only

Decoder-only transformers have become the go-to architecture for Large Language Models (LLMs), popularized by OpenAI’s GPT models.

A decoder-only model can tackle any task an encoder-decoder can handle. Instead of processing source data through a separate encoder, all data flows through the decoder. The decoder omits the cross-attention layer since there’s no encoder to attend to. Its architecture mirrors that of the encoder-only transformer, with the key difference being causal attention.

The encoder-decoder architecture can be viewed as a constrained version of the decoder-only architecture. The separate encoding of source data in encoder-decoder models represents a form of inductive bias that decoder-only architectures generalize away from.

Encoder-decoder models require paired source and target text sequences for training, which can limit their flexibility. In contrast, decoder-only models can process a single input sequence, making them more versatile and adaptable to various tasks.

Decoder-only architectures have surpassed encoder-decoder models in popularity due to their simplicity and versatility. However, encoder-decoder models still offer unique advantages, such as the ability to train on encoder-specific objectives or fine-tune the encoder for downstream tasks.

Output Processing

The architecture of the output processing is simple. There is a final linear layer that maps embeddings of \(d_{model}\) to the size of the prediction. The output of the linear layer and how it is applied depends on the task the model is trained on.

Transformer models can be trained with different objectives and losses based on the use case and architecture type. We will discuss different types of objectives and how they are used for training. We will also explain how inference works under these objectives.

Next Token Prediction

Encoder-Decoder and Decoder-only transformers are primarily trained on next token prediction. This task involves predicting the subsequent word in a sequence based on the preceding words. The output layer is applied to each embedding, with the output size matching the input vocabulary size. A softmax function then creates a probability distribution over the token vocabulary, from which the next token is sampled.

During training, the decoder learns to predict the next word given the context. Causal masking ensures that for each token, the model can’t see the next or subsequent tokens in the sequence.

For each token, the ground truth preceding tokens are used as context. This is known as teacher forcing, as the generated tokens aren’t used as context. However, at inference time, autoregressive decoding is used instead. We start with a context which is the set of input tokens. For chatbots, this might be the user’s input; for translation, the source text. The model predicts the next token, which is then added to the context, and another token is sampled. This process iterates until a special token is produced or a predetermined limit is reached.

At each step, the model outputs a probability distribution for the next token. This allows for sampling random sequences with a tunable temperature parameter, a common input for LLM APIs. Top-k sampling is another technique where you only consider the top-k tokens after the softmax. The probabilities are renormalized before sampling again. This prevents low probability tokens from ever getting predicted. Top-p or nucleus sampling is similar but possible tokens are selected so that their probabilities do not exceed p. This is more robust to changing entropies / confidence in the model’s predictions. These parameters are tuned to strike a balance between creativity and quality.

The paper Attention Is All You Need and many other NLP applications use beam search for text generation. Beam search is a search algorithm widely used in text decoding. Wu et al. 2016 provides insights into how it is used for transformer decoding. With beam search, instead of greedily selecting the most probable next token, beam search maintains the top \(k\) most likely sequences, where \(k\) is the beam width. At each step, the model expands these \(k\) sequences with their most probable next tokens. To prevent the number of sequences from growing exponentially (\(k^n\)), only the top \(k\) paths are retained after each step. The path probability is calculated as the sum of negative log likelihoods, normalized by dividing by \(length^{\alpha}\). This method allows the model to explore multiple promising paths simultaneously, avoiding local optima that might occur from selecting single tokens. To manage computational costs and control output length, beam search requires constraints. The beam size determines the number of computed paths and token selections per step, and the length penalty \(\alpha\) exponentially penalizes longer sequences. Attention Is All You Need implements this with a beam size of 4 and an \(\alpha\) value of 0.6.

Next token prediction training is highly parallelizable due to teacher forcing. A single forward pass generates predictions and losses for each token in the input. However, inference is an iterative process requiring a forward pass for each generated token.

Masked Language Modeling

Masked Language Modeling (MLM) is a training objective introduced by BERT. This method trains encoder-only transformers. MLM training follows these steps:

1. Randomly select 15% of the input tokens for potential masking.
2. Of these selected tokens:
   • • 80% are replaced with a special [MASK] token
   • • 10% are replaced with a random token
   • • 10% are left unchanged
3. Apply the linear output layer to all masked predictions. Use cross-entropy loss to predict the correct token, regardless of how it was masked.

The intuition behind this technique differs fundamentally from next token prediction in that it’s bidirectional. As it’s an encoder-only architecture with full self-attention, tokens to the left and right are used to predict the masked tokens accurately.

MLM is a pretraining method that doesn’t yield directly interpretable output. The model can be used for embedding representations, where the output embedding is aggregated and used in another application. The model can also be fine-tuned on a scalar prediction task.

Scalar Predictions

Encoder-only transformers support a wide variety of losses in addition to MLM. While MLM is a per-token objective where outputs are generated from multiple tokens, many objectives require using an output linear layer on a singular embedding to get a single output.

There are multiple ways to achieve this:

• • Special output token
  • BERT and ViT add a token to the input. The output embedding from this input is used for predictions.
• • Output pooling
  • Alternatively, you can take the average of all the embeddings and apply the output layer on this pooled embedding.
• • Attentive probing
  • Attention can process the output. A learnable query vector attends to all token embeddings, producing a weighted sum that is then used for the output layer. This is essentially a cross-attention with a fixed number of query embeddings.

Once you have a singular output embedding, it can be processed by the output linear layer, and then any loss function relevant to the objective.

Conclusion

This blog post covered the basic components of early transformers. In part 2, we will cover more recent innovations that further optimize these models and enable new capabilities.

Many of the current state-of-the-art results come from contrastive self-supervised learning. In contrastive SSL, an image from the dataset is transformed using multiple data augmentations, including cropping, color distortion, and flipping. These augmentations are inputted into a neural network to get representations. A contrastive loss is applied to push representations from the same image (different augmentations) closer together, and representations from different images are pushed further apart. Training with this objective will learn image representations that encode the content of the image. The data augmentations are needed for the network to avoid trivial solutions to optimize the contrastive loss, which would be comparing pixel values instead of the image’s semantic content.

Contrastive learning is limited by its dependence on data augmentations. These augmentations are hacky and unnatural. They can change the meaning of the image. The act of augmenting an image destroys some amount of information.

Example from Purushwalkam et al showing that cropping can remove semantic information in images.

Many AI researchers believe that humans learn through self-supervised learning. However, humans most likely don’t learn with data augmentations. It is unlikely we mentally do crops and color distortions. It is more likely we learn by tracking objects through time. For example, if you are watching a dog, the dog looks different from one point in time to another. The dog is likely in a different pose, different location, and has different lighting conditions. Through time, we can get natural data augmentations, but can they help train better image representation models?

The blog post explores whether video can improve self supervised computer vision models. We look at some papers with different approaches to the problem.

Image vs Video Representations

Image and video representations are two distinct research problems. Image representation learning aims to learn fixed-size embeddings of images, while video representation learning aims to learn fixed-size embeddings of videos. There are many papers on applying contrastive learning techniques to videos; however, these involve applying data augmentations to videos. In this post, we will also discuss using videos to learn image representations so that we can utilize the concept of obtaining natural data augmentations from videos.

Learning Video Representations

Video representation learning can be understood through its downstream tasks. One of the most common datasets is Kinetics. Which consists of short clips of human actions which are to be classified into classes such as “saluting” and “cutting cake”.

Spatiotemporal Contrastive Video Representation Learning

This work is very similar to SimCLR. It uses contrastive learning to learn representations of videos. Like SimCLR, InfoNCE loss is used to bring closer the representations of positive pairs and repel those of negative pairs. Positive pairs are data augmentations of the same video. There are two main differences with SimCLR:

1) 3D CNN model architecture for the video encoder to produce representations

3D-ResNets are used instead of 2D-ResNets.

This is a natural extension of CNNs to handle video that treats time as a third dimension. This architecture obtains a fixed-size representation from a fixed number of frames of a certain resolution.

2) Spatiotemporal data augmentations

In addition to the standard image data augmentations used in other contrastive SSL methods (color jittering, cropping, etc.), the authors add spatial and temporal augmentations. These are designed specifically for videos. It is important to carefully design the data augmentations used for videos. If the augmentations are too aggressive, the representations will be invariant to useful information. If the augmentations are too weak, the representations may learn trivial solutions.

Temporal Augmentations

The temporal interval is sampled first. This interval is the time difference between the two augmentations. Smaller intervals have more probability, which is desired because temporally distant augmentations might be too far apart. We want the content of the two augmentations to be the same. The further apart in the video they are, the less likely this is.

Spatial Augmentations

Applying spatial augmentations to each frame independently has the disadvantage of destroying motion between frames. If each frame is randomly cropped, the objects will randomly move around frame to frame. We want to preserve the consistency between frames. The solution to this is to sample a frame-level augmentation once per video clip and apply the same transformation to each frame. If each frame is cropped to exactly the same pixels, the motion will be perfectly preserved. For contrastive learning, we use data augmentations to make negative examples more different. We want two video clips to be different, but there is no need to make the frames of a single clip different from each other.

This paper is a great example of how image SSL techniques can be extended to videos. It just requires different model architectures and data sampling / augmentation techniques. There are some gaps to this approach:

1. It can only produce representations for fixed-size video clips. The resulting model can’t be used for downstream image tasks.
2. Temporal data augmentations can still destroy some temporal information in the video, especially fine-grained details.
3. Spatial data augmentations such as cropping are required and can be used to learn unwanted spatial invariances.

Learning Image Representations from Video

Self-Supervised Learning of Video-Induced Visual Invariances (VIVI)

They developed a video-based self-supervised learning framework for image representations and evaluated it on the VTAB image representation benchmark.

This uses the YouTube 8M dataset. This is a larger dataset than Kinetics and contains longer and more complex videos. Videos are composed of frames. In this work, they also utilize “shots,” which are sequences of continuous frames within a video. Shots have a high-level relationship with each other. Frames within the same shot are used as positive pairs, and shots from different shots are used as negative pairs.

Shot embeddings are defined as pooled (mean pooled or attention pooled) frame embeddings. These are used for shot order prediction. An LSTM or MLP is used to predict the next shot embedding given the current shot embedding. This models the relationship between shots in a video. An alternative to shot embedding prediction, would be to contrast shot embeddings between videos.

This paper uses co-training with the ImageNet supervised classification task. There is a long way to go since YouTube-8M is much larger than ImageNet, so ImageNet should not be needed for good results. This gap may be due to the data itself. ImageNet is a relatively clean dataset where the object in each image is usually centered. Video frames from YouTube are much noisier.

Demystifying Contrastive Self-Supervised Learning: Invariances, Augmentations and Dataset Biases

In this paper, the authors observe that the cropping required for contrastive SSL hurts performance on downstream tasks such as object detection. They make a distinction between scene-centric and object-centric datasets. ImageNet is object-centric, which means that a single object is presented in the image and it is centered. The data in video often contains multiple objects in different parts of the frame. For certain tasks, ImageNet is a far superior training dataset due to this bias.

They pretrain and evaluate an SSL model (MOCOv2) with the MSCOCO dataset (scene-centric) and MSCOCO bounding box cropped dataset (object-centric). The results show the effect of cropping when pretraining. They find that object-centric pretraining leads to better object-centric evaluation, while scene-centric pretraining leads to better scene-centric evaluation.

This shows that for scene-centric evaluation tasks like object detection, cropping while pretraining is harmful.

To replace cropping, videos can be used to get “temporal transformations”.

Pairs are formed by frames consecutive frames subsampled by $k$. We do not want frames that are too similar too each other in time, as their difference would be too insignificant. $k$ can be set depending on the frame rate of the video. If it is equal to 60, every 60th frame is considered and adjacent frames are pairs.

The dataset is then formed by applying transformations to the pairs of frames. For this work, the same transformations as MOCO are used. MOCO doesn’t require negative examples to train. Positive example pairs are formed by applying a transformation on the pairs of frames. Although this work explores using video for natural transformations, it still relies on the traditional data augmentations used in SSL.

This contrastive learning setup uses the whole frame, but the authors want to train on object-centric data to make the representations more robust for object recognition.

They use region tracking to make the frames of the video object-centric. They track the same object across multiple frames and use these cropped versions of the frames. This way, the model learns how objects change in a video while ignoring the scene-wide changes. For example, the position of an object in a scene is ignored.

VideoMAE

This work extends the MAE work (SSL with Vision Transformers blog post) to video. Rather than using contrastive learning, masked patches are predicted in the video.

The video is represented as a 3D image. The video is split into cubes (in this paper, it’s size 2 × 16 × 16). These cubes are treated like patches in ViT, but with an added time dimension. Each cube in the video is linearly projected and treated as a token embedding in the transformer. A position embedding is also added to the cube embedding. These cubes are treated as tokens in the transformer.

For each video, a tube mask is applied. This means the same patch in multiple consecutive frames of the video is masked. This is meant to make the SSL task harder, as patches don’t change much frame to frame. There is a high temporal correlation that needs to be broken in designing the objective. A high masking ratio, along with the tube masking strategy, ensures the SSL task is difficult and forces the model to learn higher-level spatiotemporal information.

An encoder processes the unmasked tokens. With a high masking ratio, we can save on computation cost by only processing the unmasked tokens. This is feasible with the transformer architecture, but not with CNNs. The encoder maps the token embeddings to a latent space. The decoder then predicts the masked cubes of the video. This can be trained with a simple reconstruction loss like MSE (mean squared error).

One interesting architecture decision is that “joint space-time attention” is just full self-attention. This means the attention captures all pairwise interactions between all tokens. It would be interesting to introduce causal attention on the time dimension. This would mean that within a frame, there is full attention. But cubes can only attend to cubes in the future. However, this type of causal masking would likely require a lower cube masking ratio to be effective.

Many other video SSL methods utilize image data in addition to video data. However, VideoMAE is able to achieve SOTA results on Kinetics-400 (video classification) without this external data. Ideally, we want video SSL methods that do not need to rely on images at all. This paper does not report results on image evaluation tasks. But this architecture would support this. We want video-pretrained models to achieve superior performance to image models on both video and image evaluation tasks. However, current video pretraining methods lag behind image-pretrained methods.

This work leverages the flexibility of the transformer architecture to directly predict elements of the video. This is simpler in that it does not require tracking objects, constructing triplets, or applying data augmentations. It is closer to how language models are trained. This also allows producing image representations and video representations, unlike the CNN methods. This is because the number of input frames to the transformer is variable.

Conclusion

Video lags behind images in representation learning for several reasons. While videos contain more information than static images, this information is spread across a much larger volume of data, making it computationally inefficient to process. The first frame of a video provides substantial information. You can learn what objects are present in the scene and the setting of the video. Subsequent frames offer diminishing returns. Future frames might just have one object moving across the frame. Because of this temporal redundancy, images are more information-dense than videos. Language models are more advanced than image models because language is a significantly more information-dense modality. Videos are an additional order of magnitude less information-dense than images. Data quality also plays a role. Video datasets, though vast, often lack the curated quality of image datasets like ImageNet, impacting the performance of video-based models. These challenges make it harder to build effective models with only video.

Looking ahead, progress in video-based AI will be driven by improved datasets, advancements in computational power, and novel modeling techniques. As these developments unfold, we can expect more sophisticated vision models that effectively incorporate video data, bridging the current gap between image and video understanding in AI systems. This may unlock new capabilities in computer vision models.

Datasets

ImageNet: 1,000,000 images, 1000 classes

Kinetics 700: 650,000 videos, 700 classes, ~10 seconds each, 1800 hours

YouTube 8M: ∼8 million videos, 500K hours of video—annotated with a vocabulary of 4800 entities

One key variation is that you often see masking in these methods. CNN-based SSL methods rely more on data augmentations to create a prediction task for the model. Masking is advantageous for several reasons outlined below, and it also aligns more with language model training (example: BERT).

• Computational efficiency
  • You do not have to process the masked regions of the image when a large portion of the image is masked.
• Data augmentations can introduce unwanted invariances and remove useful information
  • For example, a data augmentation that strongly distorts the color may result in representations that do not encode color.

Masking is more naturally enabled by the transformer architecture. There is a reason that masking-based SSL training hasn’t worked well with CNNs.

By examining these different methods, we’ll discuss what makes transformers work for vision.

DINO

This paper (Emerging Properties in Self-Supervised Vision Transformers) by Caron et al. introduces a new self-supervised training method called DINO, which they apply to vision transformers. They argue that transformers are better than CNNs for images with SSL training, more so than with supervised training. Transformers can match the performance of CNNs with supervised training, albeit with more training cost. However, they have more useful properties with SSL training. This follows our intuition that SSL and transformers are a natural combination.

DINO takes inspiration from BYOL but introduces two key innovations:

1. A novel loss function that enables direct matching between student and teacher outputs
2. Elimination of the prediction layer on the student, simplifying the architecture

These changes result in a self-distillation approach that proves particularly effective with vision transformers.

1. Two views of an image \(x\), \(x_1\) and \(x_2\) are generated through data augmentations.
   1. A multi crop strategy is used in which two large global views are generated along with a set of smaller cropped local views. The teacher only processes global views, while the student processes all views, with the constraint that the loss is not trying to match the same views to each other. This method was introduced in the SwAV paper, and helps the model learn local to global correspondences. Restricting the teacher to only global views also encourages the encoders to output global representations.
   2. Are position embeddings used?
2. The views are passed to their respective encoder (teacher/student)
3. The teacher encoding is “centered”.
   1. Perhaps centering allows this method to work without having the predictor layer. The center is a exponential moving average of the teacher encoding (of both views). This vector is subtracted from the teacher’s encoding before the softmax. A temperature is also applied with the softmax to achieve a “sharpening”. These methods help the teacher avoid collapse. Centering ensures that a single component of the vector doesn’t dominate. Sharpening ensures that it doesn’t collapse to a uniform vector.
4. Softmax is applied to each encoding. The student is trained with a cross entropy loss to match the teacher. The teachers weights are updated as an exponential moving average of the student.

This paper compares the performance of DINO with ResNet and ViT architectures against SOTA SSL methods such as BYOL, MoCov2, and SwAV. The combination of DINO and ViT has the most significant advantage. Interestingly, it is 6.6% better than ViT with BYOL training on linear ImageNet evaluation, despite minor differences in the methods. The SSL methods that are used for comparison were developed for CNN architectures, which put them at a disadvantage. DINO is designed for transformers, but what about it makes it work better with transformers? One possible explanation is that transformers handle different resolutions of images better. Higher resolution images results in more image patches generated in the transformer. The computation also scales quadratically in the attention operations with respect to the number of patches. For ResNet, the computation increases linearly.

The two main “emerging properties” they observe is that DINO ViT features are useful for dense predictions such as semantic segmentation. Another property is that k nearest neighbors on the output encodings, without any finetuning. This enables image retrieval applications.

They observe the teacher outperforms the student in DINO training. This is not observed with other SSL methods. They cite “Polyak-Ruppert averaging” as an explantation of this. This means the teacher simulates an ensemble model with its momentum weights.

The multi-crop strategy enforces that the inputs be rectangular. This makes this method compatible with CNNs in addition to ViTs. DINO shows that SSL is effective with vision transformers. However, it is designed in a way that makes the training method compatible with CNNs. This leads to some very interesting comparisons between the properties of SSL CNN and ViT models. The other works we will discuss take advantage of the flexibility of the transformer architecture, at the cost of CNN compatibility.

DINOv2: Learning Robust Visual Features without Supervision scales DINO using a 1 billion parameter ViT model along with a larger proprietary dataset. They used an interesting data processing pipeline to combine curated and uncurated data, to get a large dataset of high quality and diverse images. This step is important because unprocessed uncurated data can be of low quality and dominated by certain modes of data and duplicated data.

There are several architectural and training changes applied on top DINO v1 that allow it to scale effectively. Notably, in addition to DINO, they add an iBOT loss. This method masks some of the input tokens of the student. In order to combine DINO and iBOT losses, they learn separate heads on the student and teacher for each loss. iBOT does BERT style pretraining of image transformers, which we will also cover in this post.

data2vec

The teacher model predicts representations from unmasked input, while the student model predicts representations from masked input. The student aims to match the teacher’s output by predicting the representations of the masked tokens. To avoid collapse, the teacher’s weights are an exponential moving average of the student’s weights.

Instead of training a multimodal model, independent models are trained on different modalities. Data2VecAudio, Data2VecText, and Data2VecVision are developed. The learning objective remains the same, but the generation of embeddings and masking strategies differ.

1. Encoding of inputs into embeddings:
   1. Text is tokenized, and learned embeddings for each token are retrieved.
   2. Images are divided into 16x16 patches and linearly projected into an embedding.
   3. Audio is encoded by a 1D convolutional neural network with multiple layers. A 16 kHz waveform is mapped to a 50 Hz representation. This means a sequence of 320 integers is mapped to a single representation.
      1. Unlike images, a multiple-layer network is used for audio, likely due to the absence of a Fourier transform.
2. Masking:
   1. Some of the student input embeddings are replaced by the MASK token embedding.
      1. Text: Random tokens are masked.
      2. Images: Embeddings corresponding to rectangular blocks are masked.
      3. Audio: Continuous spans of embeddings are masked.
3. Addition of position encoding.
4. Both the teacher and student transformer models receive the input.
5. Representations at different layers are distilled from the teacher to the student. Outputs from the masked tokens of the top \(K\) transformer blocks are normalized and averaged into a single vector.
6. A regression loss (Smooth L1) is applied to the averaged vectors of each network.
   1. The loss transitions from a squared loss to an L2 loss when the error margin goes below the hyperparameter \(\beta\). The L2 loss is only applied when the student and teacher predictions are close. This loss is designed to be less sensitive to outliers.

1. The students weights are updated with SGD. The teacher’s weights are updated as a EMA of the students weights: \(\Delta \leftarrow \tau \Delta + (1-\tau)\theta\)
   1. \(\Delta\) represents the teacher’s parameters, while \(\theta\) represents the student’s parameters.

The position encoding and feature encoder weights are shared between the two models. However, the teacher’s transformer weights are specified through an exponential moving average.

data2vec 2.0

Data2Vec 2.0 introduces several architectural and loss function changes that lead to a significant speed up in training.

They use target representations for multiple masked predictions of a sample. This is more computationally efficient because we only need to run the teacher model once to train with \(M\) different masks of the input instead of 1. Further efficiency gains are implemented through not processing the masked parts of the image with the student, and sharing the feature encoder output across all masks.

They use a L2 loss instead of a smooth L1 loss. This is a simplification of the earlier loss. They also use a convolutional decoder to predict the masked representations rather than a transformer.

They also introduce inverse block masking. Rather than masking blocks. Blocks are chosen to be unmasked areas. The representations outside of the block will be predicted. There are multiple blocks which may overlap. A mask consists of multiple blocks. Training includes multiple masks for each target.

They also add a linear attention bias (ALiBi). This essentially modifies self attention to increase the bias for query key pairs that are far apart. This enables faster training by providing an inductive bias.

Masked Autoencoders Are Scalable Vision Learners

This paper uses a simple autoencoder architecture to learn image representations. Parts of the images are masked, and the model is tasked to predict what is in the masked regions. This model can be trained through this notebook.

1. The image is split into patches, as done in Vision Transformers.
2. Using a mask ratio (75%–95%), patches are selected randomly without replacement.
3. The unmasked patches are input into the encoder. Note that the mask tokens do not get processed by the encoder (difference from BERT). The encoder uses a vanilla ViT architecture, where the unmasked patches are linearly projected into token embeddings which get processed by transformer blocks. The output is a ViT-processed embedding for each unmasked patch. Each patch has an added position embedding.
4. The encoded tokens and the masked tokens are combined as an input to the decoder. The mask tokens map to a learned embedding. This embedding will be the same at all positions because it is not transformed by the encoder. At this stage, position embeddings are added to the full set.
   1. Note that for unmasked tokens, position embeddings are added twice, once before the encoder and once before the decoder.
5. The decoder reconstructs the unmasked image from the set of patch embeddings. The decoder is trained by a mean squared error loss with respect to the unmasked input image.

This architecture builds on the vision transformer. An alternative is to use CNNs. This would involve directly setting pixels in the input image to zero, learning a vector representation, and then decoding it back to the image. The reason this fails is that it aims to globally decode an image. With transformers, you first predict representations of the masked patches, and then decode into the image patch. This breaks it down into two easier problems. Also, with CNNs, you can’t explicitly encode masked regions like you can with a ViT. Having a mask token more explicitly indicates the mask.

They mask a very high percentage of patches (80%). This reduces spatial redundancy and forces the model to learn more higher-level and useful features. With a lower mask ratio, the model might learn to represent small local changes, like color and lighting variation. It doesn’t need to understand the higher-level structure of the image, because it’s mostly already there. This is a notable change from language models. BERT masks 15% of tokens. MAE and related works mask a majority of the image (75%+).

The model uses the ImageNet-1K dataset for pretraining and evaluation. Evaluation is done by either finetuning the full encoder model or using a linear probe (training one MLP layer on the output of the encoder) on the task of classification.

One interesting result is that the performance of finetuning and linear probing has different trends when ablating the masking ratio. Linear probing accuracy increases linearly with masking ratio until 75%. Finetuning has relatively consistent performance between 40% and 80%.

Having a deep decoder allows for the representations to be more abstract, because the decoder has more capacity for reconstruction. A shallower decoder would lead to the encoder having to represent more of the details needed for reconstruction. This is less relevant for finetuning than it is for linear probing, as during finetuning, the encoder can shift from focusing on reconstruction to recognition. In my opinion, linear probing results are more interesting since the goal is to build useful representations that can be used for various tasks. Finetuning offers just a marginal improvement over just training on the classification task directly without pretraining at all. However, linear probing discourages learning nonlinear features in the representation. To address this, the authors evaluate “partial finetuning” in which the last few blocks of the transformer are finetuned.

Excluding mask tokens from the input and using a lightweight decoder makes this model very efficient to train. Using mask tokens in the encoder also creates a domain shift between pretraining and downstream tasks, which hurts performance. This is because a large portion of the pretraining input will be mask tokens, which is significantly different from what the model will see downstream.

BEiT: BERT Pre-Training of Image Transformers

This approach is most similar to BERT / NLP SSL models.

A fundamental difference in applying SSL to images compared to text is that images are continuous. Text has a finite number of tokens. You can use a softmax to get a probability distribution across all tokens. In ViTs, patches of an image are treated as tokens. However, you can’t get an explicit probability distribution over all possible image patches. BEiT addresses this problem by training a discrete variational autoencoder (dVAE) to learn discrete visual tokens. These discrete tokens are an approximation or compression of image patches.

The main difference between this and a vanilla ViT architecture is the usage of discrete visual tokens.

There are two steps to training:

1. Tokenizer and Decoder are trained as a VAE to learn discrete visual tokens
2. The discrete tokens from the learned tokenizer are used to pretrain a BEiT encoder.

Why aren’t the tokens used as the input directly? The softmax distribution of tokens could be used as a soft label for the BEiT encoder.

The transformer training task is named masked image modeling (MIM), as it is designed after BERT’s masked language modeling (MLM). 40% of the tokens are masked. Similar to other methods, BEiT masks a large portion of the image to make the pretraining task sufficiently difficult.

Conclusion

The landscape of self-supervised learning for image processing is undergoing a significant transformation. While it originated with Convolutional Neural Networks (CNNs), a strong coupling with transformer-based architectures is emerging and may lead the way for further advancements.

In the AI research community, Yann LeCun has a unique and often controversial perspective. As of 2024, LLMs and Generative AI are the main focus areas of the field of AI. We’ve all been impressed by the performance of LLMs in various contexts, and generative systems like OpenAI’s Sora. However, it is not clear where these advances fit in the long term goal of achieving and surpassing human level intelligence, which many call AGI.

In his position paper A Path Towards Autonomous Machine Intelligence and his many recent talks (linked below), Yann presents an alternative framework for achieving artificial intelligence. He also proposes a new architecture for a predictive world model: Joint Embedding Predictive Architecture (JEPA).

This blog post will dive deep into Yann’s vision for AI, the JEPA architecture, current research, and energy-based models. We will go deep into the technical aspects of these ideas, as well as give my opinions, along with interesting references. I will also cover recent research advances such as V-JEPA

This is a long post, feel free to jump to the sections about JEPA, I-JEPA, and V-JEPA.

Relevant Talks by Yann LeCun

From Machine Learning to Autonomous Intelligence

Objective-Driven AI: Towards Machines that can Learn, Reason, and Plan

Problems with Current AI

The JEPA architecture aims to address current AI challenges. To contextualize these issues, we’ll examine Yann LeCun’s criticisms of popular AI trends as of 2024.

Recent years have seen tremendous excitement around Large Language Models (LLMs) and Generative AI. LLMs are pretrained using autoregressive self-supervised learning, predicting the next token given preceding ones. They’re trained on vast datasets of text and code from the internet and books, often fine-tuned with supervised learning or reinforcement learning. Generative AI broadly refers to creation of multimodal media from inputs, such as text-to-image generation.

However, these models face significant limitations:

1. Factuality / Hallucinations: When uncertain, models often generate plausible-sounding but false information. They’re optimized for probabilistic likelihood, not factual accuracy.
2. Limited Reasoning: While techniques like Chain of Thought prompting improve LLM’s ability to reason, they’re restricted to solving the selected type of problem and approaches to solving them without improving generalized reasoning abilities.
3. Lack of Planning: LLMs predict one step at a time, lacking effective long-term planning crucial for tasks requiring sustained goal-oriented behavior.

Despite impressive advancements, the challenge of autonomous driving illustrates the gap between current AI and human-level intelligence. As LeCun notes, humans can learn driving basics in about 20 hours. In contrast, self-driving car development has consumed billions of dollars, extensive data collection, and decades of effort, yet still hasn’t achieved human-level performance.

Even achieving Level 5 autonomy wouldn’t signify true human-level AI or Artificial General Intelligence (AGI). Such intelligence would involve learning to drive from scratch within a day, using only data collected during that experience, without relying on massive pre-existing datasets for finetuning. Realizing this level of adaptable intelligence might require several more decades of research.

Common Sense

The limitations in AI models can often be attributed to a lack of common sense. Common sense can be defined as thinking and acting in a reasonable manner. Humans and many animals have this ability. This includes avoiding egregiously dangerous or incorrect actions. Expanding on the autonomous driving example, AV systems need to be trained to deal with new situations safely. When learning to drive, humans utilize their common sense to know to not do dangerous things like driving off the road or into other cars. This is not obvious to current AV systems, so they require a large amount of training data to avoid these actions.

LLMs similarly demonstrate a lack of common sense through nonsensical or illogical outputs. Common sense is a vague term. One definition is that it is a lower bound on the types of errors an agent makes. For AI to be trustworthy, it needs this foundational level of understanding.

Common sense can also be viewed as a collection of world models. These models enable quick learning of new skills, avoidance of dangerous mistakes in novel situations, and prediction of outcomes in unfamiliar scenarios. Essentially, we use world models to generalize our experiences.

How Humans Learn

Humans acquire a basic understanding of the world during early infancy, but we’re also born with some innate knowledge. The brain isn’t randomly initialized; it’s evolved, pre-trained, and fine-tuned throughout life. This differs significantly from artificial neural networks, which start with random initializations and have far weaker inductive biases than humans or animals. Life is generally pre-programmed to behave in a certain way from birth. More intelligent life is able to learn more and not purely rely on innate knowledge.

Understanding the extent to which babies acquire common sense during infancy is crucial for AI development. If common sense is largely innate, the focus should be on massive datasets mimicking evolutionary timescales. If it’s primarily learned, priority should be given to models that excel at quick learning from limited data.

A baby’s experience, while not comparable to evolutionary timescales, still represents a substantial dataset. If a baby is awake for 8 hours a day, in four months they have seen about 960 hours of data. This data is also augmented by other sensory signals and dense biological supervision (pain, hunger, emotions). This is around the same length as the Kinetics 400 video dataset. This is still dwarfed by the millions of hours of video that self driving cars are using.

This Nature paper by Orhan and Lake explores learning from infant-perspective data. They demonstrate that computer vision models can be trained on noisy, less diverse datasets collected from infant headcams. These egocentric datasets are far noisier and less diverse than standard image/video datasets, but AI models without strong inductive biases can learn from them.

Emmanuel Dupoux’s diagram, presented by Yann LeCun, suggests that babies often understand concepts like object permanence, solidity, and biological motion by around four months. While presented as quick learning, it’s important to note the significant amount of data processing that occurs during this time.

We don’t yet know precisely how much data AI systems would need to learn the same concepts as babies. It’s likely that the data efficiency gap is relatively small for basic concepts that babies learn. For instance, object permanence could probably be learned from 960 hours of video data. However, it becomes evident that this gap grows substantially with age and with the complexity of the knowledge being assessed. The challenges in developing fully autonomous vehicles clearly demonstrate how large this data efficiency gap can become.

In addition to the lack of common sense, we mention three other fundamental gaps in the ability of current AI: hallucinations, lack of planning, and lack of reasoning.

Learning to Think

The question of whether Large Language Models (LLMs) can truly reason and plan is a contentious topic in the AI community. While these models exhibit behaviors that resemble reasoning and planning, skeptics argue that they merely replicate patterns from their training data.

To frame this discussion, let’s consider reasoning and planning as forms of “thinking”, which we will define as a variable length internal process that precedes any outputs.. Current deep learning models employ two primary mechanisms for this kind of processing:

1. Depth: Each layer in a neural network can be viewed as a step in the thinking process. However, this depth is typically fixed, with some recent work exploring dynamic depth adjustment based on input complexity. Despite these advances, maximum depth and other constraints still limit the model’s flexibility.
2. Sequential Generation: Decoder-based LLMs, such as GPT, generate text one token at a time. Each step in this process involves some degree of computation that could be interpreted as thinking. Prompt engineering techniques leverage this sequential nature to guide the model towards desired outputs. A key limitation of this approach is that the model must produce a token at each step, preventing purely internal information processing.

While these properties enable models to create the illusion of thought, significant advancements are necessary to achieve more effective reasoning and planning capabilities.

Many researchers draw parallels between AI and the two-system model of thinking proposed by Daniel Kahneman. System 1 thinking is fast and intuitive, providing immediate responses without conscious deliberation. System 2, in contrast, is slower and more deliberate, engaging in deeper cognitive processing. Current machine learning models, including LLMs, primarily operate in a System 1 mode by processing information in a single pass without the ability to plan ahead. While they excel at pattern recognition, they lack true reasoning or planning capabilities.

This inability to plan contributes to factual errors in LLM outputs. Each generated word carries a risk of inaccuracy, with the probability of errors increasing exponentially as the output length grows. The sequential nature of token generation means that early mistakes can compound, potentially invalidating the entire output. This stands in stark contrast to human speech, where we typically plan our utterances at a higher level before vocalization, minimizing such errors. In this context, reasoning can be viewed as the planning of speech. Without the capacity to reason or plan effectively, LLMs essentially “speak without thinking.”

In the JEPA paper, Yann LeCun proposes frameworks for models that can think. Learning to think may address the fundamental problems in current AI models and represent a crucial step towards achieving more human-like intelligence in AI.

Modality

Recent advancements have expanded LLMs to include multimodal processing and outputs, but they remain primarily language-centric. This raises questions about the sufficiency of language alone for AI and the investment needed in visual understanding. Could visual comprehension help ground AI in reality, improving common sense and reducing hallucinations?

Language serves as a compressed representation of the complex concepts humans experience. Its expressive power is vast, capable of describing intricate scientific theories and nuanced emotions. Yet, language alone may not suffice for complete understanding.

Humans interpret language within the context of shared reality. It functions as a highly efficient medium for transmitting information through the relatively narrow bandwidth of speech. When we process language, our brains rely on prior knowledge and experiences. While some of this prior information can be acquired through text, a significant portion stems from visual and physical interactions with the world.

Currently, it does seem that language models are more capable than vision models. Language models currently outperform visual models due to information density, data requirements, and data availability.

In a given data point there is a certain amount of explicit information in the form of bits. But then there is relevant information that is useful. For example, if you take an image of the park, a lot of bits are used to represent the position of every blade of grass. But that is not useful in most scenarios. Language is very compressed. While there are some filler words that don’t add much information, the ratio of knowledge to bits is high. However, for images, most of the bits are not useful. This means you need orders of magnitude more bits of data to learn equivalent knowledge. Video models are further behind because you need another order of magnitude more bits since consecutive frames in video are mostly redundant.

While language-based AI leads, scenarios exist where visual learning could catch up. One scenario in which visual learning could overtake language is that we will have a large number of robots / autonomous vehicles interacting with the world while collecting visual data. Language will be data constrained with the rate of new text generation limiting scaling. In a world with a lot of robots, the knowledge gained from the visual world and the size of the available datasets may exceed that of text. However, this is all very speculative. We don’t know how important vision or grounding is for intelligence.

A Framework for Building Human-Level AI

Yann proposes a high level architecture for building an AI system that is aimed at addressing the problems we outlined. This is a design for an intelligent agent that can perceive the world,

We will then explore the various challenges that must be addressed to construct such an architecture. Currently, this is merely a theoretical architecture. Building certain components remains an open problem, and assembling all the modules will pose an additional challenge.

This architecture contains different proposed components. We will explain these components and their relationships.

Configurator: Configures input from all other modules and configures them for the task at hand. It tells the perception module what information to extract.

Perception: Estimates the current state of the world from different sensory signals.

World module: Estimates missing information about the state of the world and predicts future states. It simulates the world and extracts relevant information as determined by the configurator.

Cost module: Measures the level of discomfort as energy. This energy is the sum of the intrinsic cost module and the trainable critic module.

Intrinsic cost: Computes a cost given the current state of the world and predicted future states. This cost can be imagined as hunger, pain, or general discomfort. This cost can be hardwired in AI agents, as done with rewards in RL.

Trainable Critic: Predicts future intrinsic energy. It has the same input as the intrinsic cost. This estimate is dependent on the intrinsic cost and cannot be hardwired. It is trained from past states and subsequent intrinsic cost, retrieved from memory.

Short term memory: Stores relevant information about past present and future states of the world along with intrinsic cost.

Actor: Proposes sequences of actions. These sequences are executed by the effectors. The world model predicts future states from the sequence which then generates a cost.

Actor

The actor proposes an optimal action or sequence of actions.

If the world model and cost are well behaved, gradient based optimization can be used to determine an optimal action sequence. If actions are discrete then dynamic programming methods such as beam search can be used.

There are two different modes in the actor. These align with Kahneman’s System 1 and 2, which we mentioned earlier.

Mode 1 Reactive Behavior: A policy module that computes an action from the state generated by perception and short-term memory. This module acts fast and produces simple decisions. A world model is needed to estimate the cost of an action. Without a world model the agent would have to perturb their actions which is not feasible. The world model can be adjusted after observing the next state.

Mode 2 Reasoning and Planning: A sequence of actions along with predicted corresponding states is generated. From this sequence of states, a cost can be computed. Planning is done by optimizing the action sequence to minimize total cost. The action sequence is then sent to the effectors which execute at least the beginning of the sequence. The states and costs are stored in short-term memory. The sequence can be optimized through gradients since the cost and world model are differentiable. Dynamic programming can also be used. Planning in this setup is essentially inference time cost optimization.

Agents may have multiple policy modules executing mode 1. In this design, the agent only has one world model, so mode 2 can only be run once. However, AIs could be designed to have multiple world models and mode 2 processes at the same time. This is similar to having multiple thoughts at the same time. However, this would be very complicated in that the different modules would have to coordinate with the effectors and other modules to avoid conflicts. Also, this may be why humans don’t think like this.

Policy modules can be learned to approximate actions from mode 2 reasoning. This is the process of learning a new skill. In humans, system 2 thinking can be done through system 1 after enough learning. For example, in chess, inexperienced players plan steps explicitly and simulate outcomes. Experienced players can instantly recognize patterns and make optimal moves.

Cost

Cost is the sum of an immutable intrinsic cost and a trainable cost or critic.

Each of these costs are the sum of different sub-costs generated by submodules. The weights of the sub-cost at each state \(u\) and \(v\) are determined by the configurator. This allows the agent to focus on different goals at different times.

The IC being immutable prevents the agent from drifting towards bad behaviors. It constrains the behavior of the agent.

\(\mathrm{TC}\) or the critic is trained to predict future intrinsic cost values. The intrinsic cost only considers the current state. The critic can be trained to predict the future cost so the agent can minimize cost in the future. The short term memory stores triplets of (time, state, intrinsic energy): \((\tau, s_{\tau}, IC(s_{\tau}))\). The critic can be trained to predict the cost of a future state or a discounted sum of future intrinsic costs. For example, the loss function of the critic could be \(\|\|\mathrm{IC}(s_{\tau+\delta}) - \mathrm{TC}(s_{\tau})\|\|^2\). This formulation trains the critic to predict the intrinsic cost of a state \(\delta\) steps in the future. \(\mathrm{IC}(s_{\tau+\delta})\) can be replaced with other targets that can be extracted from the sequence of triplets. However, it cannot depend on the future trainable cost itself.

Configurator

The configurator controls the other components of the system. If these components are implemented as transformers, they can be easily configured by adding tokens. The configurator would inject tokens to steer these components in certain directions. For example, it may influence certain types of actions from the actor, or for perception to focus on certain properties.

The configurator is also responsible for setting the weights of the cost terms. This will allow for the agent to focus on different subgoals at different times. The unanswered question is how the configurator can learn to decompose a complex task into subgoals.

World Model

In JEPA, the purpose of the world model is to predict future representations of the state of the world. There are three main issues

1. Diversity of the state sequences the model is able to observe during training
2. The world isn’t fully predictable, so the model has to predict multiple plausible state representations following an action
3. Predictions must be made at different time scales and abstractions

Self-Supervised Learning / Energy-Based Models

In order to train a world model, Yann LeCun proposes an SSL energy-based model (EBM).

\(x\) and \(y\) can be considered videos, where \(y\) follows x. EBMs learn an energy function \(F(x,y)\) that take low values when \(x\) and \(y\) are compatible and high if not. Compatible in this context means that \(y\) is a plausible continuation of \(x\).

This is different from generative models in that \(y\) is not directly predicted from \(x\). There is a large space of values of \(y\) that can follow \(x\). Predicting exactly what will happen is an intractable problem. However, it is feasible to understand what is possible and what is not. Being good at this task requires an understanding of the world and common sense. A value of \(y\) that defies the laws of physics should result in a high energy value.

However, planning requires predictions of future states. Although \(y\) can’t be predicted directly, we can predict future representations of \(y\). We can get representations from an encoder: \(s_x = g_x(x)\), \(s_y = g_y(y)\)

The encoder will be trained such that the representations are maximally informative about each other, and that \(s_y\) can easily be predicted from \(s_x\). We can make predictions on this representation to enable planning.

A latent variable can be introduced to handle uncertainty. A latent variable is just an arbitrary random variable. It is the source of randomness that is transformed to a useful distribution. Here we want to map the latent variable to the large space of possible values \(s_y\) can take.

A latent-variable EBM (LVEBM) is represented as \(E_w(x, y, z)\).

The energy function can be determined by find the \(z\) value that minimizes the energy. \(F_w(x,y) = \min_{z \in \mathcal{Z} }E_w(x,y,z)\)

The EBM collapses when all pairs have the same low energy. This can happen when the latent variable has too much information capacity. This happens because \(z\) can vary along a larger space. This means that the space for which the energy of \(y\) is low is correspondingly large. If it is too large then the energies of \(y\) collapse. If the \(z\) dimension is the same as the representation dimension, the model can ignore \(y\) entirely and set \(s_y\) to equal \(z\).

The paper describes a high data density region. This refers to \((x, y)\) pairs that are commonly seen in the real data distribution. We want to lower energy in this region, but keep it high outside of it. Collapse is when the energy is low inside and outside of this region which makes the EBM useless.

There are two training methods used to prevent collapse.

Contrastive methods: Collapse is avoided by increasing the energy with respect to negative examples. It requires some method to generate examples to contrast against. The number of contrastive examples needed grows exponentially with respect to the dimension of the representation.

Regularized methods: In these methods, the loss is regularized to minimize the space in \(y\) where the energies are lowered. These are less likely to be affected by the curse of dimensionality. Contrastive architectures can be regularized. For example, the latent dimension can be constrained.

Joint Embedding Predictive Architecture

JEPA is an EBM that performs predictions in the representation space. The energy is the error in predicting \(s_y\) from \(s_x\).

JEPA needs multi-modality, which in this context means to represent multiple possible values of \(y\). There are two ways it can be achieved.

Encoder invariance: This means that \(s_y\) will be the same for different values of \(y\). The encoder ignores aspects of the state that may vary.

Latent variable predictor: Varying \(z\) will lead to different plausible predictions of \(s_y\).

There are four criteria that can be used to train this architecture without contrastive loss:

1. Maximize the information content of \(s_x\) about \(x\): \(-I(s_x)\)
2. Maximize the information content of \(s_x\) about \(y\): \(-I(s_y)\)
3. Make \(s_y\) predictable from \(s_x\): \(D(s_y, \tilde{s_y})\)
4. Minimize the information content of the latent variable with a regularizer: \(R(z)\)

Hierarchical JEPA (H-JEPA)

There is a trade off between information loss in the encoding and the predictability of the encodings. If a representation contains most of the information of the input, it would be hard to predict. A more abstract and higher level representation would be lower in dimension and more predictable. Higher dimension representations are also more suitable for longer term predictions.

H-JEPA (Hierarchical JEPA) enhances JEPA’s abstraction capabilities by splitting the architecture into two parts. The first JEPA handles low-level representations for short-term predictions, while the second operates at a higher abstraction level for longer-term forecasts. This two-tier structure, though innovative, is arbitrary. True intelligence requires multiple levels of abstraction. However, it is not clear how many levels of abstraction are needed. We may even need variable levels of abstraction. Different situations have different levels of complexity.

This architecture can enable higher level planning. In JEPA-2, we can sample from the latent variable for several time steps. Directed search / pruning can be employed in order to efficiently search. This search can be used to determine an optimal action.

This kind of search would be different in JEPA-1 or without H-JEPA because the latent dimension would be too large to efficiently sample from. Abstraction is needed to enable this kind of planning.

World Model Architecture

The world is unpredictable but the agent itself is predictable to the agent. This may motivate a model of self (ego model) that does not have a latent variable.

The state of the world varies only slightly between time steps. Rather than regenerating, it can be updated in memory. With this architecture, the world model will only output the change in the state. This can be implemented with an attention-like mechanism.

1. The world model outputs query value pairs: \((q[i], v[i])\)
2. The world model retrieves a value from memory using the query
   • \[\mathrm{Mem}(q) = \sum_jc_jv_j\]
     • The value retrieved from memory is a weighted sum of all values.
   • \[\tilde{c}_j = \mathrm{Match}(k_j,q)\]
     • Measures dissimilarity between the key and query.
   • \[c = \mathrm{Normalize}(\tilde{c})\]
     • This is often a softmax.
   • \[v_j = \mathrm{Update}(r,v_j,c_j)\]
     • Value is updated using the current value and new value.
     • The update function can be \(cr+(1-c)v\)

Data Streams

In building a world model, we have to consider the fundamental differences in the type of data that humans and AI models process. Yann lists 5 modes of information gathering that an agent can use to learn its world model.

1. Passive observation: sensor stream without control
2. Action foveation: The agent can direct attention within the data stream
3. Passive agency: Observing another agent’s actions and causal effects
4. Active Egomotion: The sensors can be configured, for example moving a camera
5. Active Agency: Sensory streams that are influenced by the agent’s actions

Current AI methods largely focus on passive observation. Other modes may be needed to reach intelligence.

AI is trained on internet data. Internet data is not experienced by the agent. Humans train on data that they experience. This is a fundamental difference. This is also why autonomous cars need so much training data. The AI driving systems don’t have other datasets that they have experienced. For example, if they trained on a large dataset of just walking around, they would need less driving data.

It is challenging to create large-scale datasets from the perspective of an agent, especially reaching the scale of internet datasets. A present-day example is autonomous car datasets. AV companies have large fleets of vehicles on the road collecting data. These are active data streams.

Objective Driven AI

The components of this architecture can be put together to build an intelligent system that follows human defined objectives.

Perception is used to generate an initial representation of the state of the world. The actor proposes a sequence of actions. The world model then predicts the state reached if the action sequence is executed. This state is then used in the objectives. The task objective defines what we want the system to do. This could be a task or particular problem. The guardrail objective makes sure the system accomplishes the task without any unwanted behavior. These guardrails would be designed for safety.

The action sequence is optimized with respect to the objects. There will be a lot of flexibility in designing the objects to get the system to behave in the way we want.

The system can also be extended to achieve hierarchal planning. The higher levels of planning produce a state that will serve as an objective for the lower level. This state can be considered as a subgoal that is necessary to achieve the higher level goal. We can have unique objectives and guardrails for each level of planning.

Latent variables are also introduced to represent the uncertainty in predictions of future states. The latent variables at the higher levels can be thought as imaginary higher level actions. However, only the lower level actions can actually be directly executed.

Towards Implementing JEPA

The JEPA paper is a position paper that describes a vision for AI that may take decades to materialize. However, since its publication in the summer of 2022, there have been a few steps in advancing the architecture. These papers particularly explore the training of JEPAs. They do not explore the other components such as planning. These JEPAs are the first steps to creating a world model.

These are essentially self supervised pretraining methods. When comparing against other works, these papers cite training speed as their advantage. They can achieve strong downstream performance with fewer pretraining epochs.

I-JEPA: Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture

Compared to other image SSL approaches, I-JEPA takes advantage of the flexibility of the transformer architecture. ViT is used because it can handle an arbitrary amount of patches in an image, without requiring a strict shape in the input like CNNs

The input image is split into \(N\) non-overlapping patches and fed into a target encoder \(f_{\theta}\) to compute patch representations. \(s_y = \{s_{y1} … s_{yN}\}\)

\(M\) possibly overlapping blocks are sampled from these representations. These blocks are basically larger sections of the image that contain multiple patches.

Context is generated by sampling a block (larger than the target blocks). When predicting a target from this context, the overlap with the target block is masked from the context. The network is trained to predict the representations of the target blocks given the context block, and position encodings for the target block. The position encodings are added to the input so that the model knows where the target is. It is just tasked with predicting representations at those positions.

This architecture avoids collapse by having exponential moving average weights in the target encoder. This is the same approach used in data2vec and BYOL.

The main hyperparameters introduced by this work is the scale and aspect ratio of the target and context blocks. Generally, a small context is used to make this task difficult, which would force the model to learn higher level and more useful features.

V-JEPA: Revisiting Feature Prediction for Learning Visual Representations from Video

V-JEPA is an extension of I-JEPA to videos. This is done by treating videos are 3d images.

1. A clip of 64 frames (~2.1 seconds of video at 30 frames per second) is extracted from the video and resized to 16 × 224 × 224 × 3.
2. The clip is split into \(L\) spatiotemporal patches of size 16x16x2 (2 is the number of consecutive frames.

3. A random mask is calculated for the context. This is a 2D that is similar to the mask in I-JEPA. This mask is then repeated across the time dimension. This repetition is necessary because the videos are short and there would be too much redundancy for the same patch at different time steps. This redundancy would make the learning task too easy. This masking creates a context image, while the target is the original image.

   1. 2 masks are sampled: one short range and one long range. The short range mask covers less area in the image and is more discontinuous. These masks are constructed by different configurations of overlapping blocks, as done in I-JEPA. The target encoder only needs to run once, even if there are multiple masks for the context. Having multiple masks leads to more efficient training.

1. The tokens are processed by a transformer encoder (linear projection of patches + multiple transformer blocks). The masked out patches do not need to be processed. There is a separate encoder for the target and context. The target encoder is an EMA of the context encoder (same as I-JEPA).
2. The predictor predicts the representations of the masked tokens by the unmasked tokens processed by the context encoder. The loss is the L1 distance between the representations of these masked tokens (from the target encoder, and the context encoder + predictor).

This is predicting gaps in short videos. It does not predict across time. Human learning is across the time dimension.

Attentive probing is used to evaluate this model on different finetuning tasks. This is needed in place of linear probing since the input size may vary. This just requires learning a query token specific to the task and a linear classifier on top of the pretrained encoder.

V-JEPA processes small sequences of frames. These short videos are essentially images with a little animation. However, that is the current state of video self-supervised learning. To achieve a model that is closer to human or even animal-level intelligence, this approach needs to scale up significantly. The resolution of the video needs to be increased. Also, the model needs to process longer durations of video and make predictions across time. For example, you should be able to predict what happens in the next 1 minute, based on the previous ten minutes of video input. Such a model could be the basis for an intelligent agent’s world model.

V-JEPA is a very interesting model that may be the start of a highly important line of research.

MC-JEPA: A Joint-Embedding Predictive Architecture for Self-Supervised Learning of Motion and Content Features

This is an extension of JEPA to include motion information. It uses an optical flow objective to learn motion from videos and uses general SSL to learn about the content of the images/videos. Optical flow is estimating the direction in which pixels move between two consecutive frames of a video.

The details of this dense flow estimation are out of the scope of this blog post. Flow estimation and content feature learning are combined as a multitask learning objective. Images are sampled for content learning, while consecutive frames are sampled from videos for flow estimation. The encoder is shared for both tasks. This is a JEPA architecture because the representations from one frame are warped to match the representations from the next frame. The same encoder is used to process both frames.

The architecture for flow estimation is hierarchal. This may be the first instantiation of an H-JEPA architecture. This architecture is based on PWC-Net. Each level has a different resolution.

The image features are sampled from ImageNet, while a video dataset is used for flow estimation. It is also possible to use frames from video as images for content learning.

This work shows that the JEPA framework is generalizable. There are a lot of ways that we could design a world model and it could include many possible objectives.

What’s Next?

The current research in JEPA represents a significant step towards Yann LeCun’s vision of building a world model capable of human-level AI. While the present focus is on creating effective representation learning models for visual data, the ultimate goal is far more ambitious. The holy grail of this research is a V-JEPA model that can predict across extended time horizons, potentially through a Hierarchical JEPA architecture capable of processing complex, lengthy videos like 10-minute YouTube clips.

To realize this vision, several crucial advancements are necessary. Firstly, we need to embrace true multimodality, incorporating audio and other modalities that are often overlooked in current video models. Scaling up V-JEPA is also essential, requiring larger video datasets and more sophisticated model architectures that can handle higher resolutions. Additionally, the development of more challenging benchmarks for video understanding is critical, as current standards fall short of the complexity seen in image or language modeling tasks.

Future iterations of V-JEPA must evolve beyond spatial masking to make predictions across various time horizons. This capability to forecast future representations based on present information is fundamental to understanding the temporal dynamics of video content. Achieving this may necessitate a hierarchical JEPA structure, where different levels handle predictions at various time scales and abstraction levels. Maybe the next JEPA paper will introduce a hierarchal video JEPA (HV-JEPA).

V-JEPA 2

[2025-07-04 Update]

V-JEPA 2 keeps the same model architecture and pretraining method as the original models. However, the model size and the datasets are scaled up. The authors also introduce some interesting post-training methods that unlock more capabilities for this model.

First we will compare the model architectures and datasets, and then delve into the post-training methods.

Model Comparison

Post Training

We have seen the transformer architecture, which started with language, take over many other modalities. A more recent trend is that LLM training methods and recipes (pretraining → SFT → RL) are being adapted to other domains. V-JEPA 2 represents another iteration of this trend. We will explore the different types of post-training implemented in V-JEPA 2.

Progressive-Resolution Training

We also see a form of long context finetuning in the progressive-resolution training. This is part of the video pretraining recipe. While most pretraining is done with 16 frames, the model is then further trained on 64 frames during the cooldown phase. This enables efficiently training a model that can process longer video clips. This is the same way that LLMs are trained to process long contexts.

Attentive Probe Finetuning

This method was implemented in V-JEPA 1. A larger attentive probe is used in comparison to the evals in V-JEPA 1.

This attentive probe finetuning leads to strong performance on action classification and object recognition evals.

V-JEPA 2 introduces an evaluation on action anticipation. The Epic-Kitchens-100 human-action anticipation task evaluates the model’s ability to predict future actions. This is represented with an action, verb, and noun. This can be considered as 3 classification tasks. An attentive probe with three learnable query embeddings is used for this. The eval task is structured so that the model predicts what action is taken 1 second into the future. This is limited in that it can only predict from a fixed set of actions, verbs, and nouns. It also can’t predict multiple actions or on a longer time frame.

LLM Conditioning

V-JEPA 2 is used as a video encoder for an LLM. This enables evaluation on Video Question Answering.

This uses the architecture and visual instruction tuning training method of LLaVA models. The JEPA embeddings are projected into the input of the LLM as non-tokenized early fusion. The visual instruction tuning method from LLaVA is used to train this architecture.

1. Only the projector is unfrozen and trained on image captioning
2. The model is trained end to end on image question answering
3. The model is trained end to end on video question answering

Action Conditioned Post-Training

The model is also evaluated for robotics. It is post-trained on 62 hours of video from the Droid dataset. This model is then deployed on Franka arms. This is a tabletop arm with a fixed camera. The actions are end-effector commands. There are 7 degrees of freedom so that effector states can be represented by the 7D vector. The actions are changes in these states. These diffs are used to form commands.

The V-JEPA 2 encoder \(E(\cdot)\) encodes each of the 16 frames into a feature map. We get 16 feature maps \((z_k)_{k \in [16]}\), where \(z_k := E(x_k) \in \mathbb{R}^{H \times W \times D}\). With 16×16 patches, we get an output of size 16×16×1408. The feature maps are interleaved with the actions and states: \((a_k,s_k,z_k)_{k \in [15]}\). Note that we have 15 of these tuples since the actions are formed by the differences between states.

This is processed by the predictor model \(P_{\phi}(\cdot)\), which is a separate 300M parameter autoregressive transformer. Given the action, state, and previous feature maps, it predicts the next feature map. The result is a sequence of predicted feature maps \((\hat{z}_{k+1})_{k \in [15]}\).

We can then apply a teacher forcing loss, similar to next token prediction SFT in LLMs. Note that \(k=0\) is skipped because there is no prediction at the first timestamp.

A two-step rollout loss (\(T=2\)) is also used. We start with the initial state and feature map, then use the first \(T\) actions to rollout to the \(T+1\) feature map. We then minimize the distance between the true \(T+1\) feature map and the one predicted by the rollout. This differs from the teacher forcing loss since we only have the true state and feature map for the initial state. The subsequent states and feature maps are predicted by the model.

We train on a combination of the teacher forcing and rollout loss.

During this training the V-JEPA 2 encoder is frozen. We train the predictor network to enable planning.

For planning we use the Cross Entropy Method (explained in my World Models blog). This method is used to minimize the goal conditioned energy function at inference time:

\(z_g\) represents the feature map of the goal state. This uses the same method as DINO-WM (World Models blog). We get an action sequence \((a_i^\star)_{i \in [T]}\) through this optimization.

An interesting future direction is to use this world model for language-specified goals. We can train an LLM to generate the goal latent embedding from text. Here is a simple design for implementation:

1. Collect a robotics dataset containing pairs of (action, video), where actions are free-form text and videos are short clips of robots performing these actions.
2. Initialize an MLLM with a JEPA video encoder and a predictor network.
3. Use the action-conditioned training method to train the predictor and the MLLM.
   1. The loss functions would need adaptation:
      1. Add a loss term that forces the MLLM-generated embeddings to be close to the representation of the final image.
      2. Include the MLLM embedding as an input to the predictor. Backpropagate gradients to both the LLM and the predictor.

V-JEPA 2 proves to be an effective world model when applied to robotics planning. This robotics application brings JEPA closer to the brain framework introduced in the original JEPA position paper, demonstrating how to build a JEPA world model and use it for real-world planning.

If you have the resources it can be easy to speed up training by adding more GPUs. However, it is important to understand the impact this scaling will have on training. Machine learning acceleration is a huge and complex field. I intend to just cover the basic intuitions to keep in mind when training a typical model.

There are two types of machine learning training parallelization: data parallelism and model parallelism.

Data Parallelism

Data parallelism splits a training batch into smaller batches for each GPU. Each GPU has its own copy of the model. Each GPU computes gradients with its own training batch. These gradients are then aggregated across all the GPUs. Each GPU can send its gradients to all other GPUs. For example, if you train with a batch size of 64 and 4 GPUs, each GPU will get a batch size of 16. It will compute gradients for this batch. Once all the GPUs are done with their computations, they can send their gradients to each other. The gradients are then averaged and applied to the model. This allows us to train a model with batch size 64 at the speed of batch size 16. However, there is additional latency in communicating the gradients and synchronizing the GPUs, but it is usually negligible compared to the gradient computations.

Another option is to skip the gradient aggregation and simply apply the updates to the model separately. This can be done by having an orchestrator take a lock on the model. Or in Hogwild, you can just update the model without any lock. This will allow some GPU batches to be dropped due to race conditions but minimizes synchronization delays.

Adding GPUs doesn’t make training steps faster. It allows you to have larger mini-batch sizes, which in turn trains models faster.

There are three variables to consider: mini-batch size, GPU batch size, and number of GPUs. Since the number of GPUs is a function of the other two variables, we will only discuss the two types of batch size and how to optimize them.

The implementation of parallelism can vary between ML frameworks: PyTorch, TensorFlow, Jax

Optimal GPU Batch Size

With GPUs, we simply want to minimize the training step time. If we operate the GPUs in the optimal GPU batch size range, we can then just set the number of GPUs to get the optimal mini-batch size. Also, note the GPU batch size has an upper limit from the memory available.

To see how GPU performance relates to speed, I timed the training steps of a ResNet50 model against ImageNet-sized batches of different sizes. I tested batch sizes of every power of two until the GPU ran out of memory.

We see that the throughput (examples per ms) is maximized at the largest possible batch size. We also see that for batch sizes of less than 2^4 or 16, the throughput is lower. GPUs are inefficient with small batches due to overhead. CPUs perform better in some settings. The takeaway is that we want to maximize the GPU utilization by fitting the largest possible batch. Some libraries have the functionality to search for the largest possible batch size given a GPU and dataset. In the flat region, GPU step time increases linearly with batch size.

Optimal Mini Batch Size

To optimize the mini-batch size, we will ignore accelerators and just focus on mini-batch gradient descent. Mini batch size is less hardware-dependent and more problem dependent. We will use ImageNet as an example, but the effects of batch size on training should be considered for every new problem.

Assuming maximize GPU usage/batch size, optimizing the mini-batch size means selecting the number of GPUs to use. The assumption is that with more GPUs, we can train a model faster. Training on more GPUs means faster training epochs. However, we are interested in the test accuracy of the model, not just completing epochs.

Consider this plot from the paper Measuring the Effects of Data Parallelism on Neural Network Training by Shallue et al.

For points on the dashed line, the number of training steps is halved whenever the batch size is doubled. This means that doubling the GPUs/TPUs would have the training time. This is ideal. In this region, you can happily speed up your model training by adding more GPUs that you may have available. However, this tradeoff changes at batch size 2^13 or 8192. From here, doubling the GPUs still speeds up model training, but the speed will be more than half. This is the point of diminishing returns. If you have the GPUs, you might as well use them but those additional GPUs are not as effective.

The paper goes into great detail on this relationship and the effects of other factors such as model architecture, optimizers, and datasets. The takeaway for this blog is that if you set the maximum GPU batch size, up to a point, adding additional GPUs will linearly speed up the training of your model.

Model Parallelism

This type of parallelism is much less commonly used. It can be used along with data parallelism. Model parallelism is when an ML model is too large to fit in the memory of one device. It is partitioned across multiple devices. This has enabled us to train larger and larger networks. For example, the GPT-3 model is about 350 GB. No single GPU can store the whole model in memory.

There are different ways of achieving model parallelism. You can split the model vertically by layer, or horizontally by splitting the tensors.

Pipeline Parallelism

The simplest solution is to process different layers of a neural network on different accelerations. A simple illustration of this:

In PyTorch, the layers are bucketed into groups of roughly equal memory so that the computations are evenly distributed across the accelerators.

A major issue with this approach is that after Layer 0 has a forward pass, it has to wait for the other layers to compute forward and backward passes. The GPU is idle for about 75% of the time. The following diagrams are from the GPipe paper by Huang et al.

The solution is to have the layer compute the next graph while it is waiting for the gradient of the current batch. This essentially combines data parallelism with model parallelism. I explained above that to maximize training speed, we want to maximize the utilization of accelerators. For large models that require model parallelism, we have an additional problem of GPU waiting time. Pipelining GPU batches helps reduce this gap.

We see that with 4 GPU batches, each GPU is idle for about 6/16 of the time. The variables here are the number of GPUs and number of GPU batches per GPU. With 1 GPU batch (no pipelining), the utilization is: \(\frac{2}{n_{GPU}* 2} = \frac{1}{n_{GPU}}\). With pipelining, we get \(\frac{n_{batches}*2}{n_{batches}*2 + 2* (n_{GPU}-1)}\). This simplifies to the following:

This equation explains the tradeoff. Increasing the number of GPU batches drives the utilization closer to 1, while increasing the number of GPUs reduces the utilization.

In an optimal setup, we split the model among as few GPUs as possible, but increase the number of batches that they process in a step. Pipeline parallelism has the added benefit of the speedups of data parallelism. This makes it a very effective solution.

In the PipeDream paper, Harlap et al. show that we can further reduce idle time by interleaving forward and backward operations.

However, this eliminates gradient synchronization. Even eliminating batches above 4, we get the same utilization as GPipe parallelism, just in a different order. For many of the backward passes, a stale version of the model parameters is used. Gradient synchronization is an important tradeoff in all types of ML parallelism.

In analyzing utilization, we have been assuming that forward and backward computations are equivalent. Backward passes tend to take more time. If we interleave operations as to always prioritize backward passes, we can get a utilization gain. From AWS Sagemaker documentation:

The idle time here is 1 forward pass and 1 backward pass.

With backward prioritization, the idle time for GPU 0 is 3 forward passes. This effect will be increased with more GPUs. There are many tradeoffs in parallel ML, such as communication between model layers, the memory overhead of forward and backward passes, different model splits, staleness, etc. We are only covering the high-level intuitions to achieve fast and effective training of large models.

What if we want to use more GPUs for data parallelism, but without splitting up the model further? We can simply run pipelines in parallel. For example, we can split the model among four GPUs but duplicate each model split twice. We can then aggregate the gradients of both pipelines in the update. This is often called hybrid model and data parallelism.

Tensor Parallelism

Instead of splitting the model into layers, we can split the layers themselves. From the Megatron-LM paper by Shoeybi et al.:

The input X has to be completely copied for each split of the model. The layer is split into two halves. The splits of the model are then aggregated in the last layers of the model. Splitting the tensors themselves offers some benefits. The latency is reduced since you can fit more layers on a GPU. This is parallel computation instead of serialized computation. You don’t have to worry about scheduling to minimize idle time.

An issue with this approach is that the activations are also separated, so you are learning a different model architecture. There is an additional cost in concatenating \(Y_1\) and \(Y_2\) for both GPUs. The Megatron-LM architecture is designed to reduce the cost of communicating between GPUs.

Conclusion

We touched the surface on the many tradeoffs, optimizations, and considerations needed for distributed and large scale ML. As models grow larger, it will become more import to understand and keep up to date with this field.

Additional Resources

https://www.youtube.com/watch?v=3XUG7cjte2U

https://lilianweng.github.io/posts/2021-09-25-train-large/

https://openai.com/blog/techniques-for-training-large-neural-networks/

https://timdettmers.com/2014/10/09/deep-learning-data-parallelism/

https://huggingface.co/docs/transformers/v4.15.0/parallelism