Tractable Transformers for Flexible Conditional Generation
Authors: Anji Liu
Xuejie Liu
Dayuan Zhao
Mathias Niepert
Yitao Liang
Guy Van den Broeck
Venue: ICML 2025 poster (2025)
Keywords: Generative Models
Conditional Generation
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Non-autoregressive (NAR) generative models are valuable because they can handle diverse conditional generation tasks in a more principled way than their autoregressive (AR) counterparts, which are constrained by sequential dependency requirements. Recent advancements in NAR models, such as diffusion language models, have demonstrated superior performance in unconditional generation compared to AR models (e.g., GPTs) of similar sizes. However, such improvements do not always lead to improved conditional generation performance. We show that a key reason for this gap is the difficulty in generalizing to conditional probability queries unseen during training. As a result, strong unconditional generation performance does not guarantee high-quality conditional generation. This paper proposes Tractable Transformers (Tracformer), a Transformer-based generative model that is more robust to different conditional generation tasks. Unlike existing models that rely solely on global contextual features derived from full inputs, Tracformers incorporate a sparse Transformer encoder to capture both local and global contextual information. This information is routed through a decoder for conditional generation. Empirical results demonstrate that Tracformers achieve state-of-the-art conditional generation performance on text modeling compared to recent diffusion and AR model baselines.

Tractable Transformers for Flexible Conditional Generation
Anji Liu 12 Xuejie Liu 33
Dayuan Zhao
Mathias Niepert?
Yitao Liang? Guy Van den Broeck'
Abstract
WikiText103 (length 128)
1BW (length 128)
Conditional Perplesity
65
70
SFDD (1AMM)
GP 2: (1244)
Conditional Porplexity
SEDI (L60M)
GPT-2 (1244)
Non-autoregressive (NAR) generative models are
MDLM (15RM)
100
MDLM (160M)
60
Tracformer (Chass 109M)
Tracdormer (Ours; 10HM)
valuable because they can handle diverse con- ditional generation tasks in a more principled way than their autoregressive (AR) counterparts, which are constrained by sequential dependency requirements. Recent advancements in NAR mod- els, such as diffusion language models, have demonstrated superior performance in uncondi- tional generation compared to AR models (c.g.. GPTs) of similar sizes. However, such improve- ments do not always lead to improved conditional generation performance. We show that a key rea- son for this gap is the difficulty in generalizing to conditional probability queries (i.e., the set of unknown variables) unseen during training. As a result, strong unconditional generation per- formance does not guarantee high-quality condi- tional generation. This paper proposes Tractable Transformers (Tracformer), a Transformer-bawe generative model that is more robust to differ- ent conditional generation tasks. Unlike exist- ing models that rely solely on global contextual features derived from full inputs, Tracformers in- corporate a sparse Transformer encoder to cap- ture both local and global contextual informa- tion. This information is routed through a de- coder for conditional generation. Empirical re- sults demonstrate that Tracformers achieve state- of-the-art conditional generation performance on text modeling compared to recent diffusion and AR model baselines. Code is available at https: //gi thub. com / juanj h ormer.
55
80
50-
60
45
40
40
0.8 Proportion prefix provided Figure I: Zero-shot conditional perplexity (L) of AR (GPT-2 (Rad- ford et al., 2019)) and NAR models (SEDD (Lou al., 2023), MDLM (Sahoo al., 2024)) models conditioned different pro- portions prefix given sequences length 128. Existing NAR models perform worse than GPT-2 despite having comparable unconditional perplexity sequences of length 1024. Their per- formance gets when provided with more context.
0.2
Proportion of prefix provided
1. Introduction
Generative AI has emerged as transformative paradigm for solving machine learning tasks. Its core strength lies in the ability of modern deep generative models to learn complex and high- nigh-dimenssonal data distributions. Autoregressive (AR) models such as GPTs (Brown et al., 2020) are among the most well-developed generative models, demonstrating exceptional performance in modeling discrete data such as language (Dubey et al., 2024; Reid et al., 2024) and protein sequences (Shin et al., 2021; Trinquier et al., 2021). However, despite their expressiveness and scalability, AR models are not best suited for many conditional generation tasks such as DNA imputation and protein sequence infilling (Tanaka et al., 2022; Hawkins- Hawkins-Hooker al., 2023; Alamdari et al., 2023) due to their inherent (autoregressive) sequential dependencies (Kaddour al., 2023).
In contrast, non -oo--uttoreereiivee (NAR) generative models are inherently more flexible, as they can condition on ar- bitrary contexts. This flexibility makes them particularly well-suited for conditional generation tasks, such as code editing (Liu et al., 2024) and DNA imputation (Stark et al., 2024; Li al., 2024; DaSilva et al., 2024). Recent advances in discrete diffusion (Lou et al., 2023; Shi et al., 2024) and semi-auten autorrrrggeeeeevveeeeeeeee models (Chen et al., 2024) have further improved the performance of NAR models across various domains. Notably, SoTA diffusion language models outper- form GPT-2 in terms of unconditional perplexity on several benchmarks (Sahoo et al., 2024).
"Equal contribution 'Department of Computer Seience, Uni- versity California, Los Angeles IInstitute for Artificial Intelli- gence. University of Stuttgart Institute for Artificial Intelligence, Peking University "School of Intelligence Science and Technology, Peking University 'Yuanpei College, Peking University. Corre- spondence to: Anji Liu <liuanji www.eeesss.e..eeeeeeeeeeeeeaaaaa www..ceuu. Yitao Liang <yi- taol @ @hol@pku.coom cdu.eeduen>, Guy Van den Broeck <guyvdb@ wewrdddess.ul....o ucl...u.dd
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Tractable Transformers
However, observe significant performance drop when using SoTA diffusion language models for conditional gen- eration tasks. As shown in Figure 1, we evaluate the condi- tional perplexity (i.e., the perplexity of the corresponding conditional likelihood) of two SoTA diffusion language models, SEDD (Lou et al., 2023) and MDLM (Sahoo et al., 2024) by providing different fractions of the prefix of length- 128 sequences from WikiText103 (Merity et al., 2022) and IBW (Chelba et al., 2013). Despite both models having comparable or better unconditional perplexity than GPT-2 on length- 1024 sequences, their conditional perplexity is significantly worse. Moreover, their performance degrades further as the length of the provided prefix context increases.
of many modern sequence modeling frameworks, highlight- ing its key components and their roles in both AR and NAR paradigms (Sec. 2.2). We use uppercase letters (e.g. X²) to represent random variables and lowercase letters (e.g.. 21) for their assignments.
2.1. Sequence Modeling
Given a sequence of T categorical variables X := (X+) sequence modeling aims to capture their joint distribution. Autoregressive modeling achieves this by factorizing the joint probability of æ using the chain rule of probability:
Pr(z)=Pr(zz) Pr(xz)))) (
(1)
We attribute this performance discrepancy to the inability of existing NAR models to generalize effectively to conditional queries unseen during training. To address this issue, we propose Tractable Transformers (Tracformers), an NAR generation model designed to handle diverse conditional generation tasks more robustly. A key insight is to learn local features that enhance generalization across different queries. Specifically, unlike existing models that rely solely on global features from all input tokens, Tracformers use novel sparse encoder to learn features multiple context levels, which are then processed by decoder for conditional generation. As shown in Figure 1, Tracformer achieves better conditional perplexity than other NAR models.
By learning each of the I' conditional distributions with a neural network, AR models such as GPT (Radford et al., 2019; Brown et al., 2020; OpenAI, 2023) and State Space Models (Gu & Dao, 2023; Fu et al., 2022) achieved state- of-the-art performance in modeling nigh-dimensional se- quences like text and protein (Nguyen et al., 2024).
However, AR modeling requires the context to be provided contiguously at the beginning, limiting its applicability for tasks that demand more flexible context handling, such as DNA imputation and protein sequence infilling.
In contrast, von-autoregressive (NAR) modeling seeks to learn conditional distributions given arbitrary contexts. For- mally, for subset of context variables Xc with C' C: [T]. NAR models aim to encode the conditional distribution Pr X xl of the remaining variables XR. where :== == CC. Here [7] denotes the set of positive integers up to T. Compared to AR models, the flexibility of NAR models makes them adaptable to a broader range of conditional generation tasks (Li al., 2022; Han al., 2023).
Empirical results on text generation show that Tracformer achieves consistently better conditional generation perfor- mance compared existing architectures such as BERT (Devlin et al., 2018) 8) and BART (Lewis et al., 2020). Further, Tracformer beats state-of-the-art (SoTA) diffusion language models on zero-shot conditional generation tasks, which are conditional variants of the zero-shot perplexity tasks used to evaluate GPT-2 (Radford et al., 2019). In addition to proposing an NAR model architecture, we emphasize the importance of directly evaluating the conditional generation performance of NAR models.
2.2. Transformer Models
The Transformer (Vaswani et al., 2017) architecture is com- monly used implement generative models for both AR and NAR generation. Given sequence of tokens {z, )T the model initially maps each token I; into continuous embedding he E Rd, where d denotes the dimensionality of the embedding space. The embeddings are then iteratively transformed through a stack of L Transformer blocks:
In summary, the contributions of this paper are two-fold: (i) We identified the problem that existing NAR models suffer from severe performance degradation in conditional gener- ation, despite having strong unconditional generation per- formance; (ii) We propose Tracformer, a novel Transformer- based architecture specially designed to improve conditional generation and generalize to conditional queries unseen dur- ing training.
E E (l,...,,4) L}, hi = Block((h
where hy represents the embedding at position in layer I and h := (h! denotes all embeddings in the layer.
2. Background
In this section, we first introduce key concepts of sequence modeling and the distinctions between autoregressive and non-autoregressivve approaches (Sec. 2.1). We then describe the Transformer architecture, which serves as the backbone
Omitting design details such as the use of layer normal- ization and residual connections, each Transformer block consists of two modules: a feed-forward neural network (FFN) and an attention module (Attn), which can be ex- pressed as h = FFN, (Attny (h))
2

---

Tractable Transformers
(a) The likelihood of "The cat is cute" differ when computed in different orders with SEDD.
(b) LI difference between the best and
Compute probability in the forward order
Compute probability in the revense order
the worst order for length-5 sequences.
0.1
log Pr(Xi The)
-6.08
log Pr(X, = cute)
-7.22
Median of best
0.08
LL: -31.58
Froquency
log PggXt(X,= cat a The)
11.61
log Pr(X33 is 244 cute)
--7.10
0.06
logPr(Xx is| X1.2 The cat)
-3.58
log Pr(XX: cat Xx.4 is cute)
-7.34
0.04
0.02
logPr(X4 cute| X1:3 The cat is)
-8.71
log Pr(Xi The X2:4 is cute) -11.77
10
15
20
25
30
1 = The cat cute)
-2.998
log Pr(X4:4 The cat cute)
-33.38
Best LL minus worst LL
Figure 2: Different conditional queries do not respect the same joint distribution in NAR models. (a) The ligllikelssoods (ILs) of the sequence "The cute" differ depending whether SEDD is queried in the forward order the reverse order. (b) Histogram of the difference between the highest and the lowest LLs of 2,000 length-5 sequences from WikiText103 when queried in all possible orders.
While the FFN is applied independently to embeddings at each position f, the attention module captures depen- dencies between different token variables. Specifically, a Iearnable function score; first computes a score value for each pair of token positions (t,)') using their respective embeddings. The scores are then normalized with the softmax function:
3. Query Generalization in NAR Models
As shown in Figure the strong unconditional genera- tion performance of SoTA NAR models does not ensure high-quality conditional generation, different conditional queries from the same model do not align with a consis- tent joint distribution. For example, for any text distri- bution, the log-likelihood (LL) of the sequence Ø1:4 := "The cat is cute" should remain the same regardless of how we break down the joint probability into conditional prob- abilities following the chain rule. However, as shown in Figure 2(a), the LLs computed with SEDD (Lou et al., 2023) in the forward and the reverse order are different, which in- dicates that the model does not follow the same underlying joint distribution when prompted with different conditional queries. Figure 2(b) further illustrates the prevalence of such inconsistency by measuring the gap between the highest and lowest LLs of 2,000 length-5 sequences from WikiText103 (Merity al., 2022), evaluated across all possible orders.
sh,e : score, hh (hi11 h
mhe exp(s(.e)
Vt,t's E [I.....T].
$6 = En ml exp(s)_er)
where m :- (m/,e e her E 1j is a mask that speci- fies the allowed dependencies between token positions.
The normalized scores {$L, Ye are then used to weigh the embeddings, with the weighted sum of these embeddings forming the final output of the attention module. Specifi- cally, the output at position t is computed as:
Attn: (h
She g((h")
As discussed in Section 2.2, tokens not given as evidence to NAR models are often represented by a special <MASK> token. In such cases, generalization to new queries means generalizing to inputs with different <MASK> distribution.
where 9: : RR > Rd is learnable mapping. When imple- menting AR models with Transformers, causal attention masks are used to ensure each token only receives context from previous tokens:
Existing models use different mask strategies during train- ing, which means they are trained to predict certain types of conditional queries, i.e., Pr(- (((-)az)) for specific subsets C C (T). For instance, discrete diffusion models (Austin et al., 2021) and BERT (Devlin et al., 2018) sample C' uni- formly at random while T5 (Raffel et al., 2020) sample random spans of variables. However, since it is infeasible to cover all conditional queries during training, NAR mod- els have achieve robust query generalization excel conditional generation.
m {o
il' otherwise.
e {1. L}.
The output h of the final layer at position is transformed to predict the conditional distribution Pr(X,Jar: (ef. Eq. (1)) through a learnable mapping.
In contrast, NAR modeling is commonly achieved by set- ting unobserved input tokens a special placeholder token <MASK>, while configuring the attention mask to allow full
visibility of all inputs (Devlin et al., 2018; Lou et al., 2023; Sahoo et al., 2024). Specifically, given contexts Tc. all tokens in XR with R := [T] are assigned the <MASK> token, and the final hidden embedding h is used to capture Pr(X,[ for every E R.
4. Tractable Transformers
The standard Transformer architecture introduced in Sec- tion 2.2 is highly expressive and scalable thanks to its attention modules that enable information transformation
3

---

Tractable Transformers
++3
t 2
f +
t+
2 =
2
n [++1
enc.2
h MIL0.2
cnc.2
dec.0 h [++
dec.0 h ++2
dec.0 h ++3
h 1
dec.0
Cross attention mask in decoder layer #2
Encoder layer #2
Sparse cross attention
Decoder layer #1
b dic,
dee
deo.
2
2
h
h
esnc, 1++11
c h ++2
h enc, L+43
dec, h h+1
h dec.1 1+-2
dec. M++3
h
dec,
tse h Fell
V
V
Encoder layer
Decoder layer #2
V
20
h
hgne..
h
enc.0
h
cnc.0 ++3
h dec
dec.2 h 1-22
deci2 1
dec
h 1+4
V
Scope size
X
XL+1 X-+-2
X1433
Xx++1
Xr+-2
X--3
X1.44
Figure 3: Tractable Transformers (Tracformers) use an encoder to learn feature embeddings (i.e., h; 1 with context (denoted as scope the left) sizes ranging from to 24, where L is the number of layers. The encoder features are then fed decoder for conditional generation. Each decoder layer contains a feed forward neural network and cross-attentis layer collect information from corresponsing encoder features. Attention layers in both the encoder and the decoder have special sparse patterns (see Secs. 4.1 and 4.2).
between arbitrary pairs of tokens. However, this global attention mechanism can hinder generalization to unseen conditional queries due to its sensitivity to changes in the distribution of the <MASK> token. In contrast, restricting the context of features h to local neighborhoods (i.e., makes h² only depends on a small subset of variables) improves robustness across different conditional queries since local features are invariant changes of mask tokens outside their respective context windows. However, using local features significantly restricts the model's 's expressiveness.
while later layers encode broader contextual information and more abstract representations. Figure 3 provides an example multi-scope encoder with base 2. Specifically, the scope of each embedding h, unc. is the set of variables whose distance from the left of X, is smaller than 2
{c' :'' 1,0 2')
MSSAs can be implemented by applying sparse attention masks to a standard attention module as introduced in Sec- tion 2.2. Specifically, the base-2 scope pattern in Figure can be achieved using the following attention masks for each I [L]. and t, E [T']
We propose an encoder-decoder architecture that effectively leverages both local and global features to achieve robust- ness against query changes while preserving the expressive- ness of global context modeling. As shown in Figure 3, the encoder learns feature embeddings with exponentially increasing context lengths in different layers by a sparse attention mechanism. The encoder embeddings are then fed to a cross-attentionnonf, decoder Transformer, which independently predicts the conditional probability of each token. the following, we introduce the backbone structure of the encoder (Sec. 4.1) and the decoder (Sec. 4.2).
{1 ift ift- St,
enc.l m,,e"'
(4)
otherwise.
Moreover, we note that only O(T) is in every attention mask m E (0, TXT are needed to get the desired expo- nentially increasing scope pattern. For example, when the base is 2, instead of attending to every token in the range L- 2 < l' S t (Eq. (4)), we only need to attend to the two tokens at the boundary, i... m,,"'
enc.! = 1 when t' = or
4.1. Multi-Scope Encoder
l' = t- 21-1, to obtain the desired variable scope in Equa- tion (3). This enables linear time (w.r.t. sequence length) implementations of the encoder, effectively avoiding the quadratic computation overhead standard Transformers.
Following Section 2.2. we define h h,mne!/ as the feature embed- dings at position t of the 1-th encoder layer (Fig. 3). Define the variable scope (or scope) of C [T] of h.mee! as the set of variables that contribute to the computation of h.se. For example, with full attention, the scope of every embed- ding spans the entire input sequence: Vt. and 1, of [T]; when using causal attention masks in Equation (2), we have o, = [t) for every and every t.
control this computational complexity, we define a hy- perparameter Nmax (Nmax > 2) that specifies the maxi- mum number of embeddings given token attend to using the attention mask mfñe." Specifically, if the range t- - 1' < t in Equation (4) contains no more than Nes. embeddings, all embeddings in the range are attended to
The encoder layers implement a Multi-Scope Self-Attention (MSSA) mechanism such that the size of the variable scopes (i.e., (0(1) grows exponentially with the layer index 1. This enables earlier layers to capture fine-grained local semantics,
According to the definition in Equation (3), we have that - = {t' 2 1,051 2/-1 - t' < 21 and of it' 211221,05-- 21) Taking the union of the two sets leads to of -11-1 J =f' :t' 2 1,0 S - 2

---

Tractable Transformers
(i.e., the corresponding mask values are 1). Otherwise, the attention mask selects exactly N embeddings from this range by sample linearly and round to the nearest integer.
decoder attention mask is defined as:
= - {
if < and I' = (mod (mod 1
dee.d ",e,
(5)
otherwise.
To ensure the existence of encoder embeddings with full context (i.e., its variable scope is [T), we choose the number of layers to be greater than or equal to log, T.
4.3. Tractable Transformers for NAR Generation
In this section, we demonstrate how to apply Tracformers to two common NAR generation paradigms- contexxual autoregressive generation and arbitrary-context generation.
4.2. Decoder
Recall from the beginning of Section that the decoder combines local and global features computed in the encoder to predict distributions. As illustrated in Figure 3, the de- coder consists of I. Transformer blocks, each containing feed-forward neural network (FFN) and cross-attention layer. Formally, given inputs /eec./-1 to the I-th decoder block, the outputs h 4 are computed as
Contextual AR Generation. Contextual AR (CAR) gen- eration refers to scenario in which, when predicting a variable Xx, the model has access to all preceding tokens #1 and a subset of future tokens. This paradigm is used when an arbitrary context øe is provided, and the model is tasked with autoregressively sampling all remaining tokens.
To implement Tracformer for CAR generation, we use two encoders, prefix encoder and suffix encoder, which cap- ture information from preceding and succeeding tokens, respectively. As shown in Figure 4(a), the prefix encoder processes the original sequence during training since pre- ceding tokens are always available in CAR generation. In contrast, the suffix encoder observes only the tokens in a chosen context set æc, while the remaining tokens are rep- resented by the <MASK> token.
dec. h 1: = FFN: (C dec.-11 enc. 1+1 Len.L-(t1))
where in the cross-attention operation, each decoder embed- ding h (Vl (T]) attends to the outputs h me. h 4-1+1 of the (I--I+))-th encoder layer. Intuitively, the decoder first cross-attend to encoder embeddings that encode global information (i.e., those with large scope sizes) before focus- ing on embeddings encoding local information (i.e., those with small scope sizes). This progressive refinement mir- rors human language generation, where high-level intent is first structured before being articulated into detailed expres- 2
The prefix encoder uses the sparse attention mask described in Section 4.1 (i.e., the sparse version of Eq. (4)), where the scope of each feature h ( includes the 2 preceding variables, including X,. We use similar attention masks in the suffix encoder such that the scope of each feature h. covers the 2' succeeding variables starting from X,. See Ap- pendix A for a formal description. We use the masks defined in Equation (5) for cross-attention between the decoder and the prefix encoder, as the attended features contain only pre- fix token information. Analogously, the suffix encoder uses the following cross-attention masks to ensure the decoder receives only suffix information from it:
Attention masks are used in cross-attention layers to ensure each decoder embedding only depends on desired input vari- ables. For example, when given the context &c, since the t-th decoder embedding h, d.dcc./ is used to predict Pr(X, (c),, h, doe, should not attend to encoder embeddings whose vari- able scope contain variables not in Xc.
l Formally, define mi. as the mask deciding whether the t-th embedding of the input to the 1-th decoder block (i.c.,
{1 if > 1 and t+1 (mod 24-141), (6) "y,""
dec.
h, ") can attend to the f'tth embedding in the corre- sponding encoder layer (i.e., h; L-t1.
otherwise.
As illustrated in Figure 4(a), given context Tc, the t-th output feature embedding of the decoder used predict the distribution Pr(X,/xc.) where C=:=Cr CU [t-1]] 1]. The overall training loss for the CAR generation model is
Different generation paradigms can be achieved using spe- dec,
cific decoder masks m = {m, h, For example, to achieve AR modeling with Traeformers, we set m,," 1 if dec,1
I' L and zero out all remaining mask entries. Additionally, we observe that attending to fewer encoder tokens in the initial decoder layers has little impact on the performance. Therefore, to improve computational efficiency, we scale the sparsity of the cross-attention in proportion to the scope length of the corresponding encoder layer. Formally, the
I((0)= = - -...-.. log Pro
(7)
LEC
where 0 is the set of learnable parameters, D is dataset, and Po is a mask strategy used to sample the context set C.,
Forr example, Wan al. (2025) highlights the effectiveness of a top-down hierarchical scheme in long text generation.
Arbitrary-Contex Generation. In arbitrary-context (AC) generation, the model is tasked to predict the distribution
5

---

Tractable Transformers
(a) Tracformer for contextual AR (CAR) generation.
(b) Tracformer for arbitrary-context (AC) generation.
Pr(
Pr(Xalza, Ta)
Pr(Xg 21, vz) Pr Pr(X2/a1,,)) xz))
Prefix encoder
Prefix encoder
Xg
X1
X2
X3
<MASK> X1 <MASK> X3
t Shared weights
Decoder
$ Shared weights
Decoder
t=0
t=1
t=2
t=3
t=0
t=1
t=2 2
t=
Suffix encoder
Suftix encoder
<MASK>
<MASK> X3
Given context: X (21
<MASK> X1
<MASK> X
Given context: x= (a
Figure 4: Tracformers for contextual AR (CAR) and arbitrary-context (AC) generation. both cases, we use a prefix encoder and suffix encoder to capture prefix and suffix information for every token, respectively. Model parameters of the two encoders are shared, The decoder receives information from features generated by both encoders. Specifically, the decoder attention masks ensure the decoder only acquires information of preceding (resp. succeeding) variables from the prefix (resp. suffix) encoder. During training, <MASK> tokens are given to both encoders if the model trained for AC generation. When training Tracformer for CAR generation, the prefix encoder always receives all inputs as preceding tokens æ1.t-1 are assumed to be given when sampling æt.
of X, given arbitrary context, which can be used to gen- erate missing tokens in an arbitrary order. Analogous to Equation (7), the objective of AC generation task
Architectures for NAR Modeling, Decoder-only trans- formers with full attention are the most widely adopted architecture for NAR modeling. Many SOTA discrete diffu- sion models use these models. Additionally, bidirectional autoregressive modeling, exemplified by models like BART (Lewis et al., 2020) and MASS (Song al., 2019), rep- resents an intermediate approach that incorporates bidi- rectional context while preserving the left-to-right autore- gressive generation process. Sun et al. (2023) developed a Transformer-based architecture for a subclass of discrete diffusion models. Liu et al. (2025) and Xu et al. (2025) com- bine diffusion models with other deep generative models, such as AR models and energy-based models.
L(0)= = - -E D,C~Po
log Pro(x))ze)
(8)
IEC
As shown in Figure 4(b), AC generation paradigm can be implemented by the same model described for CAR genera- tion, with the only difference that inputs to both the prefix encoder and the suffix encoder use the mask token if token does not exist in the context. Correspondingly, the I-th out- put feature embedding predicts the distribution Pr( (r.).
6. Experiment
Training and Inference Efficiency. Thanks to its sparse attention modules in both the encoder and the decoder, Trac- formers enjoy efficient training and inference. In particular, KV-caching (Pope al., 2023) can be used to amortize inference cost. See Appendix E for detailed discussion.
In this section, we aim to empirically evaluate Tracformer's effectiveness in both conditional and unconditional genera- tion. Specifically, our experiments are designed to answer two key questions: (i) How does Tracformer compare to other NAR architectures in terms of conditional generation performance? (ii) Can Tracformer scale effectively and outperform existing SoTA generative models in both condi- tional and unconditional tasks? To this end, we conduct two sets of experiments: In Section 6.1, we compare Tracformer with a range of NAR architectures on WikiText (Merity et al., 2022), LAMBADA (Paperno et al., 2016), and One Billion Words (IBW) (Chelba et al., 2013) datasets to evaluate its performance across diverse conditional queries. In Sec- tion 6.2, we scale Tracformer to Open (Gokaslan & Cohen, 2019) and benchmark it against SoTA discrete diffusion models, focusing on zero-shot conditional and unconditional performance. These experiments comprehen- sively evaluate Tracformer's advantages and its potential to serve as a more effective backbone for NAR generation. Note that the perplexity numbers of diffusion model base- lines computed using the corresponding ELBO since the exact log- logelikel is intractable to compute.
5. Related Work
NAR Modeling Techniques. BERT (Devlin et al., 2018; Warner et al., 2024) is one of the first Transformer models designed for NAR generation. It proposes use special mask token to indicate unknown tokens and task the model to predict them given the observed tokens. Built on top of this mask prediction principle, discrete diffusion models (Austin et al., 2021) improve NAR generation performance by designing better learning objectives (Campbell et al., 2022; Lou et al., 2023; Sahoo et al., 2024) and mask strate- gies (Shi et al., 2024). Instead of recovering sequences from mask tokens, some discrete diffusion models learn to recover from uniformly sampled sequences (Lou et al., 2023). Another thread of work incorporates autoregressive or cemi-autoregressive biases to the denoising process of diffusion models, intending to combine the expressiveness of autoregressive modeling and the ability to perform NAR generation (Chen et al., 2024; Han et al., 2023).
6

---

Tractable Transformers
(a)
CAR (Coometrieeg 50,)
(b)
CAR (Grometrk(u=10)) 10:)
(c)
AC (Gommntrio(g
(d)
AC (Geomeettee 10))
Conditional Perplexity
20
21-
Tracfforer (Ours)
BART
Conditional Perplexity
22
20-
Traeformer (Ours)
BART
Conditional Perplexity
800
900
BERI-bidir
BERT
Tracformer (Ours)
Conditional Papleexity
400
300
BERT
Tracformrr (Ours)
BERT-bidir
19
18-
700
200
16
600
18-
14
500
100
17
12
400
0
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Mask ratio
Mask ratio
Mask ratio
Mask ratio
Figure 5: Validation conditional perplexity (4) of CAR/AC generation tasks with varying masking strategies, evaluated on the 10 validation set with sequence length of 1024. (a) and (b) show CAR generation with span lengths sampled from Geometric(j 50) and Geometric(I 10). respectively. (c) and (d) correspond to AC generation with the same span length distributions in CAR. Across all mask ratios and span lengths, Tracformer consistently outperforms all baselines, demonstrating strong generalization capabilities.
Table I: Evaluation of CAR text infilling performance using the MAUVE and BERT Score with different mask ranges. Higher scores indicate better performance. Tracformer consistently out- performs BART under conditions.
AC generation tasks (see Appx. B.2 for details).
Training Setup. For the CAR task, Tracformer and BART are both trained using the CAR objective defined in Equa- tion (7). During training, we set the mask strategy (i.e.. Pc) to sample spans whose lengths follow a geometric dis- tribution with mean μ 50 and apply a total mask ratio of 50% (see Appx. B. All models are trained on the WikiText103, LAMBADA, and IBW. We present the results of WikiText103 in the paper and defer other results to Ap- pendix B.3. The experiment setup for the AC generation task follows the CAR task, where the training objective is defined in Equation (8). See Appendix B.2 for more details.
MAUVE (1)
BERT Score (1)
Mask ranges
Tracformer BART
Tracformer BART
[0.250.75]
0.960
0
0,464
0.414
[0.5.1.0]
0.114
0.016
0.488
0.386
[0.10.4] & [0660.99
0.931
0.889
0,370
0.306
[004]& (0.50..)
0.946
0.940
0,274
0.216
[0025]] (0.75,100)
0.177
0.063
0.457
0.401
10,0.1|& & (0.2,0.5 & [0.7.1.0]
0.107
0.023
0.286
0.208
6.1. Comparison of Architectures for NAR Modeling
Empirical Insights from the CAR Generation Results. We evaluate both models' generalization capabilities on CAR tasks through two specific tests: ratio generalization and span generalization. For ratio generalization evalua- tion, the span masking strategy remains consistent with the training setup, where the span length is sampled from Geometric(j: = 50), but the total mask ratio is varied be- tween 0.1 and 0.9. Figure 5(a) shows the conditional per- plexity (PPL) for both models on the WikiTeext13 valida- tion set. The results exhibit a U-shape, where performance improves as the mask ratio approaches the training mask ratio. Tracformer consistently outperforms BART across all mask ratios. The performance gap between the two models is small near the training mask ratio, but BART's PPL increases as the mask ratio deviates further, whereas Tracformer maintains robust generalization performance.
Recall from Section 4 that we propose Tracformer's s encoder- decoder architecture as an effective approach to leveraging both local and global features, enabling robust performance across varying queries. To empirically validate this, we com- pare Tracformer with baseline models employing different transformer architectures on identical NAR modeling tasks, focusing on their conditional generalization capabilities.
Baselines. As discussed in Section 4.3, Tracformer sup- ports two distinct NAR generation modes: contextual AR (CAR) generation and arbitrary-context (AC) generation, each requiring different baseline comparisons. For CAR generation, we use BART (Lewis et al., 2020) as a baseline, as it represents a widely adopted encoder-decoder architec- ture for CAR tasks. BART's encoder captures bidirectional global context, while its autoregressive decoder generates outputs sequentially via cross-attention. This fundamen- tal design principle is shared by many CAR models, such as MASS (Song et al., 2019) and PALM (Bi et al., 2020), making BART a strong representative baseline.
To evaluate span generalization performance, we alter the mean span lengths of the adopted span masking strategy during inference, resulting in a mismatch Pe between train- ing and evaluation. Figure 5(b) illustrates the conditional PPL when the mean span length is reduced to 10. As ex- pected, shorter span lengths reduce the overall task dif- ficulty, resulting in lower PPL for all models. However, Tracformer demonstrates significantly better performance improvements compared to BART, highlighting its ability to generalize effectively to unseen masking patterns. We also include Geometric(j = 3) span mask results in Figure 7(a)
For AC generation, we compare Tracformer against two baselines: (i) BERT (Devlin et al., 2018), which is the pre- dominantly used backbone for modern diffusion models, and (ii) BERT-bidir, a variant of BERT that incorporates both forward and reverse AR encoders using dense self- attention. These baselines provide a contrast between stan- dard bidirectional encoding and autoregressive modeling in
7

---

Tractable Transformers
Table 2:
Zero-shot conditional perplexity on WikiText103 and IBW using six fixed prompts. Tracformer operates in the CAR generation mode, while the conditional PPLs of SEDD and MDLM are computed following the derivation in Appendix G. Tracformer consistently outperforms state-of-the-art discrete diffusion models SEDD and MDLM across various masking ranges.
Table 3: Model size, number of non- eoneemidd parameters, and
number of trained tokens for the models evaluated in this paper. Tracformer has smaller full model size and fewer non-embedding parameters compared to SEDD and MDLM.
Model
Full Model Size
Embed Params
#Trained Tokens
GPT2
124M
85M
WikiText103
IBW
SEDD
169M
90M
210B
Mask ranges
MDLM
169M
90M
524B
Tracforner SEDD MDLM
Tracformer SEDD MDLM
Tracformer
109M
79M
295B
[0.25,075]
29.38
37.83
30.36
41.98
53.72
55.19
[0.10.4] & (0660.9)
29.67
37.17
30.55
38.52
49.38
49.48
Table 4: Zero-shot unconditional perplexity across various datasets, Tracformer, despite its smaller size, achieves competitive or supe- rior performance compared to larger models.
[0,04] & (0.5.088)
34.73
43.86
35.84
45.56
58,22
58.41
[0,0225] & (0.75.1)
37.73
46.65
40.11
49.14
66.96
66.16
[0.20033 & (0.4.0.6)
24.31
28.31
22.97
31.76
38.76
38.97
[0.1,09]]
33.88
44.92
36.69
43.98
57.96
59.14
Model
WikiText103
WikiTexx2 Lambada
PTB
1BW
GPT-2 (124M)
41.60
42.32
45.04
138.43
75.20
of Appendix B.3, where a similar trend is observed.
DSPM (169M)
75.16
77.28
93:47
200.82
138.92
PLAID (169M)
50.86
51.80
57.28
142.60
91.12
Beyond perplexity, we assess the quality of generated text in CAR tasks using six fixed mask ranges (i.e., text infilling). Metrics such as MAUVE (Pillutla et al., 2021) and BERT Score (Zhang et al., 2020) are used to quantify the similarity between the generated and original text. See Appendix B.3 for further evaluation details. Table 1 presents the results, where Tracformer consistently achieves higher scores than all baselines, demonstrating Tracformer's s ability to produce high-quality text under diverse conditions.
SEDD (169M)
40.62
41.84
50.92
114.24
79.29
MDLM (169M)
37.01
36.75
48.46
96.40
67.94
Tracformer (109M)
43.27
43.82
58.10
166. 10 51.34
two span length distributions: Geometric(jr = =10) and DLogistic = 15, a = 3). On IBW, Tracformer consis- tently outperforms both SEDD and MDLM across all mask ratios. On WikiText103, MDLM shows a slight advan- tage at low mask ratios, but as the masking ratio increases, Tracformer achieves the best performance, demonstrating stronger generalization under more challenging conditions.
Empirical Insights from the AC Generation Results.
As shown in Figure 5(c) and (d), Tracformer achieves sig- nificantly lower conditional perplexity than both baselines across all mask ratios and span lengths, demonstrating supe- rior generalization performance consistently. Additionally, we conduct an ablation study on Tracformer's multi-scope attention, with detailed analysis provided in Appendix D.
Notably, despite its smaller model size (Tab. 3), Tracformer not only excels in fixed-mask conditional PPL (Tab. 2) but also adapts better to diverse masking distributions.
For unconditional generation, Table reports the zero-shot unconditional PPL across various datasets. Tracformer re- mains highly competitive, achieving results comparable to or better than larger models. We additionally include text samples generated by Tracformer in Appendix F for both conditional and unconditional generation. These findings further reinforce our earlier conclusion: while SoTA NAR diffusion models exhibit strong unconditional PPL, their in- ability to generalize effectively to conditional queries unseen during training limits their broader applications. Tracformer, with its robust performance in both conditional and uncon- ditional tasks, demonstrates its potential as a scalable and versatile generative model.
6.2. Benchmarking Against SOTA NAR Models
In this section, we test the scalability of Tracformer by training it on the OpenWebTexx (Gokaslan & Cohen, 2019) dataset and comparing its performance against various SoTA generative models of similar scale, including GPT-2 (Rad- ford et al., 2019) and SoTA discrete diffusion models (Lou et al., 2023; Sahoo et al., 2024; Austin et al., 2021; Gul- rajani & Hashimoto, 2024). Tracformer is trained using the CAR generation objective defined in Equation (7), with similar masking strategies Pe to Section 6.1 (see Appx. C). For evaluation, we benchmark zero-shot conditional and
unconditional performance on multiple datasets.
7. Conclusion and Limitations
For conditional generation, we compare Tracformer with SoTA discrete diffusion models, SEDD and MDLM. Table 2 presents the zero-shot conditional PPL results using fixed masking strategies, where Tracformer consistently outper- forms both basclines across all mask ranges. To further assess generalization, we evaluate varied masking strategies, with Figure 6 reporting results from aatio-generalization and gpan-generalization tests. Specifically, we consider
We propose Tracformer, a Transformer-basee architec- ture for flexible and generalizable conditional generation. Through extensive experiments, we demonstrate that Trac- former's 's multi-scope attention mechanism and specialized encoder-decoder design enable robust conditional genera- tion performance. However, due to resource limitations, we only train Tracformer at the GPT-2 (base) scale, While
8

---

Tractable Transformers
WikiText103 (Gumerzielp 10))
WikiText103 (Digihtete (6,0=3)
1BW (Gommtrielp
1BW (DIogisti(p=11..05)))
45
45
60
Conditional Perplexity
40
SEDD
MDLM
Conditional Perplexity
40
MDLM
SEDD
Conditional Perplexity
50
SEDD
MDLM
Conditional Perplexity
MDLM
SEDD
35
Tracformer (Ours)
35
Trneformer (Ouns)
Tieffomer (Ours)
50
Traeformer (Ours)
30
25
30
40
20
25
15
20
30
20
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
01.4
0.6
0.8
0.2
0.4
0.6
0.8
Mask ratio
Mask ratio
Mask ratio
Mask ratio
Figure Conditional perplexity under varied masking strategies, evaluating aetorageneraliztiion and span ppan-ennrralization We evaluate two different span length distributions: Geometric( = 10) and DLogistic(je = 15, a = 3). Tracformer demonstrates superior conditional generalization capabilities compared to SEDD and MDLM across different masking strategies.
our current results already establish Tracformer highly promising NAR architecture, future work will focus on scal- ing the model further to fully explore its potential. Moreover, given that AC generation closely aligns with one-step diffu- sion training, Tracformer could serve strong backbone for modern diffusion models. Future research will investi- gate this direction, leveraging Tracformer's capabilities to enhance diffusion-based generative frameworks.
Den Berg, R. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34: 14411991-1799,, 2021. 3, 6,
Bi, B., Li, C., Wu, C., Yan, M., Wang, W., Huang, S., Huang, F., and Si, L. Palm: Pre-training an au- buucooddnnkattreegreesss language model for context- conditioned generation. arXiv preprint A5X 2020. 7
Impact Statement
Brown, T. B., Mann, B., Ryder, N., et al. Language models are few-shot learners, Advances in Neural Information Processing Systems, 33:1877-1901 2020. 1,2
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Campbell, A., Benton, J., De Bortoli, V., Rainforth, T., Deli- giannidis, G., and Doucet, continuous time frame- work for discrete denoising models. Advances Neural Information Processing Systems, 35:28266-2822 2022.
Acknowledgements
6
This work was funded in part by the National Sci- ence and Technology Major Project (2022ZD01 14902), the DARPA ANSR program under award FA8750-23- 2-0004, the DARPA CODORD program under award HR001 12590089, the DARPA SAFRON program under the Deutsche Forschungsgemein.
Chelba, C., Mikolov, T., Schuster, M., Ge, Q., Brants, T., Koehn, P., and Robinson, T. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint 1r12v:1312.3003 2013. 6
award FRR0011253014 schaft (DFG, German Research Foundation) under Ger- many's ' Excellence Strategy EXC 2075 390740016, and gifts from Adobe Research, Cisco Research, and Amazon. We acknowledge the support of the Stuttgart Center for Simulation Science (SimTech)). MN thanks IMPRS-IS (International Max Planck Research School for Intelligent Systems) for the support. Approved for public release; distribution is unlimited.
Chen, B., Monsó, D. M., Du, Y.. Simchowitz, M., Tedrake, R., and Sitzmann, Diffusion forcing: Next-token pre- diction meets full-sequence diffusion. In The Thirty- eighth Annual Conference Neural Information Pro- cessing Systems, 2024. 1,6
NSF grant
DaSilva, F., Senan, S., Patel, Z. M., Reddy, A. J., Gabbita, S., Nussbaum, Z., Córdova, C. M. V., Wenteler, A., We- ber, N., Tunjic, T. M., et al. Dna-diffusion: Leveraging generative models for controlling chromatin accessibility and gene expression via synthetic regulatory elements. bioRxiv, 2024. 1
References
Alamdari, S., Thakkar, N., van den Berg, R., Tenenholtz, N.,
Strome, B., Moses, A., Lu, A. X., Fusi, N., Amini, A. P., and Yang, K. K. Protein generation with evolutionary diffusion: sequence is all you need. BioRxiv, pp. 2023-09, 2023. 1
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. Proceedings of NAACL-HLT, pp- 4171-4186. Association for Computational Linguistics, 2018. 3, 6,7
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van

---

Tractable Transformers
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21788, 2024.
Li, Z., Ni, Y., Beardall, W. A., Xia, G., Das, A., Stan, G.-B., and Zhao, Y. Discdiff: Latent diffusion model for dna sequence generation. arXiv preprint xxX::2442.006099 2024.
Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., and Re, C. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh Inter- national Conference on Learning Representations, 2022.
Liu, A., Broadrick, O., Niepert. M., and Van den Broeck, G. Discrete copula diffusion. The Thirteenth International Conference on Learning Representations, 2025.
Liu, F., Fu, Z., Li, G., Jin, Z., Liu, H., Hao, Y., and Zhang, L. Non-autoregressive line-level code completion. ACM Transactions on Software Engineering and Methodology, 2024.
Gokaslan, A. and Cohen, V. Openwehtext corpus. http:// Skylion007. www.m/00...ttuu.... www.gg.ttt.ooooo hwwabb///o/peeeeeeetttrrpppss 2019.
6,
Lou, A., Meng, C., and Ermon, Discrete diffusion lan- guage modeling by estimating the ratios of the data distri- bution. arXiv preprint arXiv:2310.1684, 2023. 1, 3,
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint 12Xiv:3322.00752. 2023.
6, 8, 12, 16
Gu, S., Chen, Y., Bao, H., et al. Diffusion-lm im- proves controllable text generation. arXiv preprint arXiv:2208. 13753, 2022. 16
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference on Learning Representations, 2022. 2, 3,
Gulrajani, I. and Hashimoto, B. Likelihood- diffu- sion language models. Advances Neural Information Processing Systems, 36, 2024.
Nguyen, E., Poli, M., Faizi, M., Thomas. A., Wornow, M., Birch-Sykes, C., Massaroli, S., Patel, A., Rabideau, C., Bengio, Y., et al. Hyenadna: Long-range genomic se- quence modeling at single nucleotide resolution. Ad- vances neural information processing systems, 36, 2024. 2
Han, X., Kumar, S., and Tsvetkov, Y. Ssd-Im: Semi- autoregressive complex-basee diffusion language model for text generation and modular control. In Proceed- ings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 11575-11596, 2023. 2, 16
OpenAL. Chatgpt (gpt-4), 2023, URL https://chat. openai... com. 2
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, Q. N., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint arXiv:1606.066 06031, 2016.
A., Visonà, G., Narendra, T., Rojas- Carulla, M., Schölkopf, B., and Schweikert, G. Getting personal with epigenetics: towards individual-specific epigenomic imputation with machine learning. Nature Communicationns, 14(1):4750, 2023.
Pillutla, K., Swayamdipta, S., Zellers, R., Thick- stun, J., Welleck, S., Choi, Y., and Harchaoui, Z. MAUVE: Measuring the gap between neural text and human text using divergence frontiers. In Ad- vances Neural Information Processing Systems, volume 34, pp. 4816-4828, 2021. URL https: /
Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv: (xXiv:666.0844 08466, 2016. 12
Kaddour, J., Harris, J., Mozes, M., Bradley, H., Raileanu, R., and McHardy, R. Challenges and applications of large language models. arXiv preprint arXiv:2307.10066 2023.
www.wnw..iii.c.. http://reepppepppppp//////////// #55071906edd48e6244452220000dee0 pdf.
8
Lewis, M., Liu, Y., Goyal, N., al. Bart: Denois- ing requence-to-sequenee pre-training for natural lan- guage generation, translation, and comprehension. arXiv preprint
Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Heek, J., Xiao, K., Agrawal, S., and Dean, Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5:606-624, 2023. 6
2020. 2, 6.7
Li, X., Thickstun, J., Gulrajani, I. Liang, P. S., and Hashimoto, T. B. Diffusion-Im improves controllable text generation. Advances in Neural Information Process- ing Systems, 35:4328 4344, 2022. 2
Radford, A., Wu, J., Amodei, D., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019. 1, 2,
10

---

Tractable Transformers
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017. URL http:://p//// papers httpppp.../pppe...eeeaapppe////a 3fee:/66776601110000ccoaa66aaa .df..
2
Reid, M., Savinov, N., Teplyashin, D., Lepikhin, D., Lilli- crap, T., Alayrac, J.-b., Soricut, R., Lazaridou, A., Firat, O., Schrittwieser, J., et al. Gemini 1.5: Unlocking multi- modal understanding across millions of tokens of context. arXiv preprint ar arXiv:2403.055330 2024.
Wan, K., Mu, H., Hao, R., Luo, H., Gu, T., and Chen, X. A cognitive writing perspective for constrained long-form text generation. arXiv preprint arXiv:2502.12550 120.228822.12588, 2025.
5
Warner, B., Chaffin, A., Clavié, B., Weller, O., Hallström, O., Taghadouini, S., Gallagher, A., Biswas, R., Ladhak, F., Aarsen, T., et al. Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference. arXiv preprint arXiv: 236.:3311,,363,, 13663, 2024.
Sahoo, S. S., Arriola, M., Schiff. Y., Gokaslan, A., Marro- quin, E., Chiu, J. T., Rush, A., and Kuleshov, V. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.07524, 2024. 2, 3, 6, 8, 12, 16, 20, 21, 22
Shi, J., Han, K., Wang, 7., Doucet, A., and Titsias, M. K. Simplified and generalized masked diffusion for discrete data. arXiv preprint arXiv:2406.044 04339, 2024. 1,6
Xu, M., Geffner, T., Kreis, K., Nie, W., Xu, Y., Leskovec, J., Ermon, S., and Vahdat, A. Energy-based diffusion language models for text generation. In The Thirteenth International Conference on Learning Representations, 2025.
Shin, J.-E., Riesselman, A. J., Kollasch, A. W., McMahon, C., Simon, E., Sander, C., Manglik, A., Kruse, A. C., and Marks, D. S. Protein design and variant prediction using autoregressive generative models. Nature communica- tions, 12(1):2403, 2021.
Zhang, T., Kishore, V., Wu, F,, Weinberger, K. Q., and Artzi, Y. BERTScore: Evaluating text generation with BERT. International Conference on Learning Rep- resentations, 2020. URL https: https://openrevien nett/ forum?id-SkeHuUVFDD 8
Song, K., Tan, X., Qin, T., Lu, J., and Liu, T.-Y. Mass: Masked sequence to sequence pre-training for language generation. arXiv e-prints, pp. arXiv-1905, 2019. 6,
Stark, H., Jing, B., Wang, C.,, Corso, G., Berger, B., Barzilay, R., and Jaakkola, T. Dirichlet flow matching with applica- tions to dna sequence design. In Forty-first International Conference on Machine Learning, 2024. 1
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 168:127063, 2024. 12
Sun, H., Yu, L., Dai, B., Schuurmans, D., and Dai, H. Score- based continuous time discrete diffusion models. In The Eleventh International Conference on Learning Repre- sentations, 2023. 6
Tanaka, K., Kato, K., Nonaka, N., and Seita, J. Efficient hla imputation from sequential snps data by transformer. arXiv preprint 2:XX::2211.0643333 2022. URL https: //arxiv org/ aaaaaba//200 /21s./ 06430.
Trinquier, J., Uguzzoni, G., Pagnani, A., Zamponi, F., and Weigt, M. Efficient generative modeling of protein se- quences using simple autoregressive models. Nature communications, 12(1):5800, 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, t.,, and Polosukhin, Attention

---

Tractable Transformers
A. Additional Details of Tracformers
This section introduces the design details of Tracformers.
Sparse Attention Masks of the Encoder. As described in Section 4.3, a prefix encoder and a suffix encoder are used for both CAR and AC generation. Specifically, the attention mask for each layer is the sparsified version of Equation (4) such that each token attends to at most N tokens. Formally, the attention map m of layer I E [4] in the prefix encoder is given by (denote = -- 2/-1 d
=bbth t):
if St otherwise,
2- 2 1}. (a,a+ d.a+ 63 [T]
if
N Mmax+
ene. 2
where S
0
2
Nmax
As shown in Section 4.3, each feature n in the suffix encoder covers the 2 succeeding variables starting from X₁:
= :t & sT,
which is analogous to Equation (3). Similar to the prefix encoder, we use the following sparse attention mask to implement the suffix encoder (denote = t, d Naw = t 21-1):
1 if E St where S otherwise,
ST,, 2f2 < N if 2/211 1 > Nmax
enc.l ma,e'
[a
d, a.a+
Sparse Attention Masks of the Decoder. As discussed in Section 4.3, the decoder cross-attends to features in both the prefix encoder and the suffix encoder, using the mask in Equations (5) and (6), respectively.
Encoder Blocks. Each encoder block contains an attention layer and an FFN following the standard design. Skip connections and layer normalization layers are used.
h h hene. = Lenc!--1 + Attn; ALtneeayyrrormm, (hT.r/ (/..c.--1))
1:T h ene.l hened FFN; (LayerNorm, (h:med))
In the attention layers, we adopt the Rotary positional encoding (Su et al., 2024) to encode information about relative positions between feature embeddings. This is widely adopted in NAR models such as discrete diffusion models (Lou et al., 2023; Sahoo al., 2024).
For the FFN, we follow the design of GPT-2 and a two-layer fully connected neural network with GeLU (Hendrycks Gimpel, 2016) activation. The input and output dimensions are both the embedding dimension d and the latent dimension size is 4 x d.
Decoder Blocks. Each decoder block consists of a cross-attention layer and an FFN:
h des,1 1:7
100oo doc,1-1 + @Crosstttu((Layeekoomm.
(h (h LayerNorm, (1111111+1)))
h dec.1 h dea,
+ (h deee..))).
1:T
1:T
Similar to the encoder, also adopt the Rotary positional encoding in cross-attention layers. We set the input features to the decoder to zero (i.e., dec.0
0). The FFNs are the same as in the encoder blocks.
B. Comparison of NAR Architectures in Small-Scale Experiments
B.1. Training Mask Strategy
The span masking strategy Pe used in our experiments implemented using the following algorithm:
12

---

Tractable Transformers
Algorithm Span Masking Strategy
1: Input: Sequence length L, mask probability p, mean span length m, distribution type (Geometric or DLogistic),
2: Initialize num_to to_mask max(1, round L)), num_mosked 0, blank ids []. spans
3: while num_masked < num_to, to__tssk do
4: Sample a random start index in [0, L)
5: if start overlaps with an existing span in spans then
6: continue
7: end if
8: if distribution == Geometric then
9: Sample span praa_lenghh Geometric(u m)
Geometric distribution with mean
10: else if distribution = DLogistic then
11:
Sample span_ength ~~ DLogistic( m,, = o))
Dlogistic distribution with mean and std
12:
end if
13:
end- = min(start + span span_length, L)
14:
if Overlap with existing spans detected then
15:
continue
16: end if
17:
Append (slart, end) to spuns, update dank_ids and num_masked end while
18:
19:
if num
num_masked masked > num__t fonm_tt__r_o to_mask then
20: Trim the last span to ensure exact masking ratio
end if
Return: spans, blank ids
Following Equation (7), the CAR training process involves applying the span masking strategy to corrupt input sequences and training the model to predict the original tokens based on the provided context. Specifically, at each training iteration, mask function used create masked spans within the input sequence. The masking process ensures that 50% of the tokens (based mask probability of 0.5) are selected in spans, where the span length is sampled from either a Geometric or DLogistic distribution. For training in Section 6.1, we use the Geometric distribution with mean span length of μ 50. The corrupted input sequence, along with mask indicating the positions of the masked tokens, is then passed to the model. The model is trained to minimize the negative log- likeehood of the original tokens at the masked positions, conditioning on the unmasked tokens as context.
the generate_spee
B.2. Model and Training Configurations
Training Hyperparameters. For both CAR and AC training tasks, the sequence length is set to 1024 tokens, with a batch size of 256. The models are optimized using Adam with 83 = 0.9, 32 = 0.95, and a weight decay of 0.1. The initial learning rate is set to 10 and follows a cosine decay schedule, with 1,000 warmup steps to stabilize the early training phase. The final learning rate is 10-5 Training is conducted for 30,000 steps.
Hyperparameters of Tracformer. As listed in Table 6, for both CAR and AC training tasks, Tracformer' is implemented with 10-layer encoder-decoder architecture, maintaining block size (i.e., maximum sequence length) of 1024 tokens. It utilizes sparse multi-scope attention with a constraint of 16 attended tokens per step, allowing for efficient context aggregation while keeping computational cost manageable. The decoder operates with a maximum stride of 1024 tokens, ensuring global context encoding. The model is configured with 9 attention heads and an embedding dimension of 576. dropout rate of 0. is applied to mitigate overfitting during training.
Hyp Hyperparameters of the Baseline Models. As the CAR baseline, BART adopts an 8-layer encoder-decoder architecture with 9 attention heads and embedding dimension of 576. This leads total of 138M parameters.
For the AC generation experiments, we use BERT as the encoder-only baseline with full attention, employing 10 layers, 12 attention heads, and an embedding dimension of 768. This configuration ensures a fair comparison by maintaining a similar model size to Tracformer. Additionally, we include BERT-bidir, which encodes bidirectional context using two separate encoders: a classical forward AR encoder with causal attention and a reverse AR encoder. While structurally similar to Tracformer's 's prefix and suffix encoders, BERT-bidder relies on dense causal attention and lacks the sparse multi-scope self-
13

---

Tractable Transformers
Table Training hyperparameters for CAR and AC tasks.
Parameter
Value
Table 6: Hyperparameters of Tracformer the Small-scale Ex- periments.
Training Semup
Sequence length
1024 tokens
Parameter
Value
Batch size
256
Structure Parameters
Training steps
30,000
Block size
1024
Optimizer
Number of encoder layers
10
Optimizer
AdamW
Number decoder layers
10
31
0.9
Max attended tokens (Nvax)
16
B2
0.95
Attention Parameters
Weight decay
0.1
Number attention heads
9
Learning Rate Schedule
Embedding dimension (d)
576
Initial learning rate
6 x 10-4
Other Parameters
Final learning rate
6 x 10-5
Dropout rate
0.1
Learning rate decay
Cosine schedule 1,000
Warmup steps
Table 7: Comparison of the model size for small-scale experiments.
Model
Full Model Size
Tracformer
109M
BART
138M
BERT
149M
BERT-bidir 170M
attention mechanism that enhances Tracformer's ' efficiency in handling conditional generation. BERT-bidder is configured with 12 layers, 12 attention heads, an embedding dimension of 768, and dropout rate of 0.1. After encoding, the final-layer features from both encoders are concatenated and passed through a two-layer MLP to produce the final representation, allowing the model integrate information from both directions before prediction. The MLP consists of two fully connected layers (config.n_emll embd is the embedding size of the model):
mlp nn. Mm.MoouleLList( hawara/iinnmmmiiiiaoooonnnnnnnn x 2, config.n_embll x 4), GELUO;
Mm.Linea(cconfignn.mm embd x 4, .mnigd, GELUO))
Table 7 shows the parameter count of the adopted baseline models.
B.3. Additional Evaluation Details and Results
Conditional PPL under Short Span Masking. Figure 7 contains complementary results of Figure 5.
Metrics for CAR Text Infilling. For the CAR text infilling task, a sequence length of 128 is used. To compute MAUVE and BERT scores, we evaluate 1,000 text sequences from the mikiText-103 validation set across all methods. Each prompt generates five samples, resulting in 5,000 generated sequences.
CAR Text Infilling Results on Other Datasets. given Table 8.
Text-infilling performance on the 1BW and LAMBADA datasets are
14

---

Tractable Transformers
Geomstrie
CAR Goometrk(u=3)
1000
Conditional Perplexity
BERT
Conditional Perplexity
20
BART
800
BERT-bidir
Tracfonner (Ours)
Tracformer (Ours)
600
15
10
200
0.2
0.4
0.6
0.8
0.4
0.6
0.8
Mask ratio
Mask ratio
Figure 7: Conditional perplexity (PPL) evaluated on the WikiText- validation set. The span length follows a Geometric(j = 3) distribution, while the mask ratio varies between 0.1 and 0.9. The sequence length is set to 1024, which aligns with the training setup. (a) Results for the CAR task. (b) Results for the AC task. Tracformer consistently outperforms all baselines.
Table 8: CAR Text infilling performance evaluated IBW and Lambada Datasets.
MAUVE (t)
BERT Score (1)
Mask ranges
Tracformer
BART
Tracformer
BART
IBW Dataset
[0.25.0:75]
0.974
0.975
0.471
0.446
[0.50,1000]
0.421
0.384
0.478
0.476
10.10.0.40] [0.60,0.90]
0.992
0,979
0.377
0.334
[0.00.044]] [0.50,0.80]
0.987
0.977
0.272
0.230
[0.00,0.2]] [0.75,1000]
0,379
0.375
0,450
0,425
[0.00,0. 10] (0:20.0.50) (0.70,1.00 1.00]
0.346
0.364
0.270
0.231
Lambada Dataset
10.25.0.751
0.897
0,913
0.301
0.289
[0.50.1
0.024
0.016
0.344
0.320
[0.10,0.40] & [0.60,0.90]
0.714
0.693
0.178
0.153
[0.00.0.40] [0.50,0.80]
0.804
0.839
0.077
0.056
[0.00,0.25] [0.75.1.00]
0.037
0.024
0.322
0.309
[0.00.0.10] 101 (0.20.0.50] & (0.70.1.00)
0.023
0.014
0.068
0.041
C. Details of Large-Scale Experiments on WebText/OpenWebTeei
C.1. Training Mask Strategy
To scale Tracformer to OpenWehTer while balaneing both conditional and unconditional generation performance, adopt mixed masking strategy. This approach integrates contextual autoregressive (CAR) training with an autoregressive (AR) objective by varying the masking strategy probabilistically during training.
At each training iteration, we select from the following masking strategies:
Unconditional training (30%): No context is provided, and the model learns in fully autoregressive (AR) manner, simulating traditional language modeling.
High-span masking (20%): A high masking probability (85%) is applied, forcing the model to rely on limited observed tokens to reconstruct missing content.
Moderate-span masking (50%) ) standard masking probability (50%) is used, following the CAR objective.
The generate_spe sewerate__pa__.ee span_mask function, as defined in Algorithm 1, is used when applying high- and moderate-span masking strategies to ensure structured token corruption. However, in unconditional training, no span masking is applied, effectively reducing the task to standard autoregressive modeling. This hybrid approach allows Tracformer to retain strong unconditional generation capabilities while improving its ability generalize to arbitrary conditional queries.
15

---

Tractable Transformers
Table CAR Text infilling performance (BLEU-4) evaluated on Wik
BLEU-4 (1)
Mask ranges
Traclomme
BART
[0.25,0.75]
0.524
0.513
[0.50.1.00]
0,540
0,519
[0: 10,0..04 & (0.60,090)
0.419
0.405
[0.00.0.40] (0.50,0.80]
0.339
0.325
[0.00,0.25] & [0.75.1000]
0.536
0,523
[0.30.0.10] [0.20,0.50] & (0.70,1.00)
0.337
0.322
The formalized procedure the mixed masking strategy presented Algorithm 2.
Algorithm Mixed Masking Strategy for WeenWeTText Training
Input: Input sequence x of length L 1024, mask token ID mid
2: Sample r Uniform (0,
3: if r 0.3 then
Unconditional training (AR)
4: Xcorrupt mia (fully masked)
5: else if < 0.5 then
> High-span masking (CAR)
6: spans, blank_ianlcccaaa _dd = generate_span_mask/ 1024, p = 0.85, m = 50, distribution type = Geometric)
7: else
b. Moderate-span masking (CAR)
8:
spans, blank ids = generate_span_ 1024, p = 0.5, m = 50, distribution type = Geometric)
9: end if
10: ccorrupt x.clone()
11: Xcorrupt bblank_iddd ids] mid
12: Return: Xcorrupt
C.2. Model and Training Configuration
Tracformer. Details of the architecture and optimization procedure for Tracformer are detailed in Table 10.
SEDD. We use the SEDD-small model with 169M parameters (including 90M non-embedding parameters) trained on WeenWeeText. The model is accessed through HuggingFace: bttps://hugging: Following the original paper (Lou al., 2023), adopt the log- linelinear noise schedule and the absorbing mask forward noising process.
wo/louaaron/sedda small.
We use the MDLM model with 169M parameters (including 90M non-embedding parameters) trained on OpenWebText. The model is accessed through HuggingFace: https://huggiin ace. co/kuleshov- Following the original paper (Sahoo et al., 2024), we adopt the log-linear noise schedule.
MDLM.
GPT. We use the GPT-2 small model with 124M parameters (including 85M non-embedding parameters). The model obtained from HuggingFace: https: httpp://huggingface wtt.aaaccc...cooeeeeeeeeee http:/ppnii-comiiittteeeee ty/ /gpt2.
C.3. Evaluation Metrics
For all conditional generation experiments in Section 6.2, we use sequences of length 128 following prior work (Lou et al., 2023; Han et al., 2023; Gu et al., 2022). Additionally, we exclude EOS tokens when evaluating conditional perplexity since the focus of these tasks is the ability of different models generate coherent text given prompt texts.
D. A Additional Ablation Studies
To further evaluate the impact of Tracformer's multi-scope attention, we conduct an ablation study by replacing it with full causal attention. In this modified architecture, the prefix encoder employs classical dense causal self-attention, while the suffix encoder uses reverse causal attention. where each token attends to all subsequent tokens within the layer. Apart from this modification, the encoder-decoder structure remains unchanged, ensuring a fair comparison.
16

---

Tractable Transformers
Table Hyperparameters of Tracformer the large-scale Tex Experiments
Parameter
Value
Structure Parameters
Block size
1024
Number of encoder layers
10
Number of decoder layers
10
Max attended tokens (Nvmax)
32
Attention Parameters
Number of attention heads Embedding dimension (d)
576
Other Parameters
Dropout rate
0.0
Optimization Parameters
Initial learning rate
6e-4
Terminal learning rate
6e-5
Learning rate schedule # warmup steps # training steps
cosine
2,000
600,000
Batch size
480
Weight decay
lc-1
Adam betas
(0.9, 0.95)
Gradient clipping maximum norm
1.0
Tables 11 present the AC generation performance under two different masking strategies: Geometric(50) with mean span length of 50, aligned with the training mask length, and Geometric( 10) with mean span length of 10, which deviates from the training distribution. The results show that low mask ratios (i.e., when more context is available), both models achieve comparable performance, with the ablation model exhibiting slight advantage. However, the mask ratio increases and the task becomes more challenging, Tracformer with multi-scope attention significantly outperforms the ablation model, demonstrating superior generalization under high-uncertainty conditions.
Another key observation is that replacing multi-scope attention with full attention introduces a notable efficiency drop. The ablation model incurs higher computational costs due the dense attention mechanism, making it less scalable compared Tracformer. This highlights the practical advantage of Tracformer's sparse multi multi-scoppe scope self-attention, which achieves favorable balance between performance and efficiency.
Finally, it is worth noting that despite replacing multi-scope attention with full attention, the ablation model still significantly outperforms the previous baseline, BERT-bidir, which has a similar number of parameters. This underscores the importance Tracformer's specialized mncoder-decoder design, beyond just its attention mechanism. The architectural improvements introduced in Tracformer contribute substantially to its ability to handle AC generation tasks, making it a more effective approach compared to conventional NAR models.
E. Training and Inference Efficiency of Tracformers
Training Efficiency. As shown in Section 4.1 and Appendix A, given Tracformer with T. layers, input length T, and Nmax- the computation cost of the encoder is O(L T' Nm which is linear w.r.t. T'. We note that a custom kernel implementation needed to fully take advantage of this sparse attention and leave it to future work.
According to the decoder attention mask (Eqs. (5) and (6)), the number of tokens required to attend is exponentially
17

---

Tractable Transformers
Table Ablation results
CAR generation task. Tracformer w/ full encoder an ablation model where the multi-scope encoder the original Tracformer is replaced with dense self-attention while maintaining the encoder- structure. The table presents results for two span length distributions: Geometric(μ = 50) and Geometric( = 10), evaluated under different mask ratios.
Mask Ratio
Model
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
Geometric(j = 50)
Tracformer
445.46
467.03
478.13
475.65
459.17
460.53
447.84
449.77
564.32
Tracformer w/ full encoder
443.65
462.90
475.25
473.85
458.77
462.30
452.11
455.23
574.73
Geometric(j = 10)
Tracformer
165.07
178.77
181.89
178.79
185.38
193.40
219.13
288.91
509.35
Tracformer w/ full encoder
160.39
174.64
177.88
176.36
183.01
191.53
217.69
288.45 48.45
519.73
smaller for initial decoder layers. Suppose L = Oloo T), then the total computation cost the decoder is
Σ O(T 2() + O(T O 0 -.
FPNS
Cross attention layers
This improves upon the cost O(T L) in classic decoder-only Transformers.
Inference Efficiency. We give an example of using KV-caching techniques in Tracformers designed for CAR generation (i.e., Fig. 4(a)). Given context Ic, we first run the suffix encoder for all tokens by feeding the <MASK> token to every token I g C. We then autoregressively decode the remaining tokens following the order æ first compute the (1-1)-th feature in each layer of the prefix encoder (i.e., (h,0e,) (2o) In the attention layers, we use KV caches for preceding tokens to avoid re-computing their features. Then, we compute the t-th feature in each layer of the decoder (i.e., (h, (h.c.,1 )/=o) to decode zt. This is possible because we have already computed features (hrae- I[=o in the prefix encoder and all features in the suffix encoder. Since there is self-attention layer in the decoder, KV caching is not needed.
ØT. When decoding ær, we
F. Text Samples from Tracformer
F.1 Unconditional Generation
The following are randomly selected unconditional text samples generated by Transformer model with sequence length of 128:
in the fishing hole outburst at or before the time that the whale had successfully spawned. Indeed, a relatively close encounter may have emboldened the whale to pull the mylar through the sand and shelled as she was able to in this example. Interestingly, the flattening events which resulted in the. Whale cutluring was caused by the thermal spark that resulted from the spin and instability of the high (or low) pressure, thus inducing the ripening state. The more the thermal increase of the high (or low) temperature was offset from the opposite direction along seismic quadrants. Interestingly. the evaporation rates at the
the Antiquarian League of Massachusetts (ASML) works with the Massachusetts League of New Hampshire Bar Associations enforce the law through nicety and compromise. Unlike the League of Massachusetts Missouri, and its members are independent, their committees are almost entirely autonomous, with sole jurisdiction solely on the matters of labor rights for inmates. They keep offsite their legal animosity to the state, only restricting their manifestations to marking weeks in anvance.And, in an unusual tumn, they hold committees to run stretchin, direct sensitivity exercises of a ceremonies- kind. These bundles of letters in the verse, order to mark the time during which batch
18

---

Tractable Transformers
The 90-day player initiative on June 2 also launched across all seven parks and recreational areas in the National Park system in Central Sierras. There, children along with the government hamstrung the operation, action hoped to be long-term solution to the larger problem of ecological loss along the Central Sierras. The USDA estimated that 20 275-acre sites on the Hudson River basin are almost colliding, forming a "very difficult interplay of life and conflict." " They include the East Bay Schuylkill River (230 versus the 81-acre Wyoming Standard Ditch), the...
was directed against the convict institution. Augustine said the parish sincerely believed the practice was against the 100 U.S. Statute. In fact if that is the practice, we would not call it a crime that was committed in Phillippines we were talking about out of state institutions. And we don't do that. So if the confession was confessionally formed it is in the whole, full church confessional as we described it."" n'giinn,, the most surprising day of the probation hearing was the fact the single author testified the confession was genuine. \n The principal tried to inscribe the patriarch,
still see themes that clouds up uncertainty about the work of first Y.) nInterestingly, note how Elsa feels an inability to see how far to go on her journey those citymarish dark castles. Darkness isn't the right word for her solution to the riddle- the life of an urban teen slowly dissolving behind the clouds, every shadow suddenly overwhelmed or broken, only to be replaced by more believable and understandable life. The crowding of darkness adds to the disorder. Of course she doesn't 't expect Snow to go though, and Elsa will have to find her own answers as to how far ...
TalkTalk Studios production. "We're re stellar!"!!,,)nua the time of the movic's release, there were quite few references to 5.0 rating and 60% top ballerinas, plus amusing examples of mixed reviews and negative news coverage. "We really appealed to the social media revolution because think that's really got to get people excited and care about their content: whether they like it or not," producer Jill Mullin has explained. \,n nn downhill slope has been one of the biggest tickets for Sony's likes-on-rumours
F22 Conditional Generation
The following are randomly selected conditional text samples generated by a Transformer model with a sequence length of 128. The prompt texts are bolded and in blue:
track-two diplomacy between U.S. and Soviet officials and nuclear scientists, which helped lead to the Compre- hensive Test Ban Treaty. MacArthur grantees also helped develop real alternatives to the IAEA, such as providing uranium to New Zealand. The centerpiece was the launch of a satellite called WUSA. \n' ()),nFiinnncn, of the IAEA declined. The international community supports the cessation of proliferation, nuclear nonproliferation, and proliferation reduction, and has improved the coordination of security policy on issues like nonproliferation and arms control.But more recently, they have borne the brunt of its economic decline, enduring lower wages
and to ascertain its peaceful means of production.Mugall has ruled Zimbabwe since independence from Britain in 1980 but faced an unprecedented challenge in Saturday's 's elections because of the unrest across his country. nn) FFFrm newly jailed opposition leader Eldion Oliverine for the Republic of Zimbabwe to the conservative populist Robert Vibert for the ruling party, the verdict could mark "new beginning" for Zimbabwe, which had once been its own country. (n n'The verdict applies to brain cancer, O. Wayne Corley, senior shareholder in the McNair Law Firm, Another question that seems never to have been raised is that
Powerton Historic District. Nearly 1 million square feet may be taken care of by the Virginia Historical Com- mission, say officials, but the built-design preferred by the Lee administration would have led to up to six months. Thousands of jobs are threatened, and towns will lose part of their economic base.OK, so today is time for some Aliens Interviews: Episode 5 - Gruff Ed Beeton interview. AnA you might expect but a few comments are here:) n Term Ear Ham Mick Deputy Assistant to the President, Mr. Ballentine was Special Assistant to the President for Legislative Affairs, where he focused on energy
19

---

Tractable Transformers
full-time practice for the Dynamo with one of their available options make his debut from the bench is looking likely to be the more familiar striker. The way players used to it is not the "main factor," Roden said.If there are rights, it is too little.He had been in San Antonio since the winter of 2010 as civil servant with pension to return to Brazil. "Brazil is the kind of economy that opened up variations in economic conditions and strengthened the key refinery centers," said Jim Rouiller, meteorologist with private weather forecaster Heaialyytess.aaauureeeeeee.e in Basking
at wavelength of 13 cm show evidence for water ice (Fig. 16()). The middle crust of the pixel is depressurised by around 0.3 Ả before a crystalline bath and column. The software eliminates any liquid after residual ice, down to the image resolution of 10 m per pixel. | On November 15, 2008, 34-kg, 5.2-pound airborne "mysterious satellite" appears on high resolution radar screen at altimeter branches at the Pangaea solar array observatory on Oak Flat, California, with a radar altimeter, video imaging system, and mass spect
ramuros, were destroyed but after the war, reconstruction took place. \n In 1948, President Elpidio Quiroga set in motion a program that could reshape Philippine history. The concept was called "The Bilateral Relations with Manila, created in 1938 by former President Manuel L. Quezon, which was named after him. The move ended any implementation first created by Pren-Prreideet Marcos. nin 1945, the Bilateral Relations obscure evolved into inter-provincial relations. Manila and Bilateral Relations were also linked to the Bilateral Relations. \ n'The Bilateral relations were based on strategic, tactical
is difficult to make sense of. Most researchers think that terrestrial vertebrae are made of tendon, but there are several proposed hypotheses.1 Two proposed vertebrae, which sometimes don't 't such as Sclerothorax and Eryops that may have been at least partly terrestrial also have long neural spines on top of their vertebrae that would have stabilized the spine and required reattachment to form non-spine rest part.3 Another speculation involving the placement of vertebrae in a notch below the flexic museles on the back of the spine would seem a more plausible explanation.3 Another proposed vertebrae
G. Proof of Conditional NELBO for MDLM
In this section, we generalize the unconditional negative evidence lower-bound (NELBO) derivation in Sahoo et al. (2024) to the conditional case. The resultant formula is then used to compute the conditional perplexity of baseline discrete diffusion models in various settings. First, we derive the expression for conditional NELBO, We then examine and simplify the various terms within the conditional NELBO. Finally, present the final expression for the conditional NELBO.
This proof builds upon and extends the results presented in Sahoo et al. (2024). Equations from the original paper are frequently referenced, and readers are encouraged to refer to the original paper for a deeper understanding of the foundational concepts and equations discussed.
G.1. Derivation of the Conditional NELBO Expression
For a sequence of L tokens, let F === {1, }} denote the set of all indices in the sequence. Let S C F' be the set containing all and only the indices of the given tokens. In this proof, we denote the entire sequence as xF (equivalent to T1.L in Sahoo et al. (2024)). the given tokens as xS, and the unknown tokens as F Following the definitions in Sahoo et al. (2024), we define s(i) = (i - T) and t(i) = i/T, and omit the i from t(i) and s(i).
Using the aforementioned notations, we can express the conditional negative lieeliiel as logg po (20f- Now, step into the details of deriving the conditional NELBO:
llog po (x )
- =llog f po po(x "2
2(rr) 1 (x )22)))
F
(10) dzfpr ((T);
F 2 10 log / a(=ho))... ,
Po Fx F-S 2F = (T))
(x
dz dz(())
(2 (0)) = (TT)
xx))
po (a P-S (0))
(()) (a:5)
< s- 9(2 ((= ... =/m x log
(2 ... dz
((2 (T) x
20

---

Tractable Transformers
Po F-S F
F ((T) (() ................................ (T) xx
(x
((0)
log
po( (x F- (0) x
pol 2 ((T) (x:5) IT:=1 po (2(52l22 of
2 x
log
(note that (resp. t) denotes s(i) (resp. t(i))),
F
()) II q( (z55)
(0)
T
(( |z,
po (x: P- 110) lxc
q 2 F
Eg
logpo(z((7) (x
+
log
+ log
Po (2 2 x55)
((0) xx55
1=1
(= 8 b x²"
Egr
log 12:
+
log
pal (= =t 5x55
Am
log q( x
+ log g(=1()) lx²)
(10)
(2 [x*
pe (xx F-8 x ((0) x))
(= bz x²))
log (( (x
g<2 xx p)
Egz
log po (((()) (x
+
log
+ log
po (z.") zi x55
g g(2 xx))
x F
=((()) x
F
2( (() (x=²)
T
(x Fl x x
E_
log
log
log (x P-9
(2())) x5
po {T) (3:
po (2
Fl |zf x3))
E,
po xx F
((()) x
DKL D lef of llpo (2 1zf xx11
Lreonns
Lditfusion
+ DKLlg(Zz(r)) lxf
l Ollpo(z x
(11)
ppioor
In Equation (9), = used to denote [=f
{(r)) for simplicity. When deriving Equation (10) from its previous step,
we make use of the following factorization:
((=f(z5)) = q( 1 xFF) 9(2, |zf x )g(2f lxx
(22
G.2. Simplification of Conditional NELBO Terms
As assumed in Section 3.5 of Sahoo et al. (2024), the forward noising process is applied independently across sequence, and the denoising process factorizes independently across tokens. These assumptions be translated the following equations:
{(2 2 "=,,,="") x²))
II t=1
q( (< x²)
(12)
pa(z F 12," =
II po (2(2)
(13)
As defined in Equation (7) of Sahoo et al. (2024), for vl V F.
pe po(2{1=f) = q(2/1=f x² af((f,tt)
(14)
Note that the original definition was for the specific case where only a single token is noised and denoised. However, since we are addressing the noising and denoising of entire token sequences, we extend the definition to accommodate sequences of tokens.
21

---

Tractable Transformers
Additionally, we define:
po po(2|1z xx)) 1000g(z()) q(21|/ x'= (12 ,,,,,,
(15)
where |x,5-5,=* 21 is shorthand for z. with z, replaced by x', for vl S. In practice, this is easily done by replacing the tokens with indices in S with the corresponding given tokens, while leaving the remaining tokens unchanged before feeding the sequence into the model zz.
According to "Carry Over Masking" in the original paper, the model 2g has the property that 2z(2,, t) zf if 2 is not masked. Therefore, for Vl E S,
Po Po(2|2z ,x)) )=)== = q(<2//// x= = 211 ,, ==(=[12[; x²)
(16)
Note that the symbol the x reused: the xf in the first line represents a random variable, while the af the second line the f-th token of the ground truth sequence, which is given one-hot vector.
As extension to (13), we further assume that:
po (0 22 x)) II po (= 20
(17)
Using these equations and the continuation of timesteps (i.e. T - oo), can further simplify the terms in (11).
G.2.1. DIFFUSION LOss
Following Section A.2.3 of the original paper, let Lr - Eiet +. ...... By ( DDD.Ll(( |=f", x llpo ( (2 lz 20
denote the diffusion loss. We can simplify with the following steps:
Cr = EEE 7. 1) (F
DKLla(>)1 20)0oo (1po(22 xx x²)]
Using (12) and (17)
i=1
= EEEt+. #.... E [ (nF)
Σ
DRI. 100 x²))
Using (16)
LEF-S
E 1) E xf m²) [T DKLla(> Ilpaa x
=
tEF-S
which is essentially the sum of the losses for each unknown token.
In the appendix of Sahoo et al. (2024), the derivation is focused on the case where only single token is involved. But since we decomposed the diffusion loss of the sequence into the sum of single-token losses, now we are able to use the results from the original paper, as long as we can address the difference between po(2[2z; x)) in this derivation and po(2, (2;) in the original paper.
According to Equation (15) from the original paper, we can derive from Equation (15) that:
po(2]1=f = a(=51=(, o0 = (r6 " 1,4) Cat 2 Qu o QJ 2 F F-S ,xs1,)))
z! QJ ab (1 F-8 .o5) 1))
Thus, the KL divergence term in the diffusion loss can be expressed as:
z Qq o Q. (( - s. =.,,1)))
DKLlg(z(12f x² [po x = DKL q(>(( q(=51=/ x [[at
(18)
of QJ a/ (1z)" 111)
while for the single-token case as in the original paper, the KL divergence term DKL (g
po (zf(=() is:
DKL[g(2/1 x: (2{122= DKL q(2//2/,, x [a]|Caat zz: Q1 2 QJ 2/
(19)
2 Q. (2f t))
22

---

Tractable Transformers
In Equation (19), the index & F S is added for notational consistency between Equation (18) and Equation (19). As see, the sole difference is that in Equation (19), ( 2 1) is substituted with (z t). Consequently, long as the former shares the same properties as the latter, the results derived for the single-token case still applies.
When simplifying the diffusion loss, only properties are required of az t) in the original paper:
1. The output of =f)( t) is probability distribution over all categories, i.e. (1, t)) 1.
2. "Zero Masking Probabilities" i.e. (a t), m) O.
3. "Carry Over Unmasking" i.e. az 1) xf when z = x.
which all happens be properties that MDLM models share, and therefore also holds for (1, 2xx)),)) Thus, by substituting (z;; with - t), we can apply the results from the single-token case and derive the following equation from Eqn. 42 of the original paper:
LT
Σ EEF-S
( +.......)
Po oF,T
Os log(a/((=)
F-S , 1), 1)),a())
(20)
at
G.2.2. PRIOR LOss
According to Equations (37) and (38) in the original paper, we can easily derive that:
(( M| 1,
po po(2((r) M| -
=
where M [m, m, m,...,m m| E RKLL and K is the number of categories.
Thus, (< (x)) and ps (z(()) (x)) are identical distributions where z((z) has a probability of I of being an entire sequence of masked tokens, resulting in prior loss of O:
Lprior DrI.19(zEmm)e (10 (1lp) (x))) 0.
(21)
G.2.3. CONTINUOUS TIME
For the continuous-time case, similar Equation (44) from the original paper, we can derive:
0
00
dilfusion
lim Lr
Σ
B,-210.11, (ef)
I-a;
log(z (1 x 1x)),),
(22)
T-000
LEF-S
According to Equation (45) from the original paper, under continuous time, we also have 20 x. Thus we have the following:
Crecoms = By(=E (10) (xF) bogpo | x x3),
= Bo(=foo) (mm)) |- log pe (x F 12
((() = x x),
Using z((0) x
By(aE
| | log
II
(x) (1 --S ,,f=,xs),t(()).. x'}}
= By(=fo | log
II
(x' x)),
(23)
fEF-5
log
II 1),
F-
O,
(24)
where deriving Equation (23) from its previous step makes use of the "Carry Over Masking" property, i.c., 2/ ((2f--5 F-S ,xx5),
?
23

---

Tractable Transformers
G.2.4. CONDITIONAL NELBO
Finally, using Equations (21), (22) and (24), we arrive at the final expression for conditional NELBO:
LELLOO Lrecons + Loo
diffussion
+ Cprirr:
dj (og( 2x0) 1), xz)
Σ
Byg.
1 - at
EEF
-
24


Reviews:
[{"review_id": "Jvt3P5GaKZ", "reviewer": "Reviewer_sirR", "summary": "This paper is motivated by the fact that Non Auto-regressive (NAR) generative models do not work very well for conditional generation. The paper proposes Tracformers, a transformer-based architecture robust for conditional generation in the more difficult NAR setting: this is done through using *multiple context levels*, in order to force the model being able to use only local subset of the context when training for conditional generation, which authors call *query generalization*. After describing their proposed architecture, they describe two settings, Contextual AR generation and Arbitrary Context (AC) generation, for which they describe the instantiation of the model and the loss. Experiments compare, in both settings, a model built with Tracformers to basic (BERT and BART) and a scaled-up version to recent (diffusion-based) generative NAR models, using conditional complexity on Wiki103, 1B and Lambada.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> The abstract is quite unclear in retrospect; terms that are well defined later (conditional probability query) appear obscure.\n\nWe 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.\n\n> 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.\n\nWhile 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.\n\n> Footnote 1 is unclear: I don't see how it matches Equation 3.\n\nWe 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}$} \n\nand\n\n$\\phi^{l-1}_{t}$ = {$t’ : t’ \\geq 1, 0 \\leq t - t’ < 2^{l-1}$}. \n\nTaking the union of the two sets leads to\n\n$\\phi^{l}_{t}$ = {$\\{t’ : t’ \\geq 1, 0 \\leq t - t’ < 2^{l}\\}$}. \n\nWe will fix it in the next version of the paper.\n\n> While this paper is generally well written, Section 4 stays difficult to follow.\n\nIn 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.\n\n[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).", "claims_and_evidence": "The paper claims that *Tracformer’s multi-scope attention mechanism and specialized encoder-decoder design enable robust conditional generation performance*, which is verified by the experiments. However, as pointed out in the conclusion, the larger model proposed is the size of a GPT-2 model, and the question remains if the reduction of interactions in self-attention in the Tracformer layer will significantly restrict the model's expressiveness at larger scales (and in particular, for larger encoded sequences).", "methods_and_evaluation_criteria": "The method relies on an encoder-decoder architecture, where:\n- The encoder uses increasing context length with the layers (exponential growth), with a maximal number of attended positions\n- The decoder work through cross-attention only, and making independent predictions at different positions; attending first encoded representations built on large scope and moving towards local information.\nThe description of the architecture relies on mask to restrict attention to the desired scope. \n\nDepending on the setting (CAR implies auto-regressive generation with partial knowledge of the future, while AC is completely NAR) and the available input context, different mask schemes and losses are implemented. \n\nEvaluation is made mainly on conditional perplexity of the generated sequences, as well as infilling performance with MAUVE and Bert score.", "theoretical_claims": "No theoretical claims are made in the paper.", "experimental_designs_or_analyses": "- The experiments compare the Tracformer with the main types of models corresponding to AR and CAR settings; and then, to state-of-the-art NAR generative models. \n- Models are compared across various mask ratios.\n- The analysis is extensive, and further ablation studies are provided in the supplementary material.", "relation_to_broader_scientific_literature": "This paper's contribution is a new architecture for NAR modeling, and could be used with dedicated NAR models, such as discrete diffusion models.", "essential_references_not_discussed": "I have no knowledge of any essential reference missing.", "other_strengths_and_weaknesses": "- This paper propose a promising idea which should be further explored.\n- While this paper is generally well written, Section 4 stays difficult to follow.", "other_comments_or_suggestions": "- The abstract is quite unclear in retrospect; terms that are well defined later (conditional probability query) appear obscure.\n- 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. \n- Footnote 1 is unclear: I don't see how it matches Equation 3.", "questions_for_authors": "I have no question for the authors.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lOGrXnpPtd", "reviewer": "Reviewer_KPNd", "summary": "This paper proposes a novel model architectural modification of Transformers and demonstrate that the proposed model outperforms baselines on non-autoregressive (NAR) conditional generation tasks, especially when the mask pattern at inference is different from the mask pattern at training.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback and for recognizing the novelty and clarity of our work.\n\n> 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.\n\nWe 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. \n\nWe 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.\n\n> the tasks, datasets, and evaluation metrics used in this paper mostly correspond to generic language modeling, emphasizing on fluency.\n\nWe 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.\n\nTo 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.\n\n| Mask ranges                     | Tracformer | BART  |\n| ------------------------------- | ---------- | ----- |\n| [0.25,0.75]                     | **0.524**  | 0.513 |\n| [0.5,1.0]                       | **0.540**  | 0.519 |\n| [0.1,0.4] & [0.6,0.9]           | **0.419**  | 0.405 |\n| [0,0.4] & [0.5,0.8]             | **0.339**  | 0.325 |\n| [0,0.25] & [0.75,1.0]           | **0.536**  | 0.523 |\n| [0,0.1] & [0.2,0.5] & [0.7,1.0] | **0.337**  | 0.322 |\n\n> The pre-trained SEDD was only able to handle sequences of length 1024. How did you calculate the log likelihood of shorter sequences?\n\nWhile 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.\n\n> Does SEDD output the exact model-predicted likelihood, or its lower bound (ELBO)?\n\nAs 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.\n\n> 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?\n\nYes, 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.\n\n> I would like to learn from the authors about their insights on my feedback shared in the above sections of this review form.\n\nWe 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.", "claims_and_evidence": "Most claims are supported by clear and convincing evidence.\n\nOne somewhat inaccurate claim is that at the end of Section 6.2, below Table 4, the authors claim that “Tracformer remains highly competitive, achieving results comparable to or better than larger models.” \n\nBut the *unconditional* generation perplexity (Table 4) of the proposed model does not compare favorably with strong baselines such as GPT-2, SEDD, and MDLM. \n\nI 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.\n\nThis limitation does not undermine the authors claim regarding *conditional* generation tasks.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense.", "theoretical_claims": "This paper does not claim to include theory as its primary contributions.", "experimental_designs_or_analyses": "I checked all experimental designs and analyses in the main paper. They are sound and valid.", "relation_to_broader_scientific_literature": "This paper contributes to improving the generalizability of non-autoregressive (NAR) conditional language generation tasks, and compares favorably against SOTA discrete diffusion language models under OOD masking patterns.", "essential_references_not_discussed": "References are sufficiently discussed.", "other_strengths_and_weaknesses": "Strengths\n\n1. originality: the multi-scale attention applied to the NAR context for better generalizability is novel. (Similarly designed sparse attention may have been used for AR Transformers to improve efficiency in long-context settings.)\n\n2. clarity: the techniques and the results are clearly described.\n\nWeaknesses\n\n1. significance: the tasks, datasets, and evaluation metrics used in this paper mostly correspond to generic language modeling, emphasizing on fluency.\n- This paper does not demonstrate that the proposed model outperforms baselines under some task-specific metrics (e.g. BLEU for translation, ROUGE for summarization, accuracy for reasoning, etc). \n- These specialized tasks are commonly benchmarked for non-autoregressive language modeling papers. They provide important additional information since in these tasks, fluency is not the only important capability.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. In Figure 2:\n- The pre-trained SEDD was only able to handle sequences of length 1024. How did you calculate the log likelihood of shorter sequences?\n- Does SEDD output the exact model-predicted likelihood, or its lower bound (ELBO)?\n\n2. 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?\n\n3. I would like to learn from the authors about their insights on my feedback shared in the above sections of this review form.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZE2FBp5EQT", "reviewer": "Reviewer_QzDn", "summary": "This paper explores why non-autoregressive (NAR) generative models often underperform in conditional tasks, despite strong unconditional performance. The authors introduce Tractable Transformers (TracFormer), which factorize conditional queries to handle partial inputs flexibly while leveraging both local and global context. Their experiments show that TracFormer achieves robust conditional generation, surpassing diffusion- and autoregressive-based baselines.\n\n## update after rebuttal\nThe authors have provided a thorough rebuttal and addressed most of my concerns. I continue to believe that my original score was an accurate evaluation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback and for acknowledging that the paper addresses a clearly defined and significant problem.\n\n> At line 630 in Appendix A, are you referring to Decoder instead of Encoder?\n\nWe 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.\n\n> In Figure 4(a), is the token identical to x₀ from the prefix encoder input understood from a teacher forcing perspective?\n\nYes, 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.\n\n> 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?\n\nThere 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.\n\n> 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?\n\nAs 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$.\n\nAdditionally, 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.\n\n> 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?\n\nThis 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.\n\n> 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?\n\nIn 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.", "claims_and_evidence": "Yes, the observations are clearly articulated, and the proposed methodology to address them is both well-founded and convincingly supported by experimental results.", "methods_and_evaluation_criteria": "The methodology clearly outlines how it aims to address the problem, and it is easy to understand how the proposed approach tackles the observed issues. Additionally, the evaluation was conducted using fair benchmarks in comparison with existing methods.", "theoretical_claims": "I have reviewed Appendix G and did not identify any major issues.", "experimental_designs_or_analyses": "The necessary experiments were well-designed, and the results were appropriately analyzed.", "relation_to_broader_scientific_literature": "Recently, there has been a trend in LLMs towards using simple transformers with tokenization, which has led to a significant reduction in architecture search even at smaller scales. In this regard, I find the scalable proposal in the paper to be impactful, although scaling in practice presents its own challenges.", "essential_references_not_discussed": "There are no comments on this matter.", "other_strengths_and_weaknesses": "The paper addresses a clearly defined and significant problem, with observations that are straightforward. The proposed methodology is well explained, making it easy to understand how it tackles the problem. Moreover, the work is supported by well-designed experiments and convincing results, and the analysis effectively addresses any potential concerns.\n\nHowever, a few clarifications were needed, and the details are provided in the \"Questions For Authors\" section.", "other_comments_or_suggestions": "At line 630 in Appendix A, are you referring to Decoder instead of Encoder?", "questions_for_authors": "1. In Figure 4(a), is the token identical to x₀ from the prefix encoder input understood from a teacher forcing perspective?\n2. In the original Transformer decoder, self-attention and cross-attention are both present. 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? Was this design choice naturally adopted from a modeling perspective, or did it stem from a separate insight by the authors?\n3. On line 242, the paper states, “Attention masks are used in cross-attention layers to ensure each decoder embedding only depends on desired input variables.” 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?\n4. 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? Is this choice made solely for modeling purposes, or does it reflect the authors’ specific intuition?\n5. 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "878asS12XK", "reviewer": "Reviewer_7ULg", "summary": "This paper introduces Tracformer, a Transformer-based generative model designed for flexible and robust conditional generation tasks. Tracformer incorporates a sparse multi-scope attention mechanism to capture both local and global contextual information efficiently. Empirical results demonstrate that Tracformer outperforms existing NAR models like SEDD and MDLM in both conditional and unconditional generation tasks, showcasing its potential as a scalable and versatile generative model.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback and for recognizing the potential of Tracformers in conditional generation tasks.\n\n> Some abbreviations are repeatedly defined throughout the paper, which leads to unnecessary redundancy.\n\nWe appreciate the reviewer’s feedback and will revise the paper to ensure that abbreviations are defined only once.\n\n> The contributions of the paper would be more clearly understood if they were summarized at the end of the Introduction section.\n\nWe 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.\n\n> 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.\n\nWhile 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.\n\nIn 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.\n\n> The results for zero-shot unconditional perplexity in Table 4 show that Tracformer outperforms the other baselines on only the 1BW dataset.\n\nWe 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.\n\nAnother 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.\n\n> The Background section could be moved to the appendix.\n\nWe 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.\n\n> …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.\n\nWe 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.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Fair", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n1) The motivation for the study is clearly introduced. Non-autoregressive (NAR) models have shown superior performance in unconditional generation compared to autoregressive (AR) models of similar sizes. However, a significant challenge with NAR models is their difficulty in generalizing to conditional probability queries that were not seen during training. In response to this issue, the authors propose the Tractable Transformer (Tracformer), a NAR model designed specifically for flexible conditional generation tasks.\n2) The experimental results demonstrate that the proposed Tracformer model outperforms several established models, including BERT, BART, SEDD, and MDLM, in terms of conditional generation performance.\n\nWeaknesses:\n1) Some abbreviations are repeatedly defined throughout the paper, which leads to unnecessary redundancy. For example, terms like non-autoregressive (NAR) and feed-forward neural network (FFN) are defined multiple times. \n2) The contributions of the paper would be more clearly understood if they were summarized at the end of the Introduction section. This would help provide the reader with a clearer context and set expectations for the rest of the paper.\n3) 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. \n4) The results for zero-shot unconditional perplexity in Table 4 show that Tracformer outperforms the other baselines on only the 1BW dataset. In contrast, MDLM demonstrates superior performance on three datasets, while Tracformer only outperforms baselines on one. \n5) The Background section could be moved to the appendix. Sequence modeling and Transformer models are well-established concepts and are common knowledge for researchers and readers familiar with this field. Additionally, 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.", "other_comments_or_suggestions": "None", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper addresses conditional generation in NAR, and proposes a model called Tractable Transformers (TracFormer). Essentially, the authors devised a multi-scope attention that uses only local subset of the context.

While reviewers believe the idea is interesting, there are a few drawbacks of the experiments. For example, the language models used are generally small, and the evaluation is mostly PPL oriented, except for the text infilling experiment which uses MAUVE and BERT scores (Table 1).
Decision: Accept (poster)