SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 13: Experiment Configuration for ConvLSTM on Moving-MNIST experiments
Hyperparameters	Moving-MNIST-300	Moving-MNIST-600
V100 Days	75	150
Params	20M	20M
Input Resolution	64 × 64	64 × 64
Latent Resolution	16 × 16	16 × 16
Batch Size	8	8
Sequence Length	300	600
LR	5 × 10-4	5 × 10-4
LR Schedule	cosine	cosine
Warmup Steps	5k	5k
Max Training Steps	300K	300K
Weight Decay	1 × 10-5	1 × 10-5
Encoder	Depths	64, 128, 256	64, 128, 256
Blocks	1	1
Decoder	Depths	64, 128, 256	64, 128, 256
Blocks	1	1
ConvLSTM	Hidden Dim	256	256
State Size	256	256
Kernel Size	3 × 3	3 × 3
Layers	8	8
Dropout	0	0
Table 14: Experiment Configuration for Transformer on Moving-MNIST experiments
Hyperparameters	Moving-MNIST-300	Moving-MNIST-600
V100 Days	25	50
Params	164M	164M
Input Resolution	64 × 64	64 × 64
Latent Resolution	8 × 8	8 × 8
Batch Size	8	8
Sequence Length	300	600
LR	5 × 10-4	1 × 10-4
LR Schedule	cosine	cosine
Warmup Steps	5k	5k
Max Training Steps	300K	300K
Weight Decay	1 × 10-5	1 × 10-5
Encoder	Depths	64, 128, 256, 512	64, 128, 256, 512
Blocks	1	1
Decoder	Depths	64, 128, 256, 512	64, 128, 256, 512
Blocks	1	1
Downsample Factor	8	8
Hidden Dim	1024	1024
Temporal Transformer	Feedforward Dim	4096	4096
Heads	16	16
Layers	8	8
Dropout	0	0
D.5 Long-Range 3D Environment Benchmarks
We follow the procedures from Yan et al. [13] and train models on the same pre-trained vector-quantized (VQ)  $16 \times 16$  codes used by the baselines evaluated in that work. Models were trained to optimize a cross-entropy reconstruction loss between the predictions and true VQ codes. The evaluation of DMLab and Habitat involves both an action-conditioned and unconditioned setting. Therefore, as in Yan et al. [13], the actions were randomly dropped out half the time during training on these datasets.
After training, we follow the procedure from Yan et al. [13] for evaluation in two different settings. The first setting involves computing PSNR, SSIM and LPIPS from 1024 samples generated by conditioning on the first 144 frames and then generating the next 156 frames while providing the model with past and future actions. The second setting does not provide actions as input (with the exception of Minecraft, which also provides actions in this setting). It involves computing FVD using 1024 samples generated by conditioning on the first 36 frames and then predicting the remaining 264 frames.
All sequence models we trained used the same number of layers as the Transformer used in the TECO-Transformer trained by Yan et al. [13]. In addition, the TECO-Transformer, TECO-S5 and TECO-ConvS5 models we trained used the exact encoder/decoder configuration and MaskGit configuration as in Yan et al. [13]. The Transformer, S5 and ConvS5 models we trained without the TECO framework were all trained using the same encoder/decoder configuration. See Tables 15-21 for more detailed experimental configuration details. See dataset-specific paragraphs below for hyperparameter tuning information.
TECO Training Framework Yan et al. [13] proposed the TECO training framework to train Transformers on long video data. For some of our experiments, we use ConvS5 layers and S5 layers as a drop-in replacement for the Transformer in this framework. We refer the reader to Yan et al. [13] for full details. Briefly, given the original VQ codes, TECO trains an additional encoder/decoder that compresses the frames to a lower latent resolution (e.g., from  $16 \times 16$  to  $8 \times 8$ ) by training an additional encoder/decoder with a codebook loss,  $\mathcal{L}_{\mathrm{VQ}}$ . In addition, a MaskGit [98] dynamics prior loss,  $\mathcal{L}_{\mathrm{prior}}$ , is used for the latent transitions. The sequence model (e.g. Transformer, S5, ConvS5) takes the latent frames (compressed into tokens in the case of Transformer and S5) and produces an output which is used along with the latents by the decoder to produce predictions and a reconstruction loss,  $\mathcal{L}_{\mathrm{recon}}$ . Models are trained to minimize the following total loss:
$$
\mathcal {L} _ {\mathrm {T E C O}} = \mathcal {L} _ {\mathrm {V Q}} + \mathcal {L} _ {\text {r e c o n}} + \mathcal {L} _ {\text {p r i o r}}. \tag {32}
$$
In addition, TECO includes the use of DropLoss [13], which drops out a percentage of random timesteps that are not decoded and therefore do not require computing the expensive  $\mathcal{L}_{\mathrm{recon}}$  and  $\mathcal{L}_{\mathrm{prior}}$  terms.
DMLab As mentioned above, the actions were randomly dropped out of sequences half the time (due to the two evaluation scenarios, action-conditioned and unconditioned). We observed that for DMLab, when provided past and future actions, models converged faster using the simple masking strategy discussed in Gu et al. [19] that masks the future inputs rather than feeding the predicted inputs (or true inputs during training) autoregressively. Therefore we trained all models (Transformer, Performer, S5, ConvS5, Teco-Transformer, TECO-S5, TECO-ConvS5) by using this strategy when the actions were provided, and using the autoregressive strategy when actions were not provided. We observed this significantly improved the LPIPS of the Transformer baselines. Note, in pilot runs for Minecraft and Habitat, we observed this strategy led to lower-quality frames and did not use it for the reported results for those datasets.
We trained each model, Transformer, Performer, S5, ConvS5, TecO-Transformer, TECO-S5, TECOConvS5, with three different learning rates  $[1 \times 10^{-4}, 5 \times 10^{-4}, 1 \times 10^{-3}]$  and selected the best run for each model. See Tables 15-20 for more experiment configuration details.
The TECO-ConvSSM ablation used the exact same settings as TECO-ConvS5, except the state kernel was initialized with a random Gaussian and a lower learning rate of  $1 \times 10^{-5}$  was required for stable training.
Table 15: Experiment Configuration for ConvS5 on DMLab
Hyperparameters	DMLab
V100 Days	150
Params	101M
Input Resolution	64 × 64
Latent Resolution	16 × 16
Batch Size	16
Sequence Length	300
LR	5 × 10-4
LR Schedule	cosine
Warmup Steps	5k
Max Training Steps	500K
Weight Decay	1 × 10-5
Encoder	Depths	256
Blocks	1
Decoder	Depths	256
Blocks	4
ConvS5	Hidden Dim (U)	512
State Size (P)	512
B Kernel Size	3 × 3
C Kernel Size	3 × 3
Layers	8
Dropout	0
Activation	ResNet
Table 16: Experiment Configuration for S5 on DMLab
Hyperparameters	DMLab
V100 Days	125
Params	140M
Input Resolution	64 × 64
Latent Resolution	16 × 16
Batch Size	16
Sequence Length	300
LR	1 × 10-3
LR Schedule	cosine
Warmup Steps	5k
Max Training Steps	500K
Weight Decay	1 × 10-5
Encoder	Depths	256
Blocks	1
Decoder	Depths	256
Blocks	4
S5	Downsample Factor	16
Hidden Dim (U)	1024
State Size (P)	1024
Layers	8
Dropout	0
Activation	GLU (half)
Table 17: Experiment Configuration for Transformer on DMLab
Hyperparameters	DMLab
V100 Days	125
Params	152M
Input Resolution	64 × 64
Latent Resolution	16 × 16
Batch Size	16
Sequence Length	300
LR	5 × 10-4
LR Schedule	cosine
Warmup Steps	5k
Max Training Steps	500K
Weight Decay	1 × 10-5
Encoder	Depths	256
Blocks	1
Decoder	Depths	256
Blocks	4
Downsample Factor	16
Hidden Dim	512
Temporal Transformer	Feedforward Dim	2048
Heads	16
Layers	8
Dropout	0
Table 18: Experiment Configuration for TECO-ConvS5 on DMLab
Hyperparameters	DMLab
V100 Days	110
Params	175M
Input Resolution	64 × 64
Latent Resolution	8 × 8
Batch Size	16
Sequence Length	300
LR	5 × 10-4
LR Schedule	cosine
Warmup Steps	5k
Max Training Steps	500K
Weight Decay	1 × 10-5
DropLoss Rate	0.9
Encoder	Depths	256, 512
Blocks	2
Codebook	Size	1024
Embedding Dim	32
Decoder	Depths	256, 512
Blocks	4
ConvS5	Hidden Dim (U)	512
State Size (P)	1024
B Kernel Size	3 × 3
C Kernel Size	3 × 3
Layers	8
Dropout	0
Activation	ResNet
MaskGit	Mask Schedule	cosine
Hidden Dim	512
Feedforward Dim	2048
Heads	8
Layers	8
Dropout	0
Table 19: Experiment Configuration for TECO-S5 on DMLab
Hyperparameters	DMLab
V100 Days	80
Params	180M
Input Resolution	64 × 64
Latent Resolution	8 × 8
Batch Size	16
Sequence Length	300
LR	1 × 10-3
LR Schedule	cosine
Warmup Steps	5k
Max Training Steps	500K
Weight Decay	1 × 10-5
DropLoss Rate	0.9
Encoder	Depths	256, 512
Blocks	2
Codebook	Size	1024
Embedding Dim	32
Decoder	Depths	256, 512
Blocks	4
S5	Downsample Factor	8
Hidden Dim (U)	2048
State Size (P)	2048
Layers	8
Dropout	0
Activation	GLU (half)
MaskGit	Mask Schedule	cosine
Hidden Dim	512
Feedforward Dim	2048
Heads	8
Layers	8
Dropout	0
Table 20: Experiment Configuration for TECO-Transformer on DMLab
Hyperparameters	DMLab
V100 Days	80
Params	173M
Input Resolution	64 × 64
Latent Resolution	8 × 8
Batch Size	16
Sequence Length	300
LR	1 × 10-4
LR Schedule	cosine
Warmup Steps	5k
Max Training Steps	500K
Weight Decay	1 × 10-5
DropLoss Rate	0.9
Encoder	Depths	256, 512
Blocks	2
Codebook	Size	1024
Embedding Dim	32
Decoder	Depths	256, 512
Blocks	4
Temporal Transformer	Downsample Factor	8
Hidden Dim	1024
Feedforward Dim	4096
Heads	16
Layers	8
Dropout	0
MaskGit	Mask Schedule	cosine
Hidden Dim	512
Feedforward Dim	2048
Heads	8
Layers	8
Dropout	0
Minecraft and Habitat For Minecraft and Habitat, we only trained TECO-ConvS5 due to the costs of training on these datasets. See dataset details in Appendix E and reported compute costs in Yan et al. [13]. For Minecraft, we evaluated two different learning rates  $[1 \times 10^{-4}, 5 \times 10^{-4}]$  and chose the best. For Habitat, we only performed one run with no further tuning. See Table 21 for further experiment configuration details.
Table 21: Experiment Configuration for TECO-ConvS5 on Minecraft and Habitat
Hyperparameters	Minecraft	Habitat
V100 Days	470	575
Params	214M	351M
Input Resolution	128 × 128	128 × 128
Latent Resolution	8 × 8	8 × 8
Batch Size	16	16
Sequence Length	300	300
LR	5 × 10-4	1 × 10-4
LR Schedule	cosine	cosine
Warmup Steps	5k	5k
Max Training Steps	1M	1M
DropLoss Rate	0.9	0.9
Encoder	Depths	256, 512	256, 512
Blocks	4	4
Codebook	Size	1024	1024
Embedding Dim	32	32
Decoder	Depths	256, 512	256, 512
Blocks	8	8
ConvS5	Hidden Dim (U)	512	512
State Size (P)	512	512
B Kernel Size	3 × 3	3 × 3
C Kernel Size	3 × 3	3 × 3
Layers	12	8
Dropout	0	0
Activation	ResNet	ResNet
MaskGit	Mask Schedule	cosine	cosine
Hidden Dim	768	1024
Feedforward Dim	3072	4096
Heads	12	16
Layers	6	16
Dropout	0	0
E Datasets
E.1 Moving-MNIST
The Moving-MNIST [54] dataset is generated by moving two  $28 \times 28$  size MNIST digits from the MNIST dataset [111] inside a  $64 \times 64$  black background. The digits begin at a random initial location, and move with constant velocity, bouncing when they reach the boundary. For each of the sequence lengths we consider, 300 and 600, we follow Wang et al. [81] and Su et al. [84] and generate 10,000 sequences for training.
E.2 DMLab
We use the DMLab long-range benchmark designed by Yan et al. [13] using the DeepMind Lab (DMLab) [99] simulator. The simulator generates random 3D mazes with random floor and wall textures. The benchmark consists of  $40\mathrm{K}$  action-conditioned, 300 frame videos at a  $64\times 64$  resolution. The videos are of an agent randomly navigating  $7\times 7$  mazes by choosing random points in the maze and navigating to them through the shortest path.
E.3 Minecraft
We use the Minecraft [100] long-range benchmark designed by Yan et al. [13]. The game features 3D worlds that contain complex terrains such as hills, forests, rivers and lakes. The benchmark was constructed by collecting 200K action-conditioned 300 frame videos at a  $128 \times 128$  resolution. The videos are in Minecraft's marsh biome and the agent iterates walking forward for a random number of steps and randomly rotating left or right. This results in parts of the scene going out of view and coming back into view later.
E.4 Habitat
We use the Habitat long-range benchmark designed by Yan et al. [13] using the Habitat simulator [101]. The simulator renders trajectories using scans of real 3D scenes. Yan et al. [13] compiled 1400 indoor scans from HM3D [112], Matterport3D [113] and Gibson [114] to generate 200K action-conditioned, 300 frame videos with a  $128 \times 128$  resolution. Yan et al. [13] used Habitat's in-built path traversal algorithm to construct action trajectories that move the agent between randomly sampled locations.