[tab] main.tex (latex)
\documentclass[11pt, a4paper, logo, onecolumn, copyright, nonumbering]{deepmind}

\usepackage[authoryear, sort&compress, round]{natbib}
\bibliographystyle{abbrvnat}
\input{math_commands}

\usepackage{floatrow}
\usepackage{varwidth}
\usepackage[]{algorithm2e}
\usepackage{algpseudocode}
\usepackage{bbm}
\usepackage{wrapfig}
\usepackage{siunitx}
\usepackage{cleveref}

% For the python code snippets
\usepackage{fancyvrb}
\usepackage{graphicx}
\usepackage{xcolor}

\usepackage[bottom]{footmisc}

\makeatletter
\def\PYG@reset{\let\PYG@it=
elax \let\PYG@bf=
elax%
    \let\PYG@ul=
elax \let\PYG@tc=
elax%
    \let\PYG@bc=
elax \let\PYG@ff=
elax}
\def\PYG@tok#1{\csname PYG@tok@#1\endcsname}
\def\PYG@toks#1+{\ifx
elax#1\empty\else%
    \PYG@tok{#1}\expandafter\PYG@toks\fi}
\def\PYG@do#1{\PYG@bc{\PYG@tc{\PYG@ul{%
    \PYG@it{\PYG@bf{\PYG@ff{#1}}}}}}}
\def\PYG#1#2{\PYG@reset\PYG@toks#1+
elax+\PYG@do{#2}}

\@namedef{PYG@tok@w}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.73,0.73}{##1}}}
\@namedef{PYG@tok@c}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cp}{\def\PYG@tc##1{	extcolor[rgb]{0.61,0.40,0.00}{##1}}}
\@namedef{PYG@tok@k}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kp}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kt}{\def\PYG@tc##1{	extcolor[rgb]{0.69,0.00,0.25}{##1}}}
\@namedef{PYG@tok@o}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@ow}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.67,0.13,1.00}{##1}}}
\@namedef{PYG@tok@nb}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@nf}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@nc}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@nn}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@ne}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.80,0.25,0.22}{##1}}}
\@namedef{PYG@tok@nv}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@no}{\def\PYG@tc##1{	extcolor[rgb]{0.53,0.00,0.00}{##1}}}
\@namedef{PYG@tok@nl}{\def\PYG@tc##1{	extcolor[rgb]{0.46,0.46,0.00}{##1}}}
\@namedef{PYG@tok@ni}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.44,0.44,0.44}{##1}}}
\@namedef{PYG@tok@na}{\def\PYG@tc##1{	extcolor[rgb]{0.41,0.47,0.13}{##1}}}
\@namedef{PYG@tok@nt}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@nd}{\def\PYG@tc##1{	extcolor[rgb]{0.67,0.13,1.00}{##1}}}
\@namedef{PYG@tok@s}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sd}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@si}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.64,0.35,0.47}{##1}}}
\@namedef{PYG@tok@se}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.67,0.36,0.12}{##1}}}
\@namedef{PYG@tok@sr}{\def\PYG@tc##1{	extcolor[rgb]{0.64,0.35,0.47}{##1}}}
\@namedef{PYG@tok@ss}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@sx}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@m}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@gh}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,0.50}{##1}}}
\@namedef{PYG@tok@gu}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.50,0.00,0.50}{##1}}}
\@namedef{PYG@tok@gd}{\def\PYG@tc##1{	extcolor[rgb]{0.63,0.00,0.00}{##1}}}
\@namedef{PYG@tok@gi}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.52,0.00}{##1}}}
\@namedef{PYG@tok@gr}{\def\PYG@tc##1{	extcolor[rgb]{0.89,0.00,0.00}{##1}}}
\@namedef{PYG@tok@ge}{\let\PYG@it=	extit}
\@namedef{PYG@tok@gs}{\let\PYG@bf=	extbf}
\@namedef{PYG@tok@gp}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,0.50}{##1}}}
\@namedef{PYG@tok@go}{\def\PYG@tc##1{	extcolor[rgb]{0.44,0.44,0.44}{##1}}}
\@namedef{PYG@tok@gt}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.27,0.87}{##1}}}
\@namedef{PYG@tok@err}{\def\PYG@bc##1{{\setlength{\fboxsep}{\string -\fboxrule}\fcolorbox[rgb]{1.00,0.00,0.00}{1,1,1}{\strut ##1}}}}
\@namedef{PYG@tok@kc}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kd}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kn}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kr}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@bp}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@fm}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@vc}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@vg}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@vi}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@vm}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@sa}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sb}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sc}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@dl}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@s2}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sh}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@s1}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@mb}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mf}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mh}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mi}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@il}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mo}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@ch}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cm}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cpf}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@c1}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cs}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}

\def\PYGZbs{\char`\\}
\def\PYGZus{\char`\_}
\def\PYGZob{\char`\{}
\def\PYGZcb{\char`\}}
\def\PYGZca{\char`\^}
\def\PYGZam{\char`\&}
\def\PYGZlt{\char`\<}
\def\PYGZgt{\char`\>}
\def\PYGZsh{\char`\#}
\def\PYGZpc{\char`\%}
\def\PYGZdl{\char`\$}
\def\PYGZhy{\char`\-}
\def\PYGZsq{\char`\'}
\def\PYGZdq{\char`\"}
\def\PYGZti{\char`\~}
% for compatibility with earlier versions
\def\PYGZat{@}
\def\PYGZlb{[}
\def\PYGZrb{]}
\makeatother

	itle{Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase}%
\correspondingauthor{[mihir,alfred,franz]@pdoom.org}

\keywords{World Modeling, JAX, Reinforcement Learning}


eportnumber{001} %

\author{
{\Authfont
Mihir Mahajan	extsuperscript{*1,2},
Alfred Nguyen	extsuperscript{*1,2},
Franz Srambical	extsuperscript{*1,2}}
\leavevmode \protect\\[1em]
{\Authfont
Stefan Bauer	extsuperscript{2}}
\leavevmode \protect\\[1em]
{\Affilfont
	extsuperscript{*}Contributed equally\protect\\
	extsuperscript{1}p(doom), 	extsuperscript{2}TUM 
}}

\begin{abstract}
    While world models are increasingly positioned as a pathway to overcoming data scarcity in domains such as robotics, open training infrastructure for world modeling remains nascent. We introduce Jasmine, a performant JAX-based world modeling codebase that scales from single hosts to hundreds of accelerators with minimal code changes. Jasmine achieves an order-of-magnitude faster reproduction of the CoinRun case study compared to prior open implementations, enabled by performance optimizations across data loading, training and checkpointing. The codebase guarantees fully reproducible training and supports diverse sharding configurations. By pairing Jasmine with curated large-scale datasets, we establish infrastructure for rigorous benchmarking pipelines across model families and architectural ablations.
\end{abstract}


\begin{document}
\maketitle

\section{Introduction}

Over the past decades, the field of deep learning has increasingly been shaped by methods that leverage vast data troves \citep{jozefowicz2016exploring, radford2018improving, radford2019language, chowdhery2022palm, raffel2020exploring, JMLR:v23:21-0998, deng2009imagenet}, and paradigms that unlock new ones \citep{guo2025deepseek, christiano2017deep, radford2021learning, srambical2025crowd-sourcing, silver2016mastering, berner2019dota}. Internet-scale pre-training, preference modeling, and reinforcement learning using verification signals offer a compelling pathway for language models to attain human-level performance \citep{LuongLockhart2025GeminiIMO, LinCheng2025GeminiICPC}, yet data is increasingly bottlenecking progress from spiky towards general intelligence. While some domains can leverage user feedback from deployed products for iterative model improvement \citep{cursor2025tab}, domains like robotics cannot afford such a privilege.

One paradigm proposed by the research community to overcome the data scarcity in those domains is that of world models \citep{ha2018world}. World models can aid frontier model development in numerous ways, but one particularly ambitious goal of the community is to train a world model to act as a simulation of the real world \citep{bruce2024genie, parker2022evolving, deepmind2025genie3, agarwal2025cosmos}, in order to train agents in that simulation \citep{hafner2025training}, via an adaptive curriculum \citep{parker2022evolving}, or otherwise. This regime requires the compounding error of the world model to be orders of magnitude smaller than when solely used for short-term look-ahead. The feasibility of such a world model in its truest sense is entirely understudied.
Jasmine provides the foundational infrastructure for future empirical investigation of how compute and data requirements scale with environment complexity for downstream agent training.

\begin{figure}
  \centering
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/samples/cropped/gt_jafar_prepend_105.png}
  \caption{Autoregressive sampling of Jafar \citep{willi2024jafar} (middle row) and Jasmine (bottom row) on the CoinRun case study with four conditioning frames (conditioning frames not shown). 
  The top row shows the ground-truth sequence. 
}

\label{fig:sampling}
\end{figure}

\section{Jasmine}
\label{sec:jasmine}



	extbf{Our contributions} in this work are threefold: i) we introduce Jasmine, a highly optimized and scalable JAX-based codebase for world modeling, which we use to reproduce Genie's CoinRun case study \citep{bruce2024genie,cobbe2019leveraging} an order of magnitude faster than prior work \citep{willi2024jafar}. This speedup is the result of infrastructure optimizations, including a fully reproducible, scalable training and data pipeline built on the JAX \citep{jax2018github} ecosystem. ii) We find that a critical modification to the original Genie architecture, prepending latent actions instead of adding them to video embeddings, is required for the world model to yield generations that faithfully reproduce the CoinRun environment. iii) Finally, we openly release the Jasmine codebase, along with pretrained checkpoints, curated datasets, model inspection notebooks, and a dataset of dense IDE interactions captured during Jasmine's development, providing the first openly published dataset of months-long software engineering.

Jasmine implements the Genie \citep{bruce2024genie} architecture, enabling training of interactive environments from unlabeled videos. The architecture includes a video tokenizer, which encodes videos into tokens, a latent action model (LAM) that extracts latent actions between video frames, and a dynamics model that predicts the tokens of the next frame based on the previous tokens and corresponding latent actions. At sampling time, the LAM is discarded and replaced by input from the user. All modules use an ST-Transformer \citep{ho2019axial} backbone which approximates full attention by performing intra-frame (spatial) followed by inter-frame (temporal) attention, thus reducing the attention sequence length.
The tokenizer uses a VQ-VAE \citep{van2017neural} to encode image patches using reconstruction, vector-quantization, and commitment losses.
To train an action-conditioned video-generation model from unlabeled videos, Genie learns latent actions \citep{schmidt2023learning}. Like the tokenizer, the LAM uses a VQ-VAE, with its codebook representing the latent actions. The model learns to distill information from future frames into this bottlenecked codebook:
Frames $x_{0:t}$ are encoded, producing latent actions $a_{0:t}$, which the decoder receives along with past frames $x_{0:t-1}$ to predict the next frame $x_{t}$. A temporal causal mask allows the entire sequence to be processed in a single forward pass.
The dynamics model is a decoder-only transformer that predicts future frames conditioned on past frames and corresponding latent actions. Genie uses MaskGIT \citep{chang2022maskgit}, which masks input-tokens at training time, similar in spirit to BERT \citep{devlin2019bert}. Unlike MaskGIT, Genie masks with probability $p \sim U(0.5, 1)$ (refer to Appendix 
ef{sec:maskgit_to_videos} for details about extending MaskGIT to videos).


\begin{figure}[htb]
  \centering
  \includegraphics[width=	extwidth]{fig/case-study-plots.pdf}
  \caption{Validation metrics of the CoinRun case study (patch size 4). While the loss (left) is similar between the default Genie configuration and our minimal modification, rollout metrics (middle and right, refer to \Cref{sec:experiment-metrics}) differ substantially.}
  \label{fig:case-study-plots}
\end{figure}



Building upon prior work \citep{willi2024jafar} that openly published a reimplementation of the Genie \citep{bruce2024genie} architecture, we release a highly optimized JAX-based world modeling codebase amenable to scale. Jasmine implements a range of baselines, including MaskGIT-based \citep{chang2022maskgit}, fully causal \citep{srambical2024going}, and diffusion-based approaches (Appendix 
ef{sec:diffusion-baseline}). The codebase depends solely  on battle-tested libraries from the Google ecosystem (JAX, NNX, Grain \citep{grain2023github}, Orbax, Optax, Treescope \citep{johnson2024penzai}, ArrayRecord \citep{ArrayRecord}), and scales from single hosts to hundreds of accelerators using XLA.
Jasmine supports complex sharding configurations in a few lines of code through Shardy \citep{openxla-shardy}. It provides asynchronous distributed checkpointing with configurable policies, process-parallel dataloading, and checkpointing of model, optimizer, and data loader states. Training runs are bitwise deterministic, yielding identical loss curves under identical seeds (Appendix 
ef{sec:bitwise_deterministic}).
To enable efficient large-scale experimentation, Jasmine integrates mixed-precision, FlashAttention via cuDNN SDPA \citep{NVIDIA_cuDNN_Attention}, activation checkpointing, host memory offloading, and index-shuffling during data loading. The codebase follows the shape suffix convention of \citet{shazeer2024shape}, aiming to provide a didactic implementation of modern world modeling architectures.

We run the reproducible case study described in \citet{bruce2024genie} by generating a dataset containing 50M transitions of CoinRun, an environment of the Procgen benchmark \citep{cobbe2019leveraging} (Appendix 
ef{sec:coinrun-case-study}). In contrast to \citet{bruce2024genie} we find that strict adoption of the architecture and hyperparameters described in their case study leads to deteriorating autoregressive generations in both Jasmine and Jafar \citep{willi2024jafar} (Figure 
ef{fig:sampling_prepend_vs_no_prepend}, middle row). However, a minimal modification of the case study setting, namely prepending latent actions instead of adding them to the video embeddings, yields autoregressive generations that faithfully simulate the CoinRun environment (Figure 
ef{fig:sampling_prepend_vs_no_prepend}, bottom row and Figure 
ef{fig:case-study-plots}). We hypothesize that this discrepancy between \citet{bruce2024genie} and our work stems from an ambiguity in extending MaskGIT to videos (refer to Appendix 
ef{sec:maskgit_to_videos} for further discussion).

\begin{figure}[htb]
  \centering
  \includegraphics[width=0.75	extwidth]{fig/1-wc-time-comparison.pdf}
  \caption{An order of magnitude faster convergence in wall-clock time in Jasmine (blue) compared to Jafar \citep{willi2024jafar} (orange). We report the train loss since Jafar does not collect validation metrics. Refer to Appendix 
ef{sec:arch-ablations} for Jasmine's validation metrics. Jasmine's lower variance stems from a subtle refinement in its batched masking logic (Appendix 
ef{sec:jafar-batched-masking}).}
  \label{fig:wc_time}
\end{figure}


Beyond openly publishing\footnote{\url{https://github.com/p-doom/jasmine}} the Jasmine codebase, we recorded every keystroke during Jasmine's development using \emph{crowd-code} \citep{srambical2025crowd-sourcing}, a VS Code/Cursor extension that enables crowd-sourcing dense IDE interaction data. To our knowledge, this represents one of the first open datasets capturing the full temporal scale of months-long software development, thus laying groundwork for future work in behaviour cloning, goal-conditioning, and verification signal mining.

\section{Experiments}

We evaluate the performance of Jasmine and analyze the impact of its core components through rigorous ablations. Using Jasmine, we reproduce the CoinRun case study at a patch size of 16 in under nine hours on a single GPU, compared to over 100 hours reported in prior work \citep{willi2024jafar} under the same setting (Figure 
ef{fig:wc_time}). We present ablations identifying the factors responsible for this speedup in Table 
ef{tab:infra-ablations}.
\Cref{sec:arch-ablations} further reports results from architectural modifications including replacing the latent action model with ground-truth actions, ablating co-training, adopting fully causal and diffusion baselines, and setting the feedforward expansion factor to four.


\paragraph{Architectural optimizations}
\label{sec:speed-arch-ablations}
We adapt Genie's architectural choices by integrating best practices from the language modeling community. Specifically, we use a feedforward expansion factor of four relative to the model dimension, following common practice in large-scale language modeling \citep{raffel2020exploring,radford2019language,brown2020language}. We simultaneously reduce network depth, resulting in lower overall parameter count than the Genie defaults, thus achieving higher throughput (Table 
ef{tab:ffn-ablation}) while maintaing competitive performance (Figure 
ef{fig:arch-ablations}). We employ the warmup-stable-decay (WSD) learning rate schedule \citep{zhai2022scaling, mahajan2018exploring}, which allows flexible training durations by resuming from a checkpoint prior to the decay phase. Unlike \citep{bruce2024genie}, we warm up from and decay the learning rate to zero, in line with established best practices \citep{zhai2022scaling}.
We further compare co-training LAM and dynamics model (as done in \citet{bruce2024genie}) with pre-training the LAM (as done in \citet{willi2024jafar}), embedding ground-truth actions instead of using the latent action model (Appendix 
ef{sec:ablation-gt}), and replacing MaskGIT with fully causal and diffusion baselines. Co-training, pre-training the LAM, and using ground-truth actions are all competitive (Figure 
ef{fig:arch-ablations}), while the fully causal baseline underperforms in the 200k steps training regime (\Cref{fig:causal-ablations}). However, our results indicate that the fully causal baseline in particular may benefit from longer training. Diffusion-forcing \citep{chen2024diffusion} outperforms MaskGIT, even when using identical per-frame sampling step counts and untuned hyperparameters (Figure 
ef{fig:diffusion-plot}, Appendix 
ef{sec:diffusion-baseline}).

\begin{table}[htb]
\begin{center}
\begin{tabular}{lcc}
	oprule
 & 	extbf{Throughput (bs=36)} & 	extbf{Throughput (bs=2048)} \\
\midrule
Jasmine-base & 1.00x	            & 1.00x \\
1x feedforward expansion & 0.93x	            & 0.79x \\
\bottomrule
\end{tabular}
\caption{Training throughput with a feedforward expansion factor of one, relative to Jasmine-base. We double the number of layers compared to Jasmine-base to roughly match the parameter count (refer to Genie's default configuration in 
ef{tab:hparams}).}
\label{tab:ffn-ablation}
\end{center}
\end{table}

\paragraph{Infrastructure optimizations}
\label{sec:speed-infra-ablations}
A substantial portion of our speedup compared to \citet{willi2024jafar} arises from our data loader design (Tables 
ef{tab:chunking-formats-1} and 
ef{tab:chunking-formats-2}). We use Grain for data loading with prefetching enabled and preprocess datasets into ArrayRecords \citep{ArrayRecord}, a file format optimized for random access indexing. The chosen chunking strategy significantly affects throughput, and we describe our configuration in Appendix 
ef{sec:throughput-ablation}. Jasmine further leverages FlashAttention \citep{dao2022flashattention} via cuDNN SDPA \citep{NVIDIA_cuDNN_Attention} and mixed precision training with bfloat16. We report throughputs when ablating mixed precision, FlashAttention and Grain in Table 
ef{tab:infra-ablations}.

\section{Related Work and Discussion}
Although research on world models with its inception decades ago \citep{sutton1991dyna} has matured over the years \citep{ha2018world,hafner2019learning,hafner2019dream,hafner2020mastering,hafner2023mastering,alonso2024diffusion}, they have only been scaled up recently \citep{hafner2025training,bruce2024genie,parkerholder2024genie2,deepmind2025genie3,valevski2024diffusion,agarwal2025cosmos,hu2023gaia,guo2025mineworld,decart2024oasis,lucid2024lucidv1,li2025hunyuan,pearce2024scaling}. While the open training ecosystem in language modeling provides mature solutions for large-scale language pretraining \citep{megatron-lm,maxtext-library}, open training infrastructure for world modeling is still nascent \citep{wayfarer_labs_owl_wms,Savov_2025_CVPR}. Closest to our work is \citet{willi2024jafar}, an open-source reproduction of Genie \citep{bruce2024genie}, which we build upon and significantly extend.

With Jasmine we make progress towards democratizing world modeling research. Alongside the codebase, we openly release checkpoints and datasets for CoinRun, Atari and Doom, as well as dense IDE interaction data collected over months of research engineering\footnote{Datasets published under the CC0 license at \url{https://huggingface.co/datasets/p-doom}}. While Jasmine greatly accelerates wall-clock convergence compared to prior work, it has yet to match throughput efficiencies of frontier language model implementations.

\subsubsection*{Acknowledgments}
We thank Matthew T. Jackson, Andrea Dittadi and Diego Marti Monso for useful discussions as well as the Jafar authors for openly publishing their repository. 
The authors gratefully acknowledge the computing time provided on the high-performance computer HoreKa by the National High-Performance Computing Center at KIT (NHR@KIT). This center is jointly supported by the Federal Ministry of Education and Research and the Ministry of Science, Research and the Arts of Baden-Württemberg, as part of the National High-Performance Computing (NHR) joint funding program (https://www.nhr-verein.de/en/our-partners). HoreKa is partly funded by the German Research Foundation (DFG).

\bibliography{bibliography}
\appendix

\begin{figure}
  \centering
  \includegraphics[width=	extwidth]{fig/0-vanilla-action-prepend.png}
  \caption{Autoregressive sampling of Jasmine when adding (middle row) and prepending actions (bottom row) on the CoinRun case study with four conditioning frames (conditioning frames not shown). 
  The top row shows the ground-truth sequence. 
}
\label{fig:sampling_prepend_vs_no_prepend}
\end{figure}

\section{Coinrun Case Study}
\label{sec:coinrun-case-study}
For the CoinRun case study, we strictly adhere to the setting of \citet{bruce2024genie} and train our models to unmask sequences of 16 frames with a resolution of 64x64 pixels per frame. To generate the dataset, we capture 50M observation frames and corresponding ground-truth actions during random agent rollouts and only use the ground-truth actions for a LAM ablation (\Cref{sec:ablation-gt}). Instead of sampling seeds from a fixed pool as described in \citet{bruce2024genie}, we initialize all episodes with a seed unique to the respective episode. Furthermore, we verify that our generated dataset contains no duplicate episodes and only find 7.46\% duplicate frames. 
We confirm that the validation and test set are disjoint from the train set and publish our script for duplication detection along with the repository\footnote{\url{{https://github.com/p-doom/jasmine/blob/main/data/jasmine_data/detect_array_record_duplicates.py}}}.
While train metrics are near-identical between Genie's configuration and our action-prepending modification, rollout quality differs significantly (\Cref{fig:case-study-plots}). We collect rollout metrics during training (\Cref{sec:experiment-metrics}) that capture this discrepancy.






\section{Ablations}
For our ablations, we reuse the CoinRun setting and ablate from Jasmine's base configuration, depicted in \Cref{tab:hparams}.
\label{sec:arch-ablations}
\begin{figure}
  \centering
  \includegraphics[width=1.0	extwidth]{fig/3-arch-ablations.pdf}
  \caption{Architectural ablations of Jasmine's base configuration (refer to \Cref{tab:hparams}) on CoinRun. We report loss (left) and rollout metrics (middle and right) of the dynamics model on a validation set.}
  \label{fig:arch-ablations}
\end{figure}

\paragraph{Co-training LAM and dynamics model}
\label{sec:ablation-co-training}
\citet{bruce2024genie} co-train the LAM and the dynamics model. However, their implementation remains unclear as the LAM is supervised on frames while the dynamics model is supervised on tokens. One approach to co-training is a combined loss function including stop-gradients that prevent gradients from flowing from the dynamics model to the LAM. However, such a combined loss formulation remains unmentioned in \citet{bruce2024genie}. 
\citet{willi2024jafar} instead train LAM and dynamics model sequentially, thus reducing memory footprint at the cost of longer total training time. 
For Jasmine's co-training implementation, we omit the LAM decoder entirely and allow gradients to flow from the dynamics model to the LAM.

\paragraph{Training with ground-truth actions}
\label{sec:ablation-gt}
We ablate the LAM (Figure 
ef{fig:arch-ablations}) by training the dynamics model using ground-truth actions captured in the environment. We use an embedding layer to map action indices to action latents, which are then used as additional input to the dynamics model.

\begin{figure}
  \centering
  \includegraphics[width=1.0	extwidth]{fig/2-causal.pdf}
  \caption{Loss (left) and rollout metrics (middle and right) of the fully causal baseline. We depict the final performance of our MaskGIT implementation (Jasmine's default configuration) for the rollout metrics. The loss indicates that the causal baseline might benefit from longer training and separately tuned hyperparameters. The losses between the two architectures are not comparable, hence we omit the MaskGIT loss.}
  \label{fig:causal-ablations}
\end{figure}

\paragraph{Throughput ablations}
\label{sec:throughput-ablation}

We ablate core components of Jasmine's infrastructure in Table 
ef{tab:infra-ablations}. Replacing Grain with the default data loader of \citet{willi2024jafar} reduces throughput by an order of magnitude. In Jasmine's base CoinRun configuration (sub-100M parameter model and a maximum attention sequence length of 16) the XLA compiler dispatches higher-throughput CUDA kernels than FlashAttention. While FlashAttention only outperforms XLA-compiled kernels at large model sizes and sequence lengths (\Cref{tab:big-model-flash-attn-ablation}), we enable it by default to reduce accelerator memory usage. At small batch sizes, the compiled train loop of our configuration achieves higher throughput using full precision. We attribute this to XLA largely operating based on heuristics, and posit that writing optimized Pallas kernels for key operations in the model forward pass will result in mixed precision outperforming full precision in throughput, even at small batch sizes. 

The ArrayRecord file format allows storing a configurable amount of records per file. In our case, each record corresponds to a sequence of frames and actions. We find that the chosen format significantly affects throughput (Tables 
ef{tab:chunking-formats-1} and 
ef{tab:chunking-formats-2}). Based on preliminary experiments, we preprocess the dataset to have 100 records per ArrayRecord file with 160 frames per record.

\begin{table}
\begin{center}
\begin{tabular}{lcc}
	oprule
 & 	extbf{Throughput (bs=36)} & 	extbf{Throughput (bs=2048)} \\
\midrule
Jasmine-base & 1.00x	            & 1.00x \\
w/o grain data loader & 0.25x     & 0.11x \\
w/o flash attention    & 1.15x	& 1.04x \\
w/o mixed precision & 1.18x	& 0.71x \\
\bottomrule
\end{tabular}
\caption{Training throughput of infrastructure ablations, relative to Jasmine-base. We report the throughput at Genie's default batch size (36), as well as at the batch size resulting in the highest throughput (2048) on a single H100 with 80GB of accelerator memory.}
\label{tab:infra-ablations}
\end{center}
\end{table}

\begin{table}
\begin{center}
\begin{tabular}{lcc}
	oprule
 & 	extbf{Throughput (frames/sec)} & 	extbf{Throughput (relative)} \\
\midrule
w/ flash attention & 36.15	         &  1.00x \\
w/o flash attention & 24.24 &  0.67x \\
\bottomrule
\end{tabular}
\caption{Training throughput using a larger model (1B parameter) and spatial sequence length (1024). We decrease the patch size to two. In this regime, FlashAttention yields higher throughput than the XLA compiler.}
\label{tab:big-model-flash-attn-ablation}
\end{center}
\end{table}

\begin{table}
\begin{center}
\begin{tabular}{llcc}
	oprule
	extbf{\# frames per record} & 	extbf{\# records per file} & 	extbf{Throughput (frames/sec)} & 	extbf{Throughput (relative)}   \\
\midrule
16          & 100 & 7,527.27 & 1.12x \\
160 (Base)  & 100 & 6,709.09 & 1.00x \\
1,600       & 100 & 3,752.73 & 0.56x \\
16,000      & 100 & 3,720.00 & 0.55x \\
160,000     & 100 & 3,785.45 & 0.56x \\

\bottomrule
\end{tabular}
\caption{Training throughput at Genie's default batch size (36) with different number of frames per record. Throughput decreases as the number of frames per record increases. We opt for 160 frames per record to be able to vary the sequence length.} 
\label{tab:chunking-formats-1}
\end{center}
\end{table}

\begin{table}
\begin{center}
\begin{tabular}{llcc}
	oprule
	extbf{\# frames per record} & 	extbf{\# records per file} & 	extbf{Throughput (frames/sec)} & 	extbf{Throughput (relative)}\\
\midrule
160         & 1         & 6,098.18 & 0.91x \\
160         & 10        & 6,643.64 & 0.99x \\
160 (Base)  & 100       & 6,709.09 & 1.00x \\
160         & 1,000     & 6,480.00 & 0.97x \\
160         & 10,000    & 6,763.64 & 1.01x \\
\bottomrule
\end{tabular}
\caption{Training throughput at Genie's default batch size (36) with different number of records per file.} 
\label{tab:chunking-formats-2}
\end{center}
\end{table}

\section{Diffusion Baseline}
\label{sec:diffusion-baseline}
We implement a diffusion baseline inspired by the Dreamer 4 architecture \citep{hafner2025training}, combining a masked autoencoder (MAE, \citet{he2022masked}) tokenizer  with an ST-DiT \citep{ho2019axial,peebles2023scalable} dynamics model trained under the diffusion-forcing objective \citep{chen2024diffusion}.
We leave implementing the shortcut objective \citep{frans2024one} to future work. 
\paragraph{Tokenizer} Following \citet{hafner2025training}, we use a MAE to compress raw video frames into continuous latents. Our autoencoder implementation uses an ST-Transformer backbone and a latent bottleneck. Before passing the latents to the decoder, we apply the tanh activation to constrain them to the range $(-1, 1)$ for downstream dynamics model training. We uniformly sample per-frame masking probabilities $p_i \sim U(0, 0.9)$. Unlike \citet{hafner2025training}, we omit the auxiliary LPIPS loss \citep{zhang2018unreasonable} and directly train on pixel-level reconstructions using mean-squared error. We find the tokenizer hyperparameters from \Cref{tab:hparams} to work well for MAE training as well.

\paragraph{Dynamics model}
We implement diffusion forcing \citep{chen2024diffusion}, sampling an independent noise level per frame during training. Analogous to \citet{hafner2025training} and Jasmine-base, we prepend latent actions and the embedded denoising step to the patch latents. Following \citet{hafner2025training}, we use x-prediction\footnote{In the diffusion literature, x-prediction often refers to supervision in latent-space rather than pixel-space. We follow that nomenclature but believe that the term z-prediction is a more accurate description.} and employ a ramp loss. During inference, frame-wise latents are autoregressively generated with 25 denoising steps per frame, while past input latents are slightly corrupted using a noise level of 0.1. We adopt the hyperparameters of Jasmine-base, and only change the learning rate to 1e-4 following \citet{peebles2023scalable}.

\begin{figure}
  \centering
  \includegraphics[width=1.0	extwidth]{fig/4-diffusion-maskgit-comparison.pdf}
  \caption{Loss (left) and rollout metrics (middle and right) of the diffusion baseline. We omit the MaskGIT loss as the losses are not comparable between the two architectures.}
  \label{fig:diffusion-plot}
\end{figure}

\section{Bitwise Determinism}
\label{sec:bitwise_deterministic}
On TPUs, bitwise determinism is guaranteed by Jasmine via proper usage of JAX's implementation of parallel random number generation via Threefry counters \citep{salmon2011parallel}. On GPUs however, an additional XLA flag (	exttt{xla\_gpu\_deterministic\_ops=true}) is needed in certain cases to guarantee identical training curves.

\section{Hyperparameter Configurations}
\label{sec:default-config}
We mention four distinct training configurations. We present hyperparameters and training settings of each configuration in \Cref{tab:hparams} and briefly describe them:

\begin{itemize}
    \item 	extbf{Genie}: The hyperparameter configuration of Appendix F of \citet{bruce2024genie}, but with a patch size of 16 (for easier comparison against Jafar). We use this configuration for the CoinRun case study (\Cref{fig:sampling_prepend_vs_no_prepend}, middle row; refer to \Cref{sec:coinrun-case-study}).
    \item 	extbf{Genie w/ prepend}: As detailed in \Cref{sec:jasmine}, we found a minimal modification to the Genie configuration to be necessary to yield generations faithful to the CoinRun environment (Figure 
ef{fig:sampling}, bottom row and \Cref{fig:sampling_prepend_vs_no_prepend}, bottom row). This setting is identical to our Genie configuration, with the exception that we prepend latent actions to the video embeddings instead of adding them.
    \item 	extbf{Jafar}: The hyperparameters used by Jafar \citep{willi2024jafar} for their CoinRun case study. This setting is identical to \citet{bruce2024genie}, but \citet{willi2024jafar} use a patch size of 16 for faster training. Unlike \citet{bruce2024genie}, they pre-train the LAM as opposed to co-training the LAM with the dynamics model. We use this configuration for the Jafar baseline runs (Figure 
ef{fig:sampling}, middle row and \Cref{fig:wc_time}). We solely run this configuration with the Jafar repository.
    \item 	extbf{Jasmine-base}: We define a base configuration for Jasmine that represents a trade-off between training speed, modeling quality and simplicity, integrating best practices from the language modeling literature. This is Jasmine's configuration in the wall-clock convergence comparison (\Cref{fig:wc_time}) between Jasmine and Jafar, as well as our architectural and infrastructure ablations (Figures 
ef{fig:arch-ablations}, 
ef{fig:sampling_multi} and Tables 
ef{tab:ffn-ablation}, 
ef{tab:infra-ablations}, 
ef{tab:big-model-flash-attn-ablation}, 
ef{tab:chunking-formats-1}, 
ef{tab:chunking-formats-2}).
\end{itemize}

\begin{table}
\begin{center}
\begin{tabular}{ll| r|r|r|r}
	oprule
 & 	extbf{Parameter} & 	extbf{Genie} & 	extbf{Genie w/ prepend} & 	extbf{Jafar} & 	extbf{Jasmine-base} \\ 
\midrule
\midrule
Tokenizer & \# blocks    & 8      & & & 4 \\
          & \# heads     & 8      & & & \\
          & model dim       & 512    & & & \\
          & ffn dim       & 512 & & & 2048 \\
          & \# codes & 1024 & & & \\
          & latent dim & 32 & & & \\
          & patch size & 16 & & & \\
          & total train steps  & 300k  & & & \\
          & learning rate   & $3 * 10^{-4}$  & & & \\
          & lr decay end & $3 * 10^{-4}$ & & &  0 \\ 
          & batch size & 48 & & & \\
\midrule
LAM       & \# blocks    & 8      & & & 4 \\
          & \# heads     & 8      & & & \\
          & model dim       & 512    & & & \\
          & ffn dim       & 512 & & & 2048 \\
          & \# codes & 6 & & & \\
          & latent dim & 32 & & & \\
          & patch size & 16 & & & \\
          & total train steps  & 200k  & & & \\
          & learning rate   & $3 * 10^{-5}$ & & & \\
          & lr decay end & $3 * 10^{-6}$ & & & 0 \\ 
          & batch size   & 48 & & & \\
\midrule
Dynamics  & \# blocks    & 12      & & & 6 \\
          & \# heads     & 8      & & & \\
          & model dim       & 512    & & & \\
          & ffn dim       & 512 & & & 2048 \\
          & total train steps  & 200k  & & & \\
          & learning rate   & $3 * 10^{-5}$  & & & \\
          & lr decay end & $3 * 10^{-6}$ & & & 0 \\ 
          & batch size   & 36 & & & \\
          & action conditioning & additive & prepend & & prepend \\
          & baseline & MaskGIT & & & \\
\midrule
\midrule
Training & optimizer    & AdamW      & & & \\
          & lr schedule & cos     & & & wsd \\
          & warmup steps  & 1k & & & \\
          & wsd decay steps & - & & & 10\% \\
          & dataset size (frames) & 50M & & & \\
          & co-training & yes & & no & \\
\midrule
\midrule
Inference & temperature & 1.0 & & & \\
          & maskgit steps & 25 & & & \\
\bottomrule
\end{tabular}
  \caption{Configurations used in our experiments. We only show the difference to our base Genie configuration.
  Note that \citet{bruce2024genie} uses a tokenizer patch size of four, and that we use an expanded dataset with 50M frames for all of our runs (to ensure that no method performs worse due to overfitting).}
  \label{tab:hparams}
\end{center}
\end{table}

\section{Extending MaskGIT to Videos}
\label{sec:maskgit_to_videos}
MaskGIT \citep{chang2022maskgit} is defined on images and there are multiple ways to extend it to videos. We follow \citet{willi2024jafar} by randomly masking tokens in the entire sequence using the uniformly sampled probability $p \sim U(0.5,1)$. An alternative would be to sample a different masking probability per frame, similar to \citet{chen2024diffusion}, or leaving $k \sim U(0,T)$ frames unmasked to closely emulate inference.

\section{Evaluation Metrics}
\label{sec:experiment-metrics}
In early experiments, we found the Genie configuration to suffer from a discrepancy in performance between validation loss and autoregressive rollouts (Figure 
ef{fig:case-study-plots}). Therefore, besides validation loss, Jasmine also tracks rollout metrics throughout training of the dynamics model. We generate a single frame using the sampling logic of the respective architecture and calculate SSIM and PSNR between the generated frame and the ground-truth. 
Although rollout metrics are calculated on a single frame, we find that they directly correlate with the model's performance in generating full rollouts. 
Validation and rollout metrics are calculated on a validation set. The rollouts in Figures 
ef{fig:sampling}, 
ef{fig:sampling_prepend_vs_no_prepend}, 
ef{fig:sampling_multi} and 
ef{fig:sampling_diffusion} are sampled from frames of a test set.

\section{Jafar's Batched Masking Logic}
\label{sec:jafar-batched-masking}
\begin{figure*} 
    \centering 
    \begin{Verbatim}[commandchars=\\\{\}]
\PYG{n}{mask\PYGZus{}prob} \PYG{o}{=} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{random}\PYG{o}{.}\PYG{n}{uniform}\PYG{p}{(}\PYG{n}{rng1}\PYG{p}{,} \PYG{n}{minval}\PYG{o}{=}\PYG{n+nb}{self}\PYG{o}{.}\PYG{n}{mask\PYGZus{}limit}\PYG{p}{)}
\PYG{n}{mask} \PYG{o}{=} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{random}\PYG{o}{.}\PYG{n}{bernoulli}\PYG{p}{(}\PYG{n}{rng2}\PYG{p}{,} \PYG{n}{mask\PYGZus{}prob}\PYG{p}{,} \PYG{n}{vid\PYGZus{}embed}\PYG{o}{.}\PYG{n}{shape}\PYG{p}{[}\PYG{o}{:}\PYG{o}{\PYGZhy{}}\PYG{l+m+mi}{1}\PYG{p}{]}\PYG{p}{)}
\PYG{n}{mask} \PYG{o}{=} \PYG{n}{mask}\PYG{o}{.}\PYG{n}{at}\PYG{p}{[}\PYG{p}{:}\PYG{p}{,} \PYG{l+m+mi}{0}\PYG{p}{]}\PYG{o}{.}\PYG{n}{set}\PYG{p}{(}\PYG{k+kc}{False}\PYG{p}{)}
\PYG{n}{vid\PYGZus{}embed} \PYG{o}{=} \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{where}\PYG{p}{(}\PYG{n}{jnp}\PYG{o}{.}\PYG{n}{expand\PYGZus{}dims}\PYG{p}{(}\PYG{n}{mask}\PYG{p}{,} \PYG{o}{\PYGZhy{}}\PYG{l+m+mi}{1}\PYG{p}{),} \PYG{n+nb}{self}\PYG{o}{.}\PYG{n}{mask\PYGZus{}token}\PYG{p}{,} \PYG{n}{vid\PYGZus{}embed}\PYG{p}{)}
\end{Verbatim} 
\caption{Code snippet from \citet{willi2024jafar} showing their batched masking logic.} \label{fig:jafar-batched-masking} 
\end{figure*}

\begin{figure*} 
    \centering 
    \begin{Verbatim}[commandchars=\\\{\}]
\PYG{n}{mask\PYGZus{}prob} \PYG{o}{=} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{random}\PYG{o}{.}\PYG{n}{uniform}\PYG{p}{(}
    \PYG{n}{\PYGZus{}rng\PYGZus{}prob}\PYG{p}{,} \PYG{n}{shape}\PYG{o}{=}\PYG{p}{(}\PYG{n}{batch\PYGZus{}size}\PYG{p}{,}\PYG{p}{),} \PYG{n}{minval}\PYG{o}{=}\PYG{n+nb}{self}\PYG{o}{.}\PYG{n}{mask\PYGZus{}limit}
\PYG{p}{)}
\PYG{n}{per\PYGZus{}sample\PYGZus{}shape} \PYG{o}{=} \PYG{n}{vid\PYGZus{}embed\PYGZus{}BTNM}\PYG{o}{.}\PYG{n}{shape}\PYG{p}{[}\PYG{l+m+mi}{1}\PYG{p}{:}\PYG{o}{\PYGZhy{}}\PYG{l+m+mi}{1}\PYG{p}{]}
\PYG{n}{mask} \PYG{o}{=} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{vmap}\PYG{p}{(}
    \PYG{k}{lambda} \PYG{n}{rng}\PYG{p}{,} \PYG{n}{prob}\PYG{p}{:} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{random}\PYG{o}{.}\PYG{n}{bernoulli}\PYG{p}{(}\PYG{n}{rng}\PYG{p}{,} \PYG{n}{prob}\PYG{p}{,} \PYG{n}{per\PYGZus{}sample\PYGZus{}shape}\PYG{p}{),}
    \PYG{n}{in\PYGZus{}axes}\PYG{o}{=}\PYG{p}{(}\PYG{l+m+mi}{0}\PYG{p}{,} \PYG{l+m+mi}{0}\PYG{p}{),}
\PYG{p}{)(}\PYG{n}{jnp}\PYG{o}{.}\PYG{n}{asarray}\PYG{p}{(}\PYG{n}{\PYGZus{}rngs\PYGZus{}mask}\PYG{p}{),} \PYG{n}{mask\PYGZus{}prob}\PYG{p}{)}
\PYG{n}{mask} \PYG{o}{=} \PYG{n}{mask}\PYG{o}{.}\PYG{n}{at}\PYG{p}{[}\PYG{p}{:}\PYG{p}{,} \PYG{l+m+mi}{0}\PYG{p}{]}\PYG{o}{.}\PYG{n}{set}\PYG{p}{(}\PYG{k+kc}{False}\PYG{p}{)}
\PYG{n}{vid\PYGZus{}embed\PYGZus{}BTNM} \PYG{o}{=} \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{where}\PYG{p}{(}
    \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{expand\PYGZus{}dims}\PYG{p}{(}\PYG{n}{mask}\PYG{p}{,} \PYG{o}{\PYGZhy{}}\PYG{l+m+mi}{1}\PYG{p}{),} \PYG{n+nb}{self}\PYG{o}{.}\PYG{n}{mask\PYGZus{}token}\PYG{o}{.}\PYG{n}{value}\PYG{p}{,} \PYG{n}{vid\PYGZus{}embed\PYGZus{}BTNM}
\PYG{p}{)}
\end{Verbatim} 
\caption{Code snippet from Jasmine showing its batched masking logic.} \label{fig:jasmine-batched-masking}
\end{figure*}

Whereas \citet{willi2024jafar} sample a single masking probability and apply the same masking pattern to all samples in a batch (Figure 
ef{fig:jafar-batched-masking}), Jasmine samples 	exttt{batch\_size} many sampling probabilities and uses per-sequence masking patterns (Figure 
ef{fig:jasmine-batched-masking}). This leads to significantly reduced loss variance (Figure 
ef{fig:wc_time}), especially in highly distributed settings.

\section{Model Inspection using Treescope}
%TODO: switch out for camera-ready
Jasmine's training loop is highly modular and supports easy model surgery as well as model inspection using Treescope \citep{johnson2024penzai}. We provide a demo notebook\footnote{\url{https://colab.research.google.com/drive/1zHkciFIZxXloJgue9F5LtFlA0m00rJIf}} alongside our repository, which illustrates debugging a common training instability (\Cref{fig:demo-notebook}).
%Jasmine's training loop is highly modular and supports easy model surgery as well as model inspection using Treescope \citep{johnson2024penzai}. We provide a demo notebook\footnote{\url{https://colab.research.google.com/drive/[RETRACTED]}} alongside our repository, which illustrates debugging a common training instability (\Cref{fig:demo-notebook}).


\begin{figure}
  \centering
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/samples/cropped/gt_jafar_prepend_100.png}
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/samples/cropped/gt_jafar_prepend_101.png}
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/samples/cropped/gt_jafar_prepend_102.png}
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/samples/cropped/gt_jafar_prepend_103.png}
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/samples/cropped/gt_jafar_prepend_104.png}
  \caption{Autoregressive rollouts on five randomly selected trajectories from the CoinRun environment. Each set of three rows corresponds to one trajectory, showing ground-truth frames (top), Jafar samples (middle), and Jasmine samples (bottom). The four conditioning frames are omitted.}
  \label{fig:sampling_multi}
\end{figure}

\begin{figure}
  \centering
  \includegraphics[width=	extwidth, trim=1 0 0 0, clip]{fig/diffusion_samples.png}
  \caption{Autoregressive rollouts on five randomly selected trajectories from the CoinRun environment. Each set of three rows corresponds to one trajectory, showing ground-truth frames (top), samples using the diffusion baseline (middle), and samples using the MaskGIT baseline (bottom). The four conditioning frames are omitted.}
  \label{fig:sampling_diffusion}
\end{figure}

\begin{figure}
  \centering
  \includegraphics[width=	extwidth]{fig/model-surgery.png}
  \caption{Treescope visualization when performing model inspection. The notebook illustrates inspecting output logits at specific checkpoints.}
  \label{fig:demo-notebook}
\end{figure}
\end{document}

\end{document}