Tractable Transformers for Flexible Conditional Generation

We thank the reviewer for their constructive feedback and for recognizing the potential of Tracformers in conditional generation tasks.

Some abbreviations are repeatedly defined throughout the paper, which leads to unnecessary redundancy.

We appreciate the reviewer’s feedback and will revise the paper to ensure that abbreviations are defined only once.

The contributions of the paper would be more clearly understood if they were summarized at the end of the Introduction section.

We thank the reviewer for the suggestion. In summary, our contributions are two-fold: (i) we identified the problem that existing NAR models suffer from severe performance degradation in conditional generation, despite having strong unconditional generation performance; (ii) we propose Tracformer, a novel Transformer-based architecture specially designed to improve conditional generation and generalize to conditional queries unseen during training. In the next version of the paper, we will add this summary at the end of the Introduction section to enhance clarity.

In the experimental section, the authors compare their proposed Tracformer model only with BERT and BART. However, more recent and powerful models, such as GPT, should be included in the comparisons presented in Figures 5 and 6.

While autoregressive models like GPT are indeed very powerful, they cannot be used in the contextual AR (CAR) and arbitrary-context (AC) generation experiments in Figures 5 and 6. This is because GPT models can only condition on prefix prompts, whereas our experiments require conditioning on context that is scattered throughout the sequence.

In our experiments, we included GPT-2 as a baseline in Table 4, as it has a similar number of parameters compared to our model and other baselines. In this comparison, all models were trained on the WebText dataset (or its open-sourced version) and evaluated for their zero-shot unconditional perplexity across multiple datasets.

The results for zero-shot unconditional perplexity in Table 4 show that Tracformer outperforms the other baselines on only the 1BW dataset.

We acknowledge that the zero-shot unconditional perplexity of Tracformer is less favorable compared to MDLM. However, the primary focus of this paper is to highlight the limitation of existing NAR models, including MDLM, which achieve strong unconditional generation performance but struggle with conditional generation. While Tracformer exhibits slightly worse unconditional generation performance, it significantly outperforms MDLM in multiple conditional generation tasks (as shown in Section 6.2). Since conditional generation is often more critical for downstream applications, we believe this aspect is more important.

Another potential factor influencing the results is that the Tracformer model used in our experiments has fewer parameters (109M vs. 169M) and was trained on fewer tokens (295B vs. 524B). This is summarized in Table 3. In future work, we plan to scale up Tracformers and explore different training objectives, such as the diffusion objective used by MDLM, to further evaluate their scalability and versatility.

The Background section could be moved to the appendix.

We thank the reviewer for the helpful suggestion. In the next version of the paper, we will move most of the background of sequence modeling and Transformers to the appendix, retaining only the essential parts needed to introduce basic notations.

…more detailed experiments and comparisons with recent models, such as GPT, should be conducted and presented in the experimental section. A deeper analysis comparing the proposed model against more state-of-the-art architectures would provide a clearer picture of its strengths and weaknesses.

We thank the reviewer for the suggestion. In Section 6.2, we compared Tracformer against state-of-the-art autoregressive models (e.g., GPT-2) as well as non-autoregressive models (e.g., SEDD and MDLM). For unconditional generation tasks, we included both autoregressive and non-autoregressive baselines. For conditional generation tasks, only non-autoregressive models were included since autoregressive models cannot condition on suffix contexts, making a direct comparison infeasible. In the next version of the paper, we will add the above discussion to the experiment section.

We thank the reviewer for their constructive feedback and for acknowledging that the paper addresses a clearly defined and significant problem.

At line 630 in Appendix A, are you referring to Decoder instead of Encoder?

We thank the reviewer for pointing out this typo, and the paragraph title should be “Sparse Attention Masks of the Decoder”. We will fix the typo in the next version of the paper.

In Figure 4(a), is the token identical to x₀ from the prefix encoder input understood from a teacher forcing perspective?

Yes, the target tokens are identical to the input. However, as described in Section 4.3 (Equations (5) and (6)), we set the cross-attention masks from the decoder to the encoder such that the decoder will not receive information about $x_0$ from the prefix encoder when predicting $X_0$. This is done by not attending to any token in the prefix encoder.

In TracFormer, is it correct to understand that the decoder combines the self-attention causal mask and cross-attention into a single cross-attention operation?

There is no self-attention in the decoder of Tracformers, and the model predicts a token only using information from the encoder(s) through the cross-attention operations. This design choice (not including self-attention layers in the decoder) was made mainly for (inference-time) efficiency considerations, as the decoder can be fully parallelized when predicting multiple tokens.

However, since equation (5) appears to merely apply a stride, I do not fully understand how it guarantees that the variable scope does not include X_C. Could you clarify this?

As shown in Figure 4, in the case of arbitrary-context (AC) generation, the masked tokens (variables not in $X_C$) are replaced by the mask token (<MASK>) so the model will not observe these tokens. In the case of contextual AR (CAR) generation, the condition $t’ < t$ in Equation (5) is used to guarantee that the model only observes previous tokens from the prefix encoder. In the next version of the paper, we will add more discussion to clarify how we guarantee the model does not observe tokens outside of $X_C$.

Additionally, we agree with the reviewer that there are other ways to design the decoder masks to ensure the model does not observe variables that it should not during training. We plan to explore other design choices in future work.

The proposed method employs two encoders (a prefix encoder and a suffix encoder). What is the fundamental advantage of using these two uni-directional encoders over a bidirectional encoder?

This design choice is made mainly because it aligns well with the contextual AR (CAR) generation paradigm, where the prefix encoder always encodes the full prefix context while the suffix encoder represents the given suffix tokens. It would be very interesting to see which (two unidirectional encoders versus one bidirectional encoder) performs better in the arbitrary-context (AC) generation paradigm, and we plan to explore it in future work. We will discuss this in the conclusion section in the next version of the paper.

In Figure 5, if the mask ratio is 1.0, does this correspond to the unconditional perplexity? If so, should we interpret that unconditional perplexity is worse than conditional perplexity?

In Figure 5, we only vary the mask ratio from 0.1 to 0.9 in increments of 0.1. Therefore, the right-most conditional perplexity values correspond to a mask ratio of 0.9. In general, as the mask ratio increases, the conditional perplexity worsens because the model has access to less contextual information.

We thank the reviewer for their constructive feedback and for recognizing Tracformers as a promising architecture for addressing the challenges NAR models face in conditional generation.

The abstract is quite unclear in retrospect; terms that are well defined later (conditional probability query) appear obscure.

We appreciate the reviewer’s suggestion. In the next version of the paper, we will clarify the term “conditional probability query” directly in the abstract, referring to the set of tokens/variables provided during NAR generation.

L238: while this remark may only be here to provide intuition for the design of the model, it would be better if it was backed by a reference.

While the main goal of this sentence is to give intuition, it can indeed be supported by references such as [1], which highlights the effectiveness of a top-down hierarchical scheme in long text generation. We will add the reference and further discuss it in the next version of the paper.

Footnote 1 is unclear: I don't see how it matches Equation 3.

We thank the reviewer for pointing out the typo in Footnote 1. According to the definition in Equation 3, we have that $\phi^{l-1}_{t-2^{l-1}}$ = {$t’ : t’ \geq 1, 0 \leq t - 2^{l-1} - t’ < 2^{l-1}$}

and

$\phi^{l-1}_{t}$ = {$t’ : t’ \geq 1, 0 \leq t - t’ < 2^{l-1}$}.

Taking the union of the two sets leads to

$\phi^{l}_{t}$ = {$\{t’ : t’ \geq 1, 0 \leq t - t’ < 2^{l}\}$}.

We will fix it in the next version of the paper.

While this paper is generally well written, Section 4 stays difficult to follow.

In our humble opinion, the main reason Section 4 is more challenging to follow than the other sections is the indexing of the variable scopes for different tokens, which is necessary to rigorously define Tracformer. To improve clarity, we will add a figure with one example Tracformer layer where each feature embedding is labeled with the corresponding variable scope. This figure will then be frequently referenced throughout the model description.

[1] Wan, Kaiyang, Honglin Mu, Rui Hao, Haoran Luo, Tianle Gu, and Xiuying Chen. "A Cognitive Writing Perspective for Constrained Long-Form Text Generation." arXiv preprint arXiv:2502.12568 (2025).

We thank the reviewer for their constructive feedback and for recognizing the novelty and clarity of our work.

I understand that the proposed model is smaller, but more experiments (e.g. slightly scaling it up to match baseline model sizes) may be needed to more convincingly demonstrate its capability of unconditional generation.

We thank the reviewer for the suggestion. As pointed out by the reviewer, we would like to highlight that although the unconditional generation performance is less favored compared to the baselines, our model achieves significantly better conditional generation performance. We will adjust the descriptions and claims about the unconditional generation performance of Tracformers versus the baseline models.

We agree with the reviewer that it would be valuable to conduct further experiments by scaling up Tracformer to assess how its performance improves with model size. Additionally, as discussed in the conclusion, another promising direction is exploring how Tracformer can be combined with recent advancements in training objectives for discrete diffusion models, such as diffusion language models. We plan to investigate these directions in future work.

the tasks, datasets, and evaluation metrics used in this paper mostly correspond to generic language modeling, emphasizing on fluency.

We thank the reviewer for the suggestion on extending the evaluation metrics. We focus mainly on metrics that reflect generative modeling performance (conditional log-likelihood/perplexity) and fluency (MAUVE, BERT score) to study the fundamental problem of NAR models having bad conditional generation performance. Although fluency alone may not be sufficient to fully characterize generation quality, it provides insights into how well the generated text aligns with natural language distributions. Perplexity, on the other hand, serves as a general measure of the model’s generative modeling performance.

To comprehensively evaluate Tracformers, we evaluated the BLEU-4 scores of Tracformer and the baseline model BART on the WikiText-103 dataset. The model was evaluated across six different masks. The evaluation setup follows Table 1. As indicated by the results, Tracformer outperforms BART consistently. We will include the results in the next version of the paper.

The pre-trained SEDD was only able to handle sequences of length 1024. How did you calculate the log likelihood of shorter sequences?

While SEDD was trained with sequences of length 1024, as described in the original paper (e.g., in their Section 5.3.2), they can handle sequences of small length. Specifically, when generating a sequence of length N, we only provide the first N tokens to the Transformer model. The authors of SEDD provided code to handle shorter sequences in their official GitHub repository, and we used their implementation. Note that all other models were also trained with sequences of length 1024, so the evaluation protocol is consistent across all models.

Does SEDD output the exact model-predicted likelihood, or its lower bound (ELBO)?

As described in their paper, SEDD, like other discrete diffusion models, can only compute the Evidence Lower Bound (ELBO) rather than the exact model-predicted likelihood. This is a common limitation in diffusion models since computing the exact likelihood is intractable.

In Table 1: What does a mask range of [0.25, 0.75] mean? Does it mean that out of the context window of 1024, the positions from 256 to 768 are masked, while other positions are visible?

Yes, the range of [0.25, 0.75] specifies the proportion of the context window that is masked. In the case of a context window of 1024 tokens, this means the positions from 256 to 768 are masked.

I would like to learn from the authors about their insights on my feedback shared in the above sections of this review form.

We thank the reviewer for their valuable comments, which we found very helpful in improving the quality of our paper. For example, as suggested by the reviewer, we added new metrics to further strengthen the experiment section. We also fully agree that text generation quality can be assessed from multiple perspectives.