PaperlayerView exact revision on arXiv

Source-linked AI summary

Communication-Inspired Tokenization for Structured Image Representations

Aram Davtyan, Yusuf Sahin, Yasaman Haghighi, Sebastian Stapf, Pablo Acuaviva, Alexandre Alahi, Paolo Favaro

arXiv:2602.20731v1cs.CVcs.AIcs.LG

TL;DR

Existing discrete image tokenizers often produce texture-focused, poorly structured tokens because they emphasize reconstruction and compression. COMiT learns a structured discrete message through attentive sequential observations and unified flow-matching encoding and decoding, with semantic alignment. The paper reports improved semantic, compositional, and relational evaluations, while noting that reconstruction fidelity remains an area for improvement.

  • Problem

    Existing discrete encoders often prioritize reconstruction and compression, yielding tokens that capture local texture instead of object-level semantic structure.

  • Method

    COMiT iteratively observes localized image crops, updates a discrete latent message, and trains unified encoding and flow-matching decoding with semantic representation alignment.

  • Results

    COMiT consistently outperforms prior one-dimensional discrete image encoders across benchmarks of semantic grounding, compositional generalization, and relational reasoning.

  • Takeaways & Limitations

    Attentive sequential tokenization organizes semantic information into interpretable, object-centric tokens, while semantic alignment provides grounding.

  • Takeaways & Limitations

    Reconstruction fidelity could be further improved, particularly for generation-focused applications.

Abstract

from arXiv · show

Discrete image tokenizers have emerged as a key component of modern vision and multimodal systems, providing a sequential interface for transformer-based architectures. However, most existing approaches remain primarily optimized for reconstruction and compression, often yielding tokens that capture local texture rather than object-level semantic structure. Inspired by the incremental and compositional nature of human communication, we introduce COMmunication inspired Tokenization (COMiT), a framework for learning structured discrete visual token sequences. COMiT constructs a latent message within a fixed token budget by iteratively observing localized image crops and recurrently updating its discrete representation. At each step, the model integrates new visual information while refining and reorganizing the existing token sequence. After several encoding iterations, the final message conditions a flow-matching decoder that reconstructs the full image. Both encoding and decoding are implemented within a single transformer model and trained end-to-end using a combination of flow-matching reconstruction and semantic representation alignment losses. Our experiments demonstrate that while semantic alignment provides grounding, attentive sequential tokenization is critical for inducing interpretable, object-centric token structure and substantially improving compositional generalization and relational reasoning over prior methods.

1. Introduction

Existing discrete image tokenizers often prioritize reconstruction and compression, producing local, poorly localized token information rather than object-level structure. COMiT addresses this gap by iteratively integrating localized observations into a unified discrete message, while experiments show that sequential tokenization is central to interpretable and compositional representations.

  • Conventional discrete encoders optimize reconstruction under compression constraints, so their tokens often capture local texture rather than object-level semantic structure.
  • One-dimensional tokenizers better match transformer sequence formats, but semantic information can remain entangled and poorly localized across tokens.
  • The method uses attentive sequential observations and a unified speaker-listener network to integrate visual information into the message.The encoder attends to different localized regions at each step, while the same network performs encoding and reconstruction.
  • COMiT frames image encoding as an iterative communication-and-reconstruction game that updates a discrete latent message from successive image crops.
  • Across benchmarks for semantic content, compositional generalization, and relational reasoning, COMiT substantially outperforms existing one-dimensional discrete image encoders.Ablations indicate complementary effects: semantic alignment improves token meaning, while sequential encoding induces more interpretable, object-centric tokens.

2. Related Work

Prior work spans recurrent attention, discrete codebooks, and flow-matching autoencoders. COMiT builds on these directions while using a unified flow objective for both encoding and decoding.

  • Attentive Encoding: Recurrent attention models iteratively select image regions and accumulate latent state, but early approaches targeted efficiency, reconstruction, or explicit inductive biases rather than abstract representation learning.
  • Image Tokenization: Image tokenization progressed from vector-quantized codebooks to masked-prediction, adversarial, and one-dimensional schemes designed for discrete visual representations.
  • Flow Matching and Diffusion Autoencoders: Flow-matching models provide stable, efficient training for continuous distributions and have also been used to train autoencoder decoders.
  • Flow Matching and Diffusion Autoencoders: COMiT is presented as the first approach to train both encoding and decoding through a unified differentiable flow objective in a single network.

3. Communication Inspired Tokenization

COMiT iteratively encodes localized image crops into a quantized latent message and uses that message to reconstruct the full image with the same transformer. Its training combines flow-matching reconstruction with semantic alignment, while randomized crop counts encourage greedy token usage.

  • Encoding: Relative crop offsets encode movements between crop locations, helping the model scan scenes without specializing message parts to constant regions.The first offset is zero; later offsets are differences between consecutive crop centers.
  • Encoding: COMiT initializes a latent message and recurrently updates it with each observed crop before quantizing the predicted message through FSQ.The quantized message is fed back with the next crop, producing the final message after K encoding steps.
  • Decoding: The final message conditions flow-based decoding of the full image, and COMiT reuses one model for both encoding and decoding.Decoding begins from a noisy interpolant and predicts the velocity used for flow matching; inference numerically integrates the flow ODE from noise.
  • Greedy Use of Tokens: Randomizing the number of aggregated crops prevents fixed capacity allocation and encourages the model to use available tokens greedily.The model cannot know whether additional crops will arrive, so it must retain essential information as the message grows.
  • Greedy Use of Tokens: Training applies stop-gradient to all message updates except the final one, reducing the computational and memory cost of backpropagating through the full sequence.The authors report only a moderate impact on preliminary small-scale experiments.
  • Semantic Alignment: SREPA distills pretrained DINOv2 [CLS] features into intermediate message-token representations using projected, pooled features and cosine similarity.The semantic alignment objective is applied during decoding to encourage semantic ordering in the learned messages.
  • Reconstruction Fidelity: Training mixes global-crop, local-crop, and empty-message decoding regimes, with pCFG ≈ 0.2 and pG ≈ 0.5 used to improve reconstruction fidelity and enable classifier-free guidance.The global crop is the full image, while local crops are smaller observations.
  • Cropping Policies: Randomized crops and crop counts permit flexible inference policies, including choices about global-crop use, crop number, and crop-processing order.The paper evaluates random, raster-scan, and adaptive ordering on 96 × 96 crops arranged over a 3 × 3 grid.

4. Experiments

The experiments evaluate COMiT’s semantic encodings, cropping policies, and emergent compositional behavior through recognition, generalization, relational, reconstruction, and grounding analyses. Across these evaluations, attentive sequential encoding and semantic alignment contribute complementary benefits, while global-only cropping offers a strong cost–performance trade-off.

  • Evaluation benchmarks: The test suite measures visual recognition, compositional generalization, and inter-object relational semantics.Recognition uses ImageNet100, compositional generalization uses disjoint object co-occurrence splits from MSCOCO, and relational semantics uses relationship prediction on Visual Genome.
  • Ablations: SREPA significantly contributes to tokenizer performance, while the ablation keeps the test-time cropping policy fixed to a single global crop.The comparison removes the SREPA loss during training and evaluates both models under identical settings on IN100.
  • Ablations: Attentive training improves recognition probing even when inference encodes only the global crop.Without local crops, the model lacks recurrent message refinement and effectively becomes a standard discrete encoder–decoder.
  • Ablations: 0.53 mean mIoU versus 0.34 shows tighter object alignment for COMiT than for the model trained without local crops.The comparison uses attention maps from the 10th layer and the same global-only cropping policy at inference.
  • Cropping policies: A single global crop provides the best overall performance–cost trade-off, while additional local crops offer modest gains on compositional generalization and inter-object relations.These cropping-policy differences occur entirely at test time without retraining, and global-only cropping is adopted by default for probing experiments.
  • Qualitative results: Sequential crop aggregation progressively reduces reconstruction ambiguity and adds observed objects to the latent message.A global crop quickly embeds the scene, while subsequent local crops focus on finer details; single-step decoding is sharp in observed regions and blurry elsewhere.

5. Conclusion and Discussion

COMiT frames one-dimensional image tokenization as iterative communication, producing structured visual representations and outperforming prior work on semantic probing. The paper identifies sequential attention as central to compositional token structure, while noting reconstruction fidelity remains improvable.

  • Contributions: COMiT constructs a discrete latent message through attentive sequential observations, organizing semantic information incrementally across tokens.The message is formed through an iterative communication-and-reconstruction process rather than a single encoding pass.
  • Results: COMiT outperforms prior 1D tokenizers on semantic probing across benchmarks covering semantic grounding, compositional generalization, and relational reasoning.The evaluation suite is designed to test these three properties in learned token sequences.
  • Results: Ablations on attention organization indicate that attentive tokenization contributes to compositional token structure aligned with meaningful scene regions.Qualitative analysis also shows progressive refinement of the latent message over encoding steps.
  • Limitations: Reconstruction fidelity could be further improved, which may matter for generation-focused applications.The paper’s primary focus is semantic organization rather than reconstruction fidelity.
  • Implications: COMiT supports flexible test-time cropping policies, opening directions for adaptive and task-dependent visual tokenization.The paper specifically mentions reinforcement learning as one possible future direction for selecting cropping policies.
  • Future work: Future work extends COMiT toward video, where temporal redundancy and long-range structure create additional challenges for discrete representation learning.The proposed direction would let latent messages accumulate motion- and action-related information over time.

Impact Statement

The paper’s methods may carry societal and ethical risks associated with generated content, training-data bias, and large-scale model costs.

  • Risks: Potential risks include misuse of generated content for misleading media, amplification of training-data biases, and environmental costs from training and deployment.These risks are presented as standard concerns for generative image models and tokenizers.

A. Additional Implementation Details

The appendix documents COMiT’s architecture, training setup, cropping and timestamp sampling, message handling, position encodings, and guidance configuration.

  • Architecture: COMiT is based on DiT, with separate AdaLN modulation layers for image and message tokens.Other architectural details, including transformer depth and dimensions, are provided in the architecture table.
  • Training: All models use ImageNet1k, Adam optimization, a 3e-4 base learning rate, warmup, and 200 training epochs on GH200 GPUs.The B and L variants use batch size 512, while XL uses batch size 256 with two accumulation steps for the same effective batch size.
  • Cropping policy: Training uses 256 × 256 images and random 96 × 96 local crops, with the crop count sampled from 1 through Kmax.The sampling distribution is skewed toward one crop rather than uniform across crop counts.
  • Flow matching: Flow-matching timestamps are sampled with a distribution skewed toward zero to prioritize high-noise regimes.The appendix contrasts this choice with the default uniform distribution on [0, 1].
  • Message handling: Message tokens are initialized from one fixed vocabulary token, embedded with a two-layer SiLU MLP, and updated using an appended buffer token.The buffer aggregates information from the existing message and the newly observed crop.
  • Position encodings: Each modality uses separate positional encodings: learned encodings for message and offset tokens and 2D sine–cosine encodings for image tokens.Each crop uses the top-left corner of the image position encodings.
  • Guidance: Classifier-free guidance uses Adaptive Projected Guidance with r = 2.5, η = 0, and β = −0.5.These settings are intended to smooth metric dependence on guidance strength and permit higher guidance values.

B.1. Test Suite

The test suite probes semantic content in frozen image-token sequences using lightweight, standardized Transformer probes across category, compositional, and relational tasks.

  • Evaluation protocol: All benchmarks use frozen image tokens and identical-capacity lightweight probes, with early stopping based on validation performance.This protocol is designed to compare the semantic information encoded by different tokenizers fairly.
  • ImageNet100: ImageNet100 measures category-level semantic information with single-label classification and validation top-1 accuracy.The probe predicts one of 100 object categories using cross-entropy.
  • MSCOCO: MSCOCO measures compositional generalization using pair-disjoint two-object images and multi-label classification.Validation object-category pairs are unseen during training, and performance is reported with top-5 accuracy.
  • Visual Genome: Visual Genome evaluates relational reasoning using image-disjoint subject–predicate–object triplets.Relations are filtered by object area and restricted to the 150 most frequent categories and predicates.
  • Probe architecture: Probes aggregate projected image tokens with a shallow Transformer and use a learnable [CLS] representation for prediction.ImageNet100 and MSCOCO use only [CLS] and image tokens, while Visual Genome additionally conditions on encoded triplet text.
  • Ablations: Reconstruction-fidelity ablations vary sampling hyperparameters and COMiT-B bottleneck size.The figure identifies these as separate ablation dimensions.

B.4. Reconstruction.

COMiT’s reconstruction quality is evaluated through sampling choices, scaling, bottleneck size, and emergent object-centric attention. Increasing model size and bottleneck capacity improves reconstruction, while semantic probing remains stable across bottleneck sizes.

  • Sampling and scaling: NFE = 10 is selected as the optimal number of evaluation steps across all model sizes, with higher NFE producing sharper details but slightly higher rFID.The sweep fixes CFG = 3.0 and evaluates NFE in {2, 4, 10, 25}.
  • Sampling and scaling: Scaling COMiT from B to L significantly reduces rFID, while scaling from L to XL yields only a moderate reconstruction improvement.The reported reconstruction setting is NFE = 10 and CFG = 7.5.
  • Bottleneck size: Bottleneck variants with larger message lengths and vocabularies maintain similar classification probing accuracy while increasing reconstruction PSNR.The tested message lengths are 64, 128, and 256 tokens, with vocabulary sizes of 1k, 16k, and 64k.

C. Analysis of the Attention Maps

COMiT’s decoding attention maps reveal object-centric structure in its discrete message tokens. Quantitative overlap with object regions is strong despite the model receiving no segmentation or category supervision.

  • Emergence of Objectness: COMiT tokens tend to attend to semantically meaningful image regions, including objects and object parts.The maps are examined in layer 24 of COMiT-XL during decoding.
  • Quantifying Objectness: 0.58 mIoU is obtained at 30% attention-map thresholding on CSSD without segmentation maps or object categories during training.Attention maps are extracted from layer 24 at denoising time t = 0.1 using COMiT-XL.

D. More Visual Results

Additional visual analyses show that COMiT’s latent messages support semantic neighborhoods, domain-shift robustness, adaptive information selection, and object-related attention. Reconstruction improves as observed crops accumulate, but unobserved content is omitted.

  • Domain generalization: COMiT remains robust under domain shifts despite being trained only on ImageNet1k.The qualitative results include rendered animations and medical images.
  • Latent-message structure: Nearest neighbors in COMiT’s latent-message space typically share semantics, indicating a well-structured representation under simple cosine-distance probing.Each image is represented by concatenated tokens forming a 1536-dimensional vector, and four nearest samples are retrieved.
  • Cropping policies: Three to four local adaptive crops are typically sufficient to capture a scene’s overall idea without including excessive detail.This behavior explains why the adaptive policy with only three crops performs well on semantic probing.
  • Cropping policies: Content absent from the crop sequence is missing from the final reconstruction, while intermediate reconstructions show iterative refinement as new information arrives.Figure 10 compares adaptive and global+adaptive cropping policies and uses 1 NFE intermediate reconstructions and 10 NFE final reconstructions.
  • Object-centric structure: COMiT-XL attention maps can achieve a 0.58 mIoU after threshold adjustment without segmentation or class-label supervision.The visualization reports the largest-IoU token attention maps.
Loading 2602.20731v1…