Selftok-Zero: Reinforcement Learning for Visual Generation via Discrete and Autoregressive Visual Tokens

Thank you very much for your valuable suggestions and positive feedback. We will revise our paper in accordance with your recommendations.

W1：Supplementary Content for Tokenizer

Thank you for pointing that out! For the theoretical justification similar to our work, please refer to [1]. The theorem in this paper proves that the loss of image attributes increases as more noise is added. This thus enables adding tokens in a reverse-autoregressive manner to make up for the loss of attributes.

We apologize for the insufficient coverage of content related to the tokenizer. Our work aims to emphasize that AR tokens can still support effective RL-based post-training for visual generation, so we condensed the content regarding the tokenizer. Thank you again sincerely for your suggestions. We will clarify this part in the revised version.

[1] Yue Z, Wang J, Sun Q, et al. Exploring diffusion time-steps for unsupervised representation learning[J]. arXiv preprint arXiv:2401.11430, 2024.

Q1：Supplementary Content for Tokenizer

We sincerely apologize for the insufficient coverage of content related to the tokenizer. We will add more details about the tokenizer in the main text to make this section more comprehensive. Thank you again for your suggestion.

Q2：RL-related Content

We sincerely apologize for the redundancy in the proof of RL theory. We will correct it in the revised version. Thank you very much for your suggestion.

Q3：Typo

Thank you for pointing out the typo, we will correct it in the revised version.

Overall, I believe this is a technically solid paper, with high impact on visual generation, with novel method, excellent evaluation and reliable reproducibility.

In the rebuttal, the authors resolved all my concerns, and therefore I will maintain my score of 5.

Thank you very much for your constructive questions and valuable suggestions. We are delighted to discuss our work with you, and our responses are as follows:

W1：Clarification on tokenizer and motivation

We are sorry that due to page limitations, we had to place some content related to the tokenizer in appendix. Thank you for your suggestion; we will revise the paper to include more details about the tokenizer in the main text to make it easier for readers to understand.

The main point of the paper is to demonstrate that autoregressive (AR) tokens is the key for reinforcement learning (RL) in visual generation. In the introduction, we explain why RL is less effective for diffusion models and AR models based on spatial tokens. To address this gap, we reveal that the reverse diffusion process is inherently AR through a recursive decomposition. This allows us to learn AR tokens that support effective visual RL.

Q1：Why is a raster-scan based solution problematic for the policy improvement optimality in RL?

As shown in Figure 1 of our paper, the raster-scan based tokens do not follow the AR causal dependencies, as they are essentially inter-dependent, $e.g.$, knowing one image patch will leak information of other patches. We prove that this violates Bellman equation in Section 2.2.

Q2：What is the key reason behind Selftok-Zero's results?

As discussed in W1, our work aims to show that AR tokens enable effective RL-based post-training for visual generation. We constructed AR tokens as visual representations based on the autoregressive diffusion process, aligning with the policy improvement theorem. Experiments on GenEval and DPG benchmarks show its effectiveness: only AR tokens meet the conditions for the existence of an optimal solution, while other non-AR tokens yield weaker RL results—highlighting our work's efficacy.

Thank you very much for your valuable comments and insightful questions. Our responses are as follows:

W1：Supplementary Content for Tokenizer

We apologize for the insufficiency of content related to the tokenizer; our work focuses on emphasizing that autoregressive (AR) tokens is the key for reinforcement learning (RL) in visual generation. Thank you again sincerely for your suggestion—we will include more details about the tokenizer in the main text.

W2：Reward Model Design and Bias

Great point! We agree that this limitation has potential influence on the policy. And we apologize that in this work, we did not explicitly address those biases, as our focus is to demonstrate that AR tokens enable effective visual RL, whereas spatial tokens are less effective in this regard. And it is worth noting that even with this potential deficiency, we still SOTA. This highlights the potential of visual RL with Selftok tokens, and we will improve the RL pipeline in our future work. This issue may need to be addressed in future work through better reward models and large-scale RL.

W3：Image Resolution

Thanks for pointing that, actually our tokenizer can be extended to 512x512, as shown in Figure C10 in Appendix. We also plan to extend it to higher resolutions in future work. Furthermore, achieving competitive results at 256x256 resolution is more challenging, as generating image details is more difficult; this further illustrates the effectiveness of our method.

Q1：Explanation of Selftok

We would love to provide a causal graph that makes the explanation easier, but couldn't do so due to rebuttal policies. We will try our best to explain intuitively to you. Recall that the collider effect arises as spatial tokens are encoded by conditioning on the original image $x_1$. In our work, the key mechanism to circumvent this problem is to introduce a random noise $x_0$ independent from $x_1$. From this $x_0$, we learn tokens that gradually guide the reverse diffusion process from $x_0$ to $x_1$. Specifically, each token is learned to enable an additional step towards $x_1$, from a position $x_t$ that is dictated by previous tokens. This naturally leads to causal dependencies among all tokens.

Q2：Reward Ablation

This issue is indeed a critical component. We will show that models trained solely with program-based rewards do have improvements on DPG benchmark, with the specific results as follows:

$*$ represents model trained only with the program-based reward

We will also continue to explore more settings related to rewards in future work. Thanks again for your valuable question.

Q3：Training Stability

Nice question! We have indeed encountered challenges that significantly impact the stability of RL training. Specifically, we found that training becomes highly unstable when the learning rate is high; therefore, we set the learning rate to $1.5e−6$. Furthermore, training also becomes unstable when there is a mismatch between the model's sampling and update frequencies, so we maintain consistency between these frequencies to avoid this issue.

We also explored the effects of three different KL divergence coefficients, with the results shown as follows:

We found that when the KL coefficient was 0, training became very unstable. We tried our best to set hyperparameters, but we were unable to stabilize the training. When the KL coefficient was 0.1, the model's convergence speed slowed down, and it was unable to exceed the GenEval score when the KL coefficient was 0.05.

Thank you very much for your constructive comments and invaluable evaluation. We are pleased to see that you have raised very critical questions regarding our work, and our responses are as follows:

W1：Robustness of the Tokenization Scheme.

Nice observation! The choice of the diffusion process $q(x_t|x_1)$ and the token schedule $k(t)$ should ensure that all decompositions $P(V_K)=P(V_{{\lt}i})P(V_{{\geq} i}|V_{{\lt}i})$ are accounted for in training. For example, consider a failure case where $x_t=x_0$ for any $t\neq1$, or a token schedule such that all tokens are allocated to $V_{{\ge}1}$, $i.e.$, $k(t)=1$, $∀𝑡 ∈ [0, 1)$. This corresponds to a trivial decomposition$P(V_{{<}1} = [\space])·P(V_K|V_{{<}1} = [\space])$, $V_K ⇔ x_0 ⇝ x_1$, and $[\space] ⇔ x_0$ where we always input the full $V_K$ to the decoder. So $V_K$ loses all the AR property.

For the forward process, we adopt rectified flow ($i.e.$, $x_t=tx_1+(1-t)x_0$), as we initialize our decoder from SD3. We are happy to try other scheduling strategies and noise types as future work.

For token schedule, as we uniformly sample $𝑡 ∈ [0, 1]$ in training, the best token schedule should be a uniform assignment $𝑘^∗(𝑡) = ⌈𝑡 × 𝐾⌉ + 1$ to ensure that every decomposition is equally respected in the training objective. We perform ablations in Token schedule of Appendix C. Interestingly, we observe a slightly better reconstruction quality by designing a schedule $𝑘(𝑡)$ that allocates fewer tokens to smaller $𝑡$, $i.e.$, $𝑘(𝑡) < 𝑘^∗(𝑡) \space for \space 𝑡< 0.5$. This aligns with the well-known trait of diffusion models: the early path $x_0 ⇝ x_𝑡$ for a small $𝑡$ has minimal impact on the reconstruction $x_𝑡 ⇝ x_1$, which can be omitted.

W2：Semantic Interpretability of Tokens

Yes, we find that tokens corresponding to smaller time-steps tend to capture the overall background, color tone or composition of the image, those at middle ones tend to capture object shapes and those at larger ones tend to capture fine-grained details and textures. This is because the diffusion process itself is tightly linked with visual semantics [1,2,3], and Selftok simply encode the process as tokens.

We would love to show you some visualizations that demonstrate the interpretability of our tokens, but we are unable to do so due to the constraints of the conference policy. We will follow your suggestion to include those figures in the revision. Meanwhile, you can check similar results in Figure 1 of [1].

[1] Yue Z, Wang J, Sun Q, et al. Exploring diffusion time-steps for unsupervised representation learning[J]. arXiv preprint arXiv:2401.11430, 2024.

[2] Preechakul K, Chatthee N, Wizadwongsa S, et al. Diffusion autoencoders: Toward a meaningful and decodable representation[C]. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 10619-10629.

[3] Zhang Z, Zhao Z, Lin Z. Unsupervised representation learning from pre-trained diffusion probabilistic models[J]. Advances in neural information processing systems, 2022, 35: 22117-22130.

W3：Reward Model Design and Bias

Great point! We agree that this could introduce potential biases. In this work, we did not explicitly address those biases, as our focus is to demonstrate that AR tokens enable effective visual RL, whereas spatial tokens are less effective in this regard. And it is worth noting that even with this potential deficiency, we still achieve state-of-the-art performance. This highlights the potential of visual RL with Selftok tokens, and we will improve the RL pipeline in our future work.

W4：Direct Comparison with Competing Works

We fine-tuned a diffusion model (SDXL) with Diffusion-DPO using the same reward model on the GenEval benchmark. To ensure the identical settings, we train SDXL with the same amount of data, and we used prompts and images to construct positive and negative pairs, whereas Selftok-Zero only use prompts. The results are as follows:

The table above shows that Diffusion-DPO does not exhibit a significant improvement(+2.5) on the GenEval benchmark, while our method leads to a +18 increase in overall score, which further demonstrates the effectiveness of our method.

W5：Quantitative Evaluation of Encoder Quality and Efficiency

The comparisons with other works, as well as the results of quantitative metrics related to the Tokenizer, are as follows:

From the table above, it can be observed that our tokenizer has achieved the best performance to date, reaching optimal levels in metrics such as rFID (0.54), PSNR (26.30), SSIM (0.805), and LPIPS (0.063). Additionally, it boasts high computational efficiency, requiring only 0.86 seconds to encode and decode an image.