Tractable probabilistic models (TPMs) are increasingly being applied in language modeling tasks, such as controlling the generation of large language models (LMs) to satisfy logical or semantic constraints. However, a thus-far overlooked aspect is the training procedure of the TPM, which optimizes for data likelihood rather than the downstream task. This raises a central question: does improving data likelihood necessarily improve downstream generation quality? In this work, we find that likelihood-trained TPMs can result in failed generations due to overly large corrections to the LM's logits. To address this, we train TPMs with LM-aligned objectives, including an ELBO-based objective and an MSE-based distillation objective. By learning directly from the LM rather than optimizing only for data likelihood, our approach produces TPMs that better align with the LM token-probability space. On a detoxification task, our results show these models avoid degeneration, maintain fluency under strong guidance, and enable stronger detoxification.

@inproceedings{HLiuTPM26,
  title     = {The Pitfalls of Maximum Likelihood Training of TPMs for Controllable Language Generation},
  author    = {Liu, Hanzhang and Zhao, William and Shao, Zilei and Wang, Benjie and Van den Broeck, Guy},
  booktitle = {Proceedings of the UAI Workshop on Tractable Probabilistic Modeling (TPM)},
  url       = "https://starai.cs.ucla.edu/papers/HLiuTPM26.pdf",
  month     = 7,
  year      = {2026},
  keywords  = {workshop}
}