ICML 2025

Tractable Transformers for Flexible Conditional Generation
Anji Liu * 1 2 Xuejie Liu * 3 4 Dayuan Zhao 5 Mathias Niepert 2 Yitao Liang 3 Guy Van den Broeck 1
Abstract
Non-autoregressive (NAR) generative models are
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 (e.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-based
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:
//github.com/liuanji/Tracformer.
*Equal contribution
1Department of Computer Science, Uni-
versity of California, Los Angeles 2Institute for Artificial Intelli-
gence, University of Stuttgart 3Institute for Artificial Intelligence,
Peking University 4School of Intelligence Science and Technology,
Peking University 5Yuanpei College, Peking Universi...