ElasticTok: Adaptive Tokenization for Image and Video

Rebuttal

Thank you for your constructive feedback. We provide detailed answers to the questions below. --- > **Choosing how to adaptively encode, or generate tokens** As mentioned by the reviewer, a primary method of adaptively selecting the number of tokens is encoding image or video according to a pre-defined threshold. Downstream generative models can still generate a variable number of tokens by conditioning on a threshold, or having no conditioning but training on a fixed threshold. For example, generating a video at a target MSE of 0.003 will have a model dynamically generate a different number of tokens per video block depending on the underlying dynamics and complexity of the generated video. This would require running ElasticTok inference at 0.003 threshold, and training on the resulting encodings. --- > **First-k tokens masking limitation** Although the left-aligned masks are simple, we believe that our experiments (e.g. Fig 4) show the effectiveness of our method on general image and video. Both training and evaluation distributions span a diverse range of visual content from the web – cartoons, realistic movies, fast transitions, scene cuts, etc. The current approach strikes a practical balance between computational overhead and effectiveness. However, we do agree that better masking schemes (e.g., content-adaptive masks or attention-guided token selection) could further boost performance, which we plan to explore in future work. --- > **For Fig.4 comparisons, do the authors first search for the ideal token length for each input for their method?** Yes, each image / video is adaptively encoded to achieve the target MSE. This allows the model to potentially allocate more tokens to harder to encode images while still retaining lower average token usage. In contrast, baseline models must always encode with the same number of tokens, which may waste tokens on easy to encode images, or be too lossy on difficult images. --- > **Do both methods use 50% of tokens during inference in Table 3 experiments?** The reviewer asked for clarification about inference token usage for both methods in Table 3. The baseline is trained with 100% of tokens, and ElasticTok is trained with a varying number of tokens. Inference in Table 3 is comparing the baseline and ElasticTok both using 100% of tokens during inference, demonstrating that ElasticTok can match baseline performance when using full context, while maintaining the additional advantage of allowing users to choose token-accuracy tradeoffs during inference if desired (e.g. able to perform more API queries if choosing to use fewer tokens).

Rebuttal

Thank you for your constructive feedback. We provide detailed answers to the questions below. --- > **Total inference time of each method** Table 4 shows the inference time in terms of NFEs, or the number of forward passes needed for the model, and the exact model used. However, the total FLOPs can be directly computed as the model FLOPs for a single forward pass times the NFEs. For example, the FLOPs for binary search on the continuous ElasticTok model is approximately (819 GFLOPs * 12 NFEs = 9.8 TFLOPs). --- > **Claims on efficient computation** We would like to clarify that the tokenizer itself does not save computation, as the encoder and decoder always operate on the full sequence length (the decoder just sees a masked / padded input). **However**, our experiments show that adaptive encoding (Fig 4) results in large savings on **token lengths of resulting encoding sizes**, which will save 2-5x tokens for downstream pretrained VLMs. This **directly translates to** at least 2-5x (more as sequence length grows due to quadratic attention) speedup in training and inference. This is especially valuable as pretrained downstream models are very large (10Bs - 100Bs params with at least ~2-5x speedup), and the extra computation from the tokenizer (200M params) is less significant and can be pre-encoded. --- > **Search method used for each evaluation** Thank you for this clarification. Unless otherwise noted, we use **binary search** as the search algorithm for all results. We have edited the paper (L239) to accurately reflect this. --- > **Details on the regressor** We have added more details in Appendix E on how the regressor is finetuned. The encoder is frozen. A small number of extra trainable transformer layers are added, along with progressive downsampling to eventually output a single scalar, which is the predicted number of tokens. The targets are normalized to [-1, 1], and regress to using MSE. The regression model is also block causal, and does see the entire historical context of the video. Motion-representations would be an accurate source of information with respect to slow / fast changes in the video, however do not generalize well to all cases, such as abrupt scene cuts (very common in videos), and frames which are just dense in visual information even if there is little motion (a very cluttered room, or a video of just static noise). Hence, we decided to use pixel representations of videos in our regression models as it contains complete information about any relevant video statistics. --- > **Citation suggestion** Thank you for this reference. We have added it as further discussion in our related works (L472-473).

Followup on Rebuttal

Dear Reviewer xRJN, Thank you for taking the time to review our paper. We appreciate your feedback and believe we have addressed the reviewer’s concerns regarding efficiency, search methods, and regressor details in our rebuttal. We would greatly appreciate your consideration of our response. Best regards, Authors

Rebuttal

Thank you for your constructive feedback. We provide detailed answers to the questions below. --- > **Lack of comparisons** Our paper focuses on developing an autoencoder-based adaptive tokenizer that can be used for downstream any-to-any multi-modal models (e.g. Chameleon [1], Omnigen [2]). Thus, we focus on comparing against SoTA baseline tokenizers that these models use – generally consisting of fixed-length autoencoders that learn discrete (VQ, FSQ), or continuous (VAE) latents using either CNNs or transformers. Since our only difference with prior SoTA tokenizers is solely fixed vs elastic (variable) length encodings, we design baseline comparison experiments that retain aspects of standard tokenizer architectures (ViT architecture, MSE + LPIPs losses) and only alter the model to learn elastic encodings. We believe that this is the most fair and direct comparison with SoTA tokenization methods, and our comparable token usage efficiency demonstrates the benefits of our method. [1] Meta AI. "Chameleon: Mixed-modal early-fusion foundation models." arXiv preprint arXiv:2405.09818 (2024). [2] Xiao, S., Wang, Y., Zhou, J., Yuan, H., Xing, X., Yan, R., ... & Liu, Z. (2024). Omnigen: Unified image generation. arXiv preprint arXiv:2409.11340. --- > **Reconstruction quality when using elastic token compression** Fig 4 shows the reconstruction quality performance of our model – ElasticTok achieves better reconstruction of images/videos (by MSE/PSNR) compared to the baseline while using fewer percentages of tokens. Note in general some reconstructions may be a little blurry since we primarily used MSE as the main optimization loss. Our proposed method is orthogonal to additional perceptual objectives such as GAN losses, which have shown to drastically improve image sharpness and can be additionally applied. --- > **Applications of elastic token compression, and compared to other compression methods** To the best of our knowledge, all existing any-to-any multi-modal models use fixed length tokenizers. This is generally undesirable, as allocating a fixed number of tokens for all images / videos will result in wasted tokens for simple images, and lossy encodings for complex images. With ElasticTok, we believe that a paradigm shift towards elastic (variable) token lengths can more efficiently allocate computational costs for downstream training and inference, which is generally by far the more expensive part in current generative models. Our model demonstrates up to 2-5x savings (Figure 4) in encodings lengths for images and long videos, which directly translates to at least 2-5x savings in training and inference costs of current generative models that traditionally used fixed length tokenizers.

Rebuttal (1/2)

Thank you for your constructive feedback. We provide detailed answers to the questions below. --- > **Paper presentation suggestions** We thank the reviewer for their thoughtful suggestions. We have addressed all points in our revised manuscript (changes highlighted in magenta): * Added RingAttention background: We have added a comprehensive background to explain RingAttention's core mechanisms and how these enable efficient processing of long sequences critical to our approach (Section2, L79-115) * Added more architectural details on the video model (Section 4.1, L215-220) * Added more details on the variables, and shapes in the algorithms with clear variable definitions, shape information, and step-by-step explanations (Algorithm 1 and 2, L187-204) * Clarified Tables 1 and 2 with more detailed discussion for each (Moved to Appendix F due to space constraints from prior changes, L850-886) --- > **Does the encoder take the masked samples as input?** Yes, the encoder does take the masked samples as input. We appreciate your careful attention to this technical detail. We have updated Figure 2 to explicitly show the masked samples flowing into the encoder block, making the architecture diagram more accurate and clearer. --- > **Masking strategy and positional embeddings** We use 1D positional embeddings over the entire sequence length, treating a patchified video as a flattened sequence of patches in raster scan order. We empirically found this simple approach effective at capturing spatiotemporal relationships, as evidenced by our model's strong performance across temporal tasks. Importantly, we **do not** left-mask the input video itself, but rather apply masking only to the resulting encoding. Thus, the encoder is given the whole unmasked video as well as the mask that **will be** applied so it can learn to put as much information as it can into the left-most tokens. This allows our encoder to process the complete unmasked video while being aware of which tokens will be masked, enabling it to optimize the information packing into the left-most tokens that will be preserved. The nature of the representations learned is hierarchical in a coarse-to-fine manner. This can be easiest seen in the progressive reconstructions (e.g. Figure 5, where we show video reconstructions using an increasing number of tokens from left to right, starting with the minimum token count). To further illustrate this coast-to-fine manner, we have added more detailed progressive reconstruction examples in Appendix G, along with specific token counts for each reconstruction stage in the bottom right of each image. --- > **Inference Time Compute Cost** Searching for masks at inference time indeed requires multiple function evaluations at video level, we would like to highlight a few points regarding how this overhead is both manageable and justified: * **Tokenizer vs Pretrained Model sizes**: Tokenizers are fairly lightweight compared to downstream vision-language models (200-500M tokenizers vs 10B-100B downstream VLMs). Thus, extra computation time of the tokenizer (1-10x depending on usable search algorithms) is less significant compared to running the downstream model. Additionally, the extra compute differences can also be made up through fewer flops due to token length reduction for the downstream model. * **Parallelizability**: Binned search methods such as 100 bins (0.5% error) or 10 bins (5% error), can be parallelized by just increasing batch size (compute different masking levels in parallel). In the best case, we can leverage more compute to have inference equivalent in speed to 1 forward pass / baseline models. * **Future work in optimized implementations**: While it is out of scope of this work, our inference method is similar to language models (block causal instead of standard causal), and can leverage many optimization techniques to speed up inference, such as quantization, paged attention / more efficient caching, more efficient kernels, etc. Overall, we do not believe that inference speed is a major bottleneck, and the benefits of adaptive tokenization outweigh the potential (addressable) drawbacks – ElasticTok enables more efficient token usage in downstream tasks, of which models are generally much larger, and more compute intensive.

Rebuttal (2/2)

--- > **Combinatorial masking combinations on long sequences** We appreciate the reviewer’s suggestions on making our combinatorial masking on long sequences more scalable and efficient. Our masking is indeed combinatorial, but it scales well to longer sequences. Our training on up to 1M tokens length sequences, the training is stable and works reasonably well on long sequences during inference. Figure 6 shows evidence of this, with monotonically increasing encoding performance on longer sequences up to 512K tokens, except the slight performance drop at 1M tokens likely due to insufficient gradient updates rather than a fundamental limitation. Intuitively, we are training on every masking task, which may help the model generalize to a wider range of unknown inference tasks, such as a more lossy encodings (higher masking amounts), less lossy encodings (lower masking amounts), or a mix (e.g. inference loss functions that only prioritize higher reconstruction quality on human faces, while lossily encoding everything else). We do agree that better masking schedules could improve the training efficiency of our model when seeking to apply our model for specific use case, perhaps through some heuristic method, or potentially bootstrapping off our current model. We believe that our current simple masking method demonstrates strong empirical results while maintaining implementation simplicity and represents a valuable contribution to the field, while leaving the door open for future optimizations. --- > **Video datasets release** Thank you for this question. Additional details about the video dataset construction, including data sources, sampling methodology, and preprocessing pipeline, are provided in Appendix C. We have eventual plans to publicly release the video dataset to support reproducibility, however, it is still in the process of being further filtered to exclude unwanted (e.g. NSFW) data (~1% of the current dataset) before release. --- > **Results (discrete vs continuous)** Both discrete and continuous models generally perform similarly, and experimental results show a mix of discrete and continuous models. We apologize for any confusion about model types in our results. The core results presented in Figure 4 were actually from the FSQ model, not the continuous model. We have now clarified this in the figure caption, and updated other table / figure captions in the paper. For comparison, Appendix F Table 3 presents parallel results from the continuous VAE model, demonstrating consistent performance across both architectures.

Question Regarding Masking Strategy and Baseline

Firstly, I appreciate the authors for putting up the rebuttal. > Importantly, we do not left-mask the input video itself, but rather apply masking only to the resulting encoding. Thus, the encoder is given the whole unmasked video as well as the mask that will be applied so it can learn to put as much information as it can into the left-most tokens. \ Initially, my understanding—due to the lack of masking details in the original submission—was that the encoder learns fixed-length embeddings, which are later masked. Therefore, as the authors mentioned in the related work, ElasticTok appeared to be an application of the Matryoshka approach to image/video reconstruction. However, based on the authors’ latest response, it seems the encoder is run repeatedly for different sampled masks at test time. This approach appears highly inefficient and results in wasted computational flops (doesn't seem a scalable approach for video going forward, if for each frame the whole encoder needs to be run again and again to find optimal mask). Why not adopt a single embedding, as in the Matryoshka approach, and perform masking directly on the encoded output? This would serve as an important baseline for comparison, and I would appreciate it if the authors could provide results for this approach.