ARM: Unifying Image Understanding, Generation, and Editing with Discrete Tokens and RL

Paper: 2606.11188 Authors: Junke Wang, Xiao Wang, Jiacheng Pan, Xuefeng Hu, Feng Li, Jingxiang Sun, Chaorui Deng, Zilong Chen, Yunpeng Chen, Kaibin Tian Categories: cs.CV

The Gap

Previous multimodal models treat image understanding (captioning, VQA) and image generation (text-to-image, editing) as separate problems with incompatible representations. Understanding models use continuous features from CLIP or DINO; generation models use VAE latents or diffusion. This forces architects to maintain two stacks of architectures, and prevents knowledge from flowing between perception and creation. Even unified attempts like Emu or SEED-LLaMA still rely on diffusion for generation, breaking the next-token-prediction paradigm that has been so successful in language.

This paper asks: can we build a single autoregressive model that does everything – understand, generate, edit – using one kind of discrete token, and improve it with reward-based tuning just like we do for LLMs? The answer is yes, provided you have the right tokenizer and a strong enough RL recipe.

Problem: Inconsistent representations block unified multimodal models
  |
  v
Assumption: Discrete semantic tokens + autoregressive next-token prediction
  can bridge understanding and generation in one architecture
  |
  v
Method: Train a multi-objective discrete visual tokenizer,
         then pretrain a 7B AR model on text + image token sequences,
         then apply RL to optimize task-level preferences
  |
  v
Evidence: WISE overall 0.50 -> 0.56; GEdit-Bench G_O 5.75 -> 6.68;
          RL also improves understanding tasks (zero-shot synergy)
  |
  v
Conclusion: Discrete autoregressive modeling + preference optimization
            is a scalable foundation for multimodal intelligence

The Increment

One sentence: Before ARM, unified multimodal autoregressive models either used continuous representations (breaking the discrete prediction paradigm) or relied on separate diffusion decoders; after ARM, a single discrete-token system with RL tuning achieves competitive results across understanding, generation, and editing – and the RL even helps the understanding side.

Core Mechanism

ARM has three components that form a pipeline. First, a discrete semantic visual tokenizer (VQ-VAE style but with extra supervision) converts an input image into a compact sequence of tokens – think of it as compressing a 256x256 image into 256 IDs, each carrying semantic meaning. Second, a 7B autoregressive transformer (decoder-only) is trained on an interleaved corpus: text-only sequences, image-only sequences, and text-image pairs. During training, the model learns to predict the next token regardless of modality. After pretraining, the model can generate images by simply continuing the text prompt with visual tokens, which the tokenizer’s decoder then reconstructs. Third, reinforcement learning (RL) fine-tunes the model on task-level objectives: for text-to-image generation, a reward model scores visual quality (e.g., Aesthetic Score) and instruction adherence (e.g., CLIP similarity); for editing, it scores edit consistency (e.g., CLIP directional similarity). The RL stage uses PPO-like optimization.

[Data flow: text or image input]
       |
       v
+------------------+       +----------------------+       +----------------+
| Discrete Visual   | ----> | Autoregressive       | ----> | RL Fine-tuning |
| Tokenizer         | image | Transformer (7B)      |       | (PPO style)    |
| (VQ-VAE + extra   | IDs   | Learns to predict     |       | Optimizes      |
|  supervision)     |       | next token (text or   |       | reward models  |
|                   |       | image ID)             |       | for each task  |
+------------------+       +----------------------+       +----------------+
       |                           |                             |
       | image -> token seq        | training objective:         | task rewards:
       | decoder: token seq->image | cross-entropy on all tokens | quality, adherence
       +---------------------------+-----------------------------+

Now the structural metaphor: think of ARM as a bilingual secretary who takes notes in a secret code. The visual tokenizer is the stenography system that turns any picture into a short string of code words (256 digits). The secretary (the 7B transformer) has two dictionaries: one for natural language words and one for code words. She reads mixed memos (text + image code sequences) and learns to predict what the next word or code word should be – this is autoregressive training. After training, when you dictate a text description, she continues writing code words, and the stenography system decodes them back into a picture. But sometimes her output isn’t great – blurry image, wrong instruction. So you hire a quality assurance manager (RL) who gives her bonuses based on how “good” the final picture looks, how well it matches the description, and (for editing tasks) how precisely the requested change was made. Over time, she learns to produce code sequences that maximize the bonus, and surprisingly, this also makes her better at writing summaries of images (understanding) – cross-task synergy from a unified reward signal.

Key Concepts

• Discrete semantic visual tokenizer: A VQ-VAE trained not only on reconstruction but also on semantic alignment: the token IDs should cluster images with similar meanings, even if pixels differ. Extra supervision comes from a language model’s embedding (e.g., CLIP text encoder) – the tokenizer’s output is forced to be close to the language representation of the image caption. This ensures that the token IDs carry semantic information that can be predicted by the autoregressive model just like words.

• Autoregressive multimodal pretraining: The model receives sequences where text tokens and visual tokens are interleaved (e.g., “A photo of <img_id1> <img_id2> …”). Because all tokens are discrete IDs (vocabulary ~ 16K text tokens + 16K visual tokens), the training is identical to a language model: predict the next token. This is simpler than previous mixed-continuous/discrete approaches and allows scaling to 7B parameters using standard LM infrastructure.

• Reinforcement learning for cross-task synergy: After supervised pretraining, ARM applies RL separately on generation and editing tasks using task-specific reward models. Crucially, the RL stage does not freeze the pretrained weights; it updates the entire model. This causes shared representations to improve for both tasks – editing RL improves text-to-image generation and vice versa, because better visual quality and instruction understanding benefit both. This is the paper’s “surprising” finding, and it mirrors the phenomenon in LLMs where RLHF on chat sometimes boosts reasoning.

Framework Shift

Before (mainstream unified models like Emu, SEED-LLaMA):

[Image] -> [Continuous VAE] -> [LLM encoder] -> [Diffusion decoder] -> [Image output]
                          different token types:   separate generation module

After (ARM):

[Image] -> [Discrete semantic tokenizer] -> [7B AR Transformer (next-token prediction)]
           all tokens are discrete IDs                  |
                                                        v
                                              +---[Token decoder]---> [Image output]
                                              +---[RL fine-tune on task rewards]

One sentence: From a hybrid continuous+discrete+diffusion architecture to a pure autoregressive discrete-only pipeline with reinforcement learning – the core shift is *unifying the token format across all modalities and tasks, and then applying task-level reward optimization just like we do with LLMs.

Expert Assessment

Problem choice: This is a real gap. The community has been chasing unified models but many still glue separate components. The question of whether a single discrete AR model can match specialized models is both scientifically interesting and practically valuable. It sits at the convergence of LLM scaling laws and multimodal learning.

Method maturity: Mostly engineering cleverness rather than a deep theoretical insight. The tokenizer with extra semantic loss is neat, the autoregressive pretraining is straightforward, and the RL is standard. The main contribution is showing it all works together, and especially the RL cross-task synergy is surprising and non-trivial. Could there be simpler approaches? Possibly a single diffusion model with discrete conditioning (e.g., UViT) – but ARM keeps the pure AR paradigm, which is attractive for integration with existing LLM stacks.

Experimental integrity: Fair baselines: they compare with Emu, SEED-LLaMA, GILL, etc., and use standard benchmarks (MS-COCO, WISE, GEdit-Bench). Numbers hold up; gains from RL are consistent and non-trivially large. One red flag: the RL reward models themselves might be overfitting to the benchmark metrics? But the paper includes human evaluation for editing, which mitigates that. Also the test set is not synthetic – they use real image captions and editing instructions.

Writing quality: The paper is generally well-structured but the tokenizer training details are a bit dense – the loss functions are listed but the rationale for each component could be clearer. The sections on RL reward design are good. If I had to rewrite one section, it would be the tokenizer architecture overview – too many equations without intuitive explanation for why each loss term matters.

Verdict: weak accept – a solid engineering paper with a surprising finding (RL cross-task synergy) that opens up a promising line for future work. Not groundbreaking but definitely worth attention if you work on unified multimodal models or RL for generation.

Takeaways

• The discrete tokenizer with semantic alignment (extra loss against CLIP text embedding) is a concrete technique you can steal for any VQ-VAE application where you want the codes to be semantically meaningful.
• The “RL on tasks, measure cross-task effect” methodology is a useful evaluation framework: don’t just report final numbers, but ablate the impact of RL on non-target tasks. This reveals hidden synergies.
• For practitioners building unified models: consider keeping the generation in the same discrete sequence space as language – it simplifies training and allows reuse of LM infrastructure (e.g., Megatron, DeepSpeed). The 7B scale is reasonable for a single node, and the RL stage adds maybe 20% additional compute.