Conditional Attribute Transformers (CAT)

What if a language model could predict not only the next token, but also its consequences?

We introduce Conditional Attribute Transformers, which jointly estimate the next token and, for each possible next token choice, sequence-level properties or outcomes.

This gives generative models three key capabilities in a single forward pass: (1) token-level attribution to downstream outcomes, (2) counterfactual reasoning about how an outcome would change under alternative next token choices, and (3) steering toward safer or more optimal outcomes through sequential next token selection.

We show that Conditional Attribute Transformers achieve state-of-the-art performance in reinforcement learning tasks and language modeling. In medical foundation models, they enable dynamic, interpretable risk estimation for downstream clinical outcomes and elucidate the tokens that drive risk, while achieving a 10⁸× speedup over traditional sampling-based approaches. As an additional benefit, we find that this joint task improves next-token prediction in baseline language models.