Conditional Attribute Estimation with Autoregressive Sequence Models

Erica Stutz
Department of Biomedical Informatics and Data Science
Yale University
New Haven, CT 06510
erica.stutz@yale.edu
&Giacomo Marino
Department of Biomedical Informatics and Data Science
Yale University
New Haven, CT 06510
giacomo.marino@yale.edu
&Daniella Meeker
Department of Biomedical Informatics and Data Science
Yale University
New Haven, CT 06510
daniella.meeker@yale.edu
&Qiao Liu
Department of Biostatistics
Yale University
New Haven, CT 06510
qiao.liu@yale.edu
&Andrew J. Loza
Department of Biomedical Informatics and Data Science, Department of Pediatrics
Yale University
New Haven, CT 06510
andrew.loza@yale.edu

Abstract

Generative models are often trained with a next-token prediction objective, yet many downstream applications require the ability to estimate or control sequence-level properties. Despite their success, next-token prediction can lead to overfitting of local patterns during training, underfitting of global structure, and requires significant downstream modifications or expensive sampling to guide or predict the global attributes of generated samples at inference time. Here, we introduce Conditional Attribute Transformers, a novel method for jointly estimating the next-token probability and the value of an attribute conditional on each potential next token selection. This framework enables three critical capabilities within a single forward pass, without modification of the input sequence: (1) per-token credit assignment across an entire sequence, by identifying how each token in a sequence is associated with an attribute’s value; (2) counterfactual analysis, by quantifying attribute differences conditional on alternative next token choices; (3) steerable generation, by decoding sequences based on a combination of next-token and attribute likelihoods. Our approach achieves state-of-the-art performance on sparse reward tasks, improves next-token prediction at sufficient model sizes, estimates attribute probabilities orders of magnitude faster than sampling, and can guide decoding of autoregressive sequence models on a range of language tasks.

1 Introduction

Generative models have demonstrated performance advances across multiple domains from language to biomedical informatics Brown et al. (2020); Ferruz et al. (2022); Waxler et al. (2025); Brixi et al. (2025). While next-token prediction is a scalable training objective Hoffmann et al. (2022), it optimizes for local coherence and can lead to greedy overfitting of local patterns and suboptimal prediction of critical branch tokens Gloeckle et al. (2024); Qi et al. (2020). Furthermore, the downstream utility of these models is frequently defined by sequence-level attributes Chang et al. (2023). In language models, there is need to control token selection to create text with specific attributes, such as correctness or helpfulness Keskar et al. (2019). In biomedical informatics, generative models are used to estimate clinically relevant sequence attributes such as disease onset, medical events, or treatment response through expensive Monte Carlo (MC) simulation Waxler et al. (2025); Renc et al. (2024); Shmatko et al. (2025). The capacity to learn representations that better capture sequence-level properties, estimate properties efficiently from partial sequences, and control generation to create sequences with specific properties has broad implications across domains.

These diverse use cases share a common mathematical requirement: the need to estimate a sequence-level attribute for a partial sequence, conditional on the selection of the next token. This capability allows for estimation of attribute likelihood for observed sequences or steerable decoding to optimize the likelihood of a particular attribute. Current methods for predicting or controlling sequence-level attributes are computationally demanding because they often require properties that may not emerge until many tokens in the future. Current approaches fall into two main categories: conditioning and base-model steering with auxiliary models. These methods have several limitations including high computational overhead, requirements for additional model training, and limited flexibility.

Here, we propose Conditional Attribute Transformers (CAT), a method for conditional attribute estimation (Fig. 1). Leveraging a reinforcement learning framework, we cast data as arising from an unknown sequential game. We develop a generative modeling of this framework with a joint objective of next-token prediction and conditional sequence-level attribute prediction in a single model through a branched architecture and shared latent space.

The specific contributions of this work are as follows:

•

We provide a framework for simultaneous estimation of next-token likelihood and a sequence-level attribute likelihood, conditional on each potential next token. We also link this objective to components of causal inference and reinforcement learning.
•

We demonstrate that this objective can be integrated into (a) the pre-training of generative decoder-only transformers with minimal computational overhead and can synergistically improve next-token perplexity, or (b) pre-trained models through fine-tuning.
•

We evaluate performance on three diverse tasks: (1) learning strategy from random gameplay, (2) predicting and controlling the likely rating for Amazon product reviews, and (3) predicting sepsis onset in a medical data set.

Refer to caption — Figure 1: CAT is a unified architecture for next-token and sequence-level attribute prediction. Tokens ( $t_{n}$ ) are processed by a shared backbone. The final latent representation is simultaneously processed by the language modeling head (Token Head) and an integrated conditional attribute model (attribute block + attribute head). Next-token cross-entropy loss ( $L_{token}$ ) is combined with the attribute loss ( $L_{attr}$ ) which can be from a binary, multinomial, or numeric attribute, delivered as a token-level sequence ( $a_{n}$ ). During training, the full conditional attribute matrix does not have to be materialized (gray) because only the attribute of the true next token is seen, although it is learned. Lower right panel shows example of continuation (dashed) or CAT steering (solid) to change a 5 star prefix into a 1 star full review.

2 Related Work

Numerous methods for controlling or estimating sequence-level attributes have been developed and fall into two classes: generation via conditioning or inference time guidance by a separate model.

Conditioning is used for generative model steering by inserting a fixed prompt or code in its input. Pre-training conditioning includes methods like CTRL Keskar et al. (2019), which prepends control codes to sequences, and Decision Transformers Chen et al. (2021), which insert reward-to-go tokens within each reward-state-action tuple using an offline reinforcement learning framework. Among post-training conditioning methods, Quark prepends sequences with reward quantile tokens Lu et al. (2022). While these models can condition generation, they do not ensure that insertion of control tokens produces sequences that remain in distribution or provide probabilistic estimates of the attribute likelihood. Furthermore, if a downstream token is selected in error, there is no means of correction. CAT, by contrast, estimates the attribute at every step without requiring modification of the input sequence, allowing for more flexible and active control of steering.

Alternatively, token generation can be guided with auxiliary models. Classifier-based methods include PPLM, which uses gradients from an external classifier to guide generation, and FUDGE, which trains a binary classifier to predict attribute selection from partial sequences Dathathri et al. (2019); Yang and Klein (2021). GeDi uses a generative discriminator to update next-token probabilities Krause et al. (2020). Director integrates the generator and classifier into a unified model by adding an attribute head alongside the language modeling head at the final latent representation Arora et al. (2022). DExperts uses two auxiliary models, one fine-tuned on the desired attribute and the other on an undesired attribute, to reweight next-token probabilities of the main model Liu et al. (2021). TRACE distills a hidden Markov model from a base language model to calculate the sequence-level attribute probability, while ILQL takes this a step further by using full Q-learning and not just single-step policy updates Weng et al. (2025); Snell et al. (2022). Many of these approaches require significant computational overhead due to the training of auxiliary models including three separate transformer models with ILQL Snell et al. (2022). Notable exceptions are TRACE and Director; however, they are limited by the expressivity of the distilled hidden Markov model and a simple linear layer. CAT maintains the full expressivity of a transformer model, yet requires substantially less computational overhead than training a full auxiliary model.

3 Methods

3.1 Definitions

We model data as arising from an unknown sequential game in which the rules, players, strategies, and differences between actions and observations are unknown. We only have access to records of its play, which include a mixed sequence of actions and observations and an associated result which we consider as the sequence-level attribute, $\alpha$ . The attribute can be binary, multinomial, or numeric. After observing these records, we aim to address two questions: (1) can an agent learn to play (i.e. to generate valid moves) and (2) can it learn to play well (i.e. to achieve a specified result)?

Each game consists of a sequence $S=[s_{1},s_{2},...,s_{N}]$ of observations or actions $s_{i}$ . We do not assume that we need to differentiate between these two during training. We use $s_{n}$ to denote a candidate next element in the sequence. For simplicity, we will consider only $s_{n}$ that are drawn from a discrete language $L$ . For a sequence $S$ , there exists an unknown function $f(S)=P(\alpha_{i}\mid S)$ that generates $\alpha_{i}\in A$ conditioned on $S$ , which we take here to be the result of the game. We assume that each sequence $S$ represents a valid game instance, but make no assumption that it is generated from a specific policy. Instead, it can be viewed as a sample from an average policy distribution $\pi^{\mu}$ .

3.2 Model

Joint Distribution:

A record of an unknown game can be drawn from the joint distribution of the sequence of gameplay and its result attribute, $P(S,\alpha_{i})$ . A standard autoregressive decomposition is:

	$\displaystyle P(S,\alpha_{i})$	$\displaystyle=P(\alpha_{i}\mid S)P(S)$		(1)
	$\displaystyle P(S)$	$\displaystyle=\prod_{i=1}^{k}P(s_{i}\mid s_{1},...,s_{i-1})$		(2)

If $\alpha_{i}$ is treated as an additional observation, a single autoregressive model can be used for actions, game state, and game result. However, the probability of $\alpha_{i}$ can only be estimated after conditioning on the full sequence. For partial sequences, MC simulation must be used to estimate $P(\alpha_{i},S)$ .

Alternative Decomposition:

We can consider $S$ to be a sequence of three parts $(S_{a},s_{n},S_{b})$ , a prefix sequence $S_{a}$ of all observations preceding $s_{n}$ , a next observation $s_{n}$ , and a suffix sequence $S_{b}$ consisting of all observations following $s_{n}$ . The original joint distribution can be expanded as:

P(\alpha_{i},S)=P(\alpha_{i},S_{a},s_{n},S_{b})

(3)

We can decompose this new expanded joint distribution as:

P(\alpha_{i},S_{a},s_{n},S_{b})=P(S_{a})\cdot P(s_{n}\mid S_{a})\cdot P(\alpha_{i},S_{b}\mid S_{a},s_{n})

(4)

We can then marginalize over $S_{b}$ :

\sum_{S_{b}}P(\alpha_{i},S_{a},s_{n},S_{b})=P(S_{a})\cdot P(s_{n}\mid S_{a})\cdot\sum_{S_{b}}P(\alpha_{i},S_{b}\mid S_{a},s_{n})

(5)

which yields:

P(\alpha_{i},S_{a},s_{n})=\underbrace{P(S_{a})}_{\mathrm{prefix}}\cdot\underbrace{P(s_{n}\mid S_{a})}_{\mathrm{sequence\ model}}\cdot\underbrace{P(\alpha_{i}\mid S_{a},s_{n})}_{\mathrm{attribute\ model}}

(6)

Since we make no assumptions on the length of $S_{b}$ , this decomposition is valid for any position of $s_{n}$ including the end of the sequence $S$ , in which case $S_{b}=\varnothing$ .

Distribution Estimation:

This decomposition can be modeled using an augmented causal transformer with two heads: a next-token prediction head for estimating $P(s_{n}\mid S_{a})$ and conditional attribute prediction head for estimating $P(\alpha_{i}\mid S_{a},s_{n})$ . The prefix $P(S_{a})$ is modeled autoregressively by the sequence head, enabling sampling from the full joint distribution $P(\alpha_{i},S)$ . Here, we use a shared model backbone $f_{\theta}$ to produce a hidden representation $H$ .

f_{\theta}(S)=H,\quad\text{where }f_{\theta}:S\to\mathbb{R}^{d}

(7)

This representation is passed to two heads:

g_{\psi}(H)=P(s_{n}\mid S),\quad\text{where }g_{\psi}:\mathbb{R}^{d}\to\Delta^{|L|}

(8)

h_{\phi}(H,s_{n})=P(\alpha_{i}\mid S,s_{n}),\quad\text{where }h_{\phi}:\mathbb{R}^{d}\times L\to\begin{cases}\Delta^{|A|}&\text{(categorical attributes)}\\ \mathbb{R}^{p}&\text{(parameterized attributes)}\end{cases}

(9)

such that the backbone model $f_{\theta}$ contains information about both next-token and attribute while $g_{\psi}$ and $h_{\phi}$ provide task-specific transformations. The function $h_{\phi}$ can be selected to produce attribute category logits or distributional parameters depending on the attribute type. Here, we demonstrate the model using binary, multinomial, and numeric attributes.

3.3 Inference

Samples of observed games can be generated via $s_{n}\sim P(s_{n}\mid S_{a})$ from the sequence model. The conditional attribute model can be used for alternative sampling strategies such as greedy decoding toward an attribute of interest. When selecting $s_{n}$ , we can restrict choices to estimated valid moves (next-token probability above a threshold $\epsilon$ ) and select one that maximizes the probability of achieving a desired result while remaining within the distribution of the training data to a certain threshold $\epsilon$ :

s_{n}^{*}=\arg\max_{s_{n}}P(\alpha_{i}\mid S_{a},s_{n})\quad\text{for all}\quad s_{n}\in P(s_{n}\mid S_{a})>\epsilon

(10)

Eq. 10 defines a greedy optimizer for maximizing the likelihood of $\alpha_{i}$ . This approach can be used to select locally optimal actions and develop a policy $\pi^{new}$ that improves upon the average policy $\pi^{\mu}$ . However, it does not guarantee attainment of the global optimum.

3.4 Relationship to Other Entities

While we introduce this model as a conditional attribute estimator, it is directly related to functions in reinforcement learning and causal inference.

Relationship to Q-Functions:

Although derived using different notation, the components of this framework map to those in Distributional Reinforcement Learning in an offline setting Bellemare et al. (2017). Consider the case where $\alpha_{i}$ represents a reward, $s_{n}$ represents an action, and $S_{a}$ represents a state. Note that here we have discussed the attribute as a terminal reward, but the derivation applies to sub-sequences as well. The sequence model $P(s_{n}\mid S_{a})$ functions as the average behavior policy $\pi_{\mu}(s_{n}\mid S_{a})$ , representing the probability of selecting an action $s_{n}$ given a state $S_{a}$ in the training distribution. The attribute model $P(\alpha_{i}\mid S_{a},s_{n})$ learns the conditional distribution over rewards. The standard state-action value function under the behavior policy, $Q^{\pi_{\mu}}(s,a)$ , is recovered by taking the expectation over this distribution:

Q^{\pi_{\mu}}(S_{a},s_{n})=\mathbb{E}_{\alpha\sim P(\cdot\mid S_{a},s_{n})}[\alpha]

(11)

We note that this work focuses on single-step policy improvement by selecting from the learned $Q$ function in a greedy (or adjusted-greedy manner) rather than learning a globally optimal $Q^{*}$ through recursive value updates Snell et al. (2022).

Relationship to Causal Models:

In the case where $s_{n}$ represents a treatment or intervention and $S_{a}$ represents the history of confounders or covariates, this approach parallels the core components of causal inference. Specifically, the sequence model $P(s_{n}\mid S_{a})$ functions as a generalized propensity score, estimating the probability of treatment assignment conditioned on covariates:

e(S_{a})=P(Z=s_{n}\mid X=S_{a})

(12)

Simultaneously, the attribute model $P(\alpha_{i}\mid S_{a},s_{n})$ functions as a conditional outcome model (or response surface), estimating the expected outcome given the covariates and treatment:

\mu(S_{a},s_{n})=\mathbb{E}[Y=\alpha_{i}\mid X=S_{a},Z=s_{n}]

(13)

While it operates on fixed covariates, focuses on single treatment assignment, and uses binary outcomes, Dragonnet is similar in that it uses a shared latent representation for prediction of a propensity and outcome model Shi et al. (2019). The recent generative causal inference approaches Liu et al. (2024b); Liu and Wong (2026b, a), further generalize this idea through latent probabilistic modeling of treatments, covariates, and outcomes, enabling flexible conditional inference and uncertainty quantification. Causal transformers use sequential data, but require separate streams for covariates, treatments, and outcomes recorded at specific time intervals Melnychuk et al. (2022). Other methods use expensive MC rollouts for estimating the outcomes associated with treatment assignments Rein et al. (2024); Xiong et al. (2024). CAT improves upon these by using full sequential data and avoiding MC simulation.

3.5 Model Architecture

We extended the nanoGPT architecture Karpathy (2023), for joint next-token and conditional attribute prediction. The forward pass of CAT matches standard transformers until the final latent embedding. This branches to a standard language modeling head and the conditional attribute transformer block (similar to how multi-token prediction is performed with DeepSeek-R1 during pretraining) Liu et al. (2024a). This is used to predict the parameters required for attribute prediction which can be binary, multinomial, or numeric. Computationally efficient training is achieved by recognizing that although next-token requires computing logits across the entire vocabulary, the conditional attribute loss only requires computing the probabilities of attributes for the true single next token. Therefore, for a vocabulary of size $V$ and an attribute dimension of $A$ , only a $1\times A$ matrix must be materialized during training instead of a $V\times A$ matrix. Cross-entropy loss is used for next-token prediction and the attribute loss can be any likelihood-based loss. We used cross-entropy loss for binary and multinomial tasks, and Guassian negative log-likelihood for regression tasks. The total training loss was defined as:

L=L_{token}+\lambda*L_{attr}

(14)

where $L_{token}$ is next-token cross-entropy loss, $L_{attr}$ is attribute loss, and $\lambda$ controls the relative contribution of the attribute loss. Optimal $\lambda$ values varied across tasks, and were selected by balancing next-token loss, attribute loss, and sampling performance. Model sizes and hyperparameters for each experiment are provided in Table A.1, A.2, and A.4.

3.6 Data Sets

Experiments were conducted on three data sets: (1) the Key-to-Door environment used in the Decision Transformers paper Chen et al. (2021) to test performance in a sparse-reward setting, (2) a language modeling data set (Amazon Reviews Hou et al. (2024)) to test performance on large action spaces, and (3) a clinical data set (PhysioNet Sepsis Reyna et al. (2019)) to test performance on a biomedical informatics task.

1.

Key-to-Door: data on 10,000 random-walk trajectories in a three-room grid world consisting of a key room, a distractor room, and a door room. The agent must pick up the key and reach the door within a fixed move budget. This environment was introduced in Mesnard et al. (2021) and later used in Chen et al. (2021) as a credit assignment benchmark.
2.

Amazon Reviews: data on 574 million Amazon product reviews from Hou et al. (2024). Each sequence consists of a product category, title, review text, and a 1-5 star rating, which serves as a multinomial attribute.
3.

PhysioNet Sepsis: ICU time-series data set with 40,336 patient sequences for early sepsis identification. Each sequence contains patient demographics followed by hourly vital-sign and laboratory-measurement tokens, with time tokens separating hourly measurements. This was derived from the 2019 PhysioNet Challenge data set Reyna et al. (2019).

4 Results

4.1 Key-to-Door: Long-Term Credit Assignment

The Key-to-Door task tests CAT’s ability to learn attribute assignment from a single terminal reward. CAT successfully learns to estimate conditional win probabilities for each potential next move (Fig. 2A) and can serve as a critic to stably classify the win probability for evolving trajectories with reduced variance compared to decision transformers (Fig. 2B)Chen et al. (2021). Using greedy decoding (Eq. 10), the model outperforms all baselines: random policy, behavior cloning, percentile behavior cloning (trained on only winning examples), conservative Q-learning, and decision transformers (Table 1).

Table 1: Win rates of evaluated methods for the Key-to-Door task. In addition to the 0.999 win rate, CAT takes the shortest manhattan-distance path 998 out of the 999 wins.

	Random Policy	Behavioral Cloning	Percentile Behavioral Cloning	Conservative Q-Learning	Decision Transformers	Conditional Attribute Transformers
Win Rate	0.031	0.016	0.951	0.133	0.946	0.999

4.2 Language Modeling: Amazon Reviews

To test scalability and performance with a large input and action space, we evaluate CAT on the Amazon Reviews dataset, using review text as the sequence and the multinomial 5-class product rating as the attribute (Fig. A.1). We evaluate (1) the impact of joint training on next-token prediction performance, (2) the performance of its role as a critic by estimating the rating from partial reviews, (3) the identification of semantic shifts caused by counterfactual adjective substitution, and (4) steering generation to create reviews with specific ratings. Full model details are reported in Appendix A.2.

4.2.1 Next-Token Prediction Performance

In contrast to prior methods where joint-training worsened next-token prediction performance Arora et al. (2022), we find that the CAT architecture enables improved next-token prediction performance with increasing model size. At small model sizes (7-72-million parameters), the conditional attribute task impairs next-token prediction. However, the 1-billion parameter CAT model achieved better performance on next-token prediction than the standard model. This synergy depended on $\lambda$ used to balance the two contributions to the joint loss (Fig. 3). Subsequent experiments use $\lambda=0.15$ to balance next-token and attribute prediction performance.

4.2.2 Review Critic Performance

CAT estimates attributes from partial sequences in a simulation-free manner. Performance of predicting the rating using CAT (conditional attribute probability for the true next token) was compared to MC simulation, an attribute head fine-tuned on a frozen standard GPT model, an attribute-only CAT model, and a version of Director which we extended for multinomial attributes (Director*) (Fig. 4A). Due to the extensive compute required, MC sampling was evaluated on 4,000 reviews and four partial review lengths, whereas other methods are evaluated on 1 million reviews.

CAT and the fine-tuned CAT model both outperformed Director*, and MC sampling using CAT’s next-token head outperformed the standard next-token model. The attribute-only CAT model underperformed most methods, except standard next-token MC simulation at a partial sequence length of 80, indicating that jointly modeling next-token and conditional attribute prediction improves performance over either task alone. CAT also provided an approximately $10^{8}\times$ speedup in partial sequence evaluation relative to MC sampling (Fig. 4B).

4.2.3 Counterfactual Estimation

Table 2: Counterfactual estimation of substituting good with alternative adjectives on predicted 1 and 5 star probabilities (

\Delta P_{1}

\Delta P_{5}

). Shown for all contexts and negated contexts (i.e. not good)

Counterfactual	$\Delta P_{1}$		$\Delta P_{5}$
Counterfactual	All Contexts	Negation in Context	All Contexts	Negation in Context
good $\rightarrow$ AMAZING	-0.01	-0.07	0.09	0.02
good $\rightarrow$ amazing	-0.01	-0.11	0.07	0.02
good $\rightarrow$ GREAT	0.00	-0.08	0.06	0.00
good $\rightarrow$ great	0.00	-0.08	0.04	0.00
good $\rightarrow$ GOOD	0.00	0.00	0.03	0.00
good $\rightarrow$ bad	0.06	-0.14	-0.08	0.06
good $\rightarrow$ BAD	0.08	-0.11	-0.10	0.03
good $\rightarrow$ horrible	0.09	-0.11	-0.11	0.02
good $\rightarrow$ HORRIBLE	0.10	-0.04	-0.19	0.01

We evaluated the semantic accuracy of CAT’s counterfactual estimates by substituting sentiment-bearing adjectives in 1,000,000 validation reviews. Replacing good with negative adjectives increased 1 star and decreased 5 star probabilities. Positive substitutions increased 5 star probabilities with little effect on 1 star probabilities. Substitutions in negated contexts (not good) had more nuanced shifts: both positive and negative substitutions reduced 1 star probability and led to no change or increased 5 star probabilities. Capitalization modulated these effects, reflecting the semantic emphasis (Table 2).

4.2.4 Guided Decoding

In steering 3 star prompts to 1 and 5 star reviews, CAT had the highest accuracy (CAT accuracy 0.64/0.77; best alternative Director* 0.58/0.65) across all evaluated models. Furthermore, generations were more fluent than models of similar accuracy (CAT perplexity 45.88/44.03; best alternative Director* 46.77/48.16) and had comparable diversity (Table 3).

Table 3: Steering 3 star prompts toward 1 and 5 star reviews.

	Model	Accuracy	Fluency ( $\downarrow$ )	Diversity ( $\uparrow$ )
	Model	Accuracy	Mean ppl.	Dist-1	Dist-2	Dist-3
1 star	CTRL	0.13	13.68	0.90	0.87	0.78
	DExperts	0.14	19.24	0.91	0.86	0.77
	Director	0.25	20.18	0.88	0.89	0.82
	Director*	0.58	46.77	0.80	0.80	0.75
	CAT	0.64	45.88	0.73	0.65	0.57
5 star	CTRL	0.41	13.45	0.90	0.86	0.77
	DExperts	0.41	18.99	0.91	0.85	0.76
	Director	0.47	18.95	0.90	0.86	0.78
	Director*	0.65	48.16	0.83	0.75	0.65
	CAT	0.77	44.03	0.79	0.84	0.82

4.3 PhysioNet Sepsis

To evaluate CAT’s critic performance beyond reinforcement learning and language, we assessed its ability to predict sepsis in ICU patients. We evaluated two sequence-level attributes: a binary label for sepsis occurrence during the ICU course, and a continuous label estimating maximum heart rate (HR) within a 6-hour sliding window. Assessed 12 hours prior to onset, CAT demonstrated comparable predictive performance compared to MC simulation of the standard next-token model and superior to Director (Fig. 5A). Although the standard model achieved a slightly higher ROC AUC, CAT provides a substantial improvement in Average Precision (AP), indicating higher precision at clinically relevant recall levels in this highly imbalanced prediction setting.

To evaluate sensitivity to counterfactual vital sign substitutions, we isolated the first temperature reading in sepsis-positive validation sequences and measured stepwise changes in predicted sepsis probability across temperature bins (Fig. 5B). Elevated temperature increased predicted sepsis risk, with substantially larger hyperthermic risk increases among older patients, consistent with their greater vulnerability to thermal dysregulation Brody (1994). A representative case illustrates CAT’s dynamic sepsis risk estimates alongside conditional forecasts of maximum HR over the subsequent six hours (Fig. 5C). Risk is assessed at each token, demonstrating an implicitly learned sensitivity to established clinical criteria. Extracting the hour with the largest increase in $P(\mathrm{sepsis})$ reveals the token-level contributions driving this risk change (Fig. 5D). Large MAP-associated risk increases following low DBP may reflect discordance between MAP, SBP, and DBP measurements in the data set (Fig. A.4).

5 Discussion

We present Conditional Attribute Transformers (CAT), a novel framework to jointly model next-token probabilities and sequence level attributes for credit assignment, counterfactual estimation, and guided decoding. We find that with large enough models, CAT improves next token prediction, suggesting that forcing the model to learn global structure can complement rather than compromise next-token prediction. This mirrors results from the multi-token prediction training objective utilized by DeepSeek Liu et al. (2024a), where forecasting future states acts as an auxiliary training objective, regularizing representations and encouraging broader structural planning. In addition, CAT’s capabilities can be extended to pretrained models through fine-tuning only the conditional attribute head.

Across a range of settings and tasks, CAT consistently outperformed baseline models. In the Key-to-Door task, CAT learns long-term credit assignment from sparse terminal rewards and achieves near-perfect performance, demonstrating its ability to identify actions that determine delayed outcomes. In language modeling, CAT effectively captures shifts in review sentiment, enables efficient partial-sequence rating prediction, and supports attribute-guided generation. In biomedical informatics, CAT estimates elevated sepsis risk well before diagnosis is recorded, supports identification and examination of the drivers of major inflection points in predicted risk, and characterizes subgroup-specific changes in risk associated with clinical variables. While CAT offers profound societal benefits, particularly by advancing explainability through counterfactual analysis, it also carries inherent risks, as its steering capabilities could be exploited to introduce harmful or malicious biases.

Limitations: This method is currently limited to generating conditional attribute probabilities over discrete action choices, but there may be settings where a continuous attribute is desired. For steering problems, current counterfactual estimates represent single-step policy improvement and yield a greedy optimum under an average behavior policy rather than a global optimum under a specific policy. This limitation is shared by prior approaches such as Director and DExperts.

Future Work: Future work includes scaling this method to larger data sets and extending it to continuous action spaces. We are currently developing approaches to enable global optimal choice selection. Furthermore, this framework could be naturally extended to any task which requires estimation or control of sequence-level properties, including the biological domain (de-novo protein design, predicting small-molecule binding functionality, predicting sequence-to-function regulatory mechanism from DNA, etc.).

Ultimately, by jointly modeling local token generation and global sequence level objectives, CAT provides a highly adaptable foundation for complex predictive and generative tasks.

6 Acknowledgements

This work was supported by the United States National Library of Medicine (grant T15 LM007056 to ES and GM) and National Institutes of Health (grants 5UL1TR001863 to DM and R00HG013661 to QL). AJL receives funding through UL1 TR001863, The Hartwell Foundation, and the ARIA foundation.

References

[1] K. Arora, K. Shuster, S. Sukhbaatar, and J. Weston (2022) Director: generator-classifiers for supervised language modeling. arXiv preprint arXiv:2206.07694. Cited by: §A.2, §2, §4.2.1.
[2] M. G. Bellemare, W. Dabney, and R. Munos (2017) A distributional perspective on reinforcement learning. External Links: 1707.06887, Link Cited by: §3.4.
[3] G. Brixi, M. G. Durrant, J. Ku, M. Poli, G. Brockman, D. Chang, G. A. Gonzalez, S. H. King, D. B. Li, A. T. Merchant, et al. (2025) Genome modeling and design across all domains of life with evo 2. BioRxiv, pp. 2025–02. Cited by: §1.
[4] G. M. Brody (1994-02) Hyperthermia and hypothermia in the elderly.. Clinics in geriatric medicine 10 (1), pp. 213–229. External Links: Link Cited by: §4.3.
[5] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. (2020) Language models are few-shot learners. Advances in neural information processing systems 33, pp. 1877–1901. Cited by: §1.
[6] Y. Chang, X. Wang, J. Wang, Y. Wu, L. Yang, K. Zhu, H. Chen, X. Yi, C. Wang, Y. Wang, W. Ye, Y. Zhang, Y. Chang, P. S. Yu, Q. Yang, and X. Xie (2023) A survey on evaluation of large language models. arXiv preprint arXiv:2307.03109. External Links: Link Cited by: §1.
[7] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch (2021) Decision transformer: reinforcement learning via sequence modeling. Advances in neural information processing systems 34, pp. 15084–15097. Cited by: §2, item 1, §3.6, §4.1.
[8] T. Chen, T. He, M. Benesty, V. Khotilovich, Y. Tang, H. Cho, K. Chen, R. Mitchell, I. Cano, T. Zhou, M. Li, J. Xie, M. Lin, Y. Geng, Y. Li, J. Yuan, and D. Cortes (2026) Xgboost: extreme gradient boosting. R package version 3.3.0.0. External Links: Link Cited by: §A.2.
[9] S. Dathathri, A. Madotto, J. Lan, J. Hung, E. Frank, P. Molino, J. Yosinski, and R. Liu (2019) Plug and play language models: a simple approach to controlled text generation. arXiv preprint arXiv:1912.02164. Cited by: §2.
[10] N. Ferruz, S. Schmidt, and B. Höcker (2022) ProtGPT2 is a deep unsupervised language model for protein design. Nature communications 13 (1), pp. 4348. Cited by: §1.
[11] F. Gloeckle, B. Y. Idrissi, B. Rozière, D. Lopez-Paz, and G. Synnaeve (2024) Better & faster large language models via multi-token prediction. External Links: 2404.19737, Link Cited by: §1.
[12] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre (2022) Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. External Links: Link Cited by: §1.
[13] Y. Hou, J. Li, Z. He, A. Yan, X. Chen, and J. McAuley (2024) Bridging language and items for retrieval and recommendation. arXiv preprint arXiv:2403.03952. Cited by: §A.2, item 2, §3.6.
[14] A. Karpathy (2023) NanoGPT: a minimalistic and educational gpt training code. Note: https://github.com/karpathy/nanoGPTGitHub repository Cited by: §3.5.
[15] N. S. Keskar, B. McCann, L. R. Varshney, C. Xiong, and R. Socher (2019) Ctrl: a conditional transformer language model for controllable generation. arXiv preprint arXiv:1909.05858. Cited by: §1, §2.
[16] B. Krause, A. D. Gotmare, B. McCann, N. S. Keskar, S. Joty, R. Socher, and N. F. Rajani (2020) GeDi: generative discriminator guided sequence generation. arXiv preprint arXiv:2009.06367. Cited by: §2.
[17] A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al. (2024) Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437. Cited by: §3.5, §5.
[18] A. Liu, M. Sap, X. Lu, S. Swayamdipta, C. Bhagavatula, N. A. Smith, and Y. Choi (2021) DExperts: decoding-time controlled text generation with experts and anti-experts. arXiv preprint arXiv:2105.03023. Cited by: §A.2, §2.
[19] Q. Liu, Z. Chen, and W. H. Wong (2024) An encoding generative modeling approach to dimension reduction and covariate adjustment in causal inference with observational studies. Proceedings of the National Academy of Sciences 121 (23), pp. e2322376121. Cited by: §3.4.
[20] Q. Liu and W. H. Wong (2026) A bayesian generative modeling approach for arbitrary conditional inference. arXiv preprint arXiv:2601.05355. Cited by: §3.4.
[21] Q. Liu and W. H. Wong (2026) An ai-powered bayesian generative modeling approach for causal inference in observational studies. Journal of the American Statistical Association (just-accepted), pp. 1–20. Cited by: §3.4.
[22] X. Lu, S. Welleck, J. Hessel, L. Jiang, L. Qin, P. West, P. Ammanabrolu, and Y. Choi (2022) Quark: controllable text generation with reinforced unlearning. Advances in neural information processing systems 35, pp. 27591–27609. Cited by: §2.
[23] V. Melnychuk, D. Frauen, and S. Feuerriegel (2022) Causal transformer for estimating counterfactual outcomes. In International conference on machine learning, pp. 15293–15329. Cited by: §3.4.
[24] T. Mesnard, T. Weber, F. Viola, S. Thakoor, A. Saade, A. Harutyunyan, W. Dabney, T. Stepleton, N. Heess, A. Guez, É. Moulines, M. Hutter, L. Buesing, and R. Munos (2021) Counterfactual credit assignment in model-free reinforcement learning. External Links: 2011.09464, Link Cited by: item 1.
[25] W. Qi, Y. Yan, Y. Gong, D. Liu, N. Duan, J. Chen, R. Zhang, and M. Zhou (2020) ProphetNet: predicting future n-gram for sequence-to-sequencepre-training. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 2401–2410. Cited by: §1.
[26] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever (2019) Language models are unsupervised multitask learners. Cited by: §A.2.
[27] S. M. Rein, J. Li, M. Hernan, and A. Beam (2024) Deep learning methods for the noniterative conditional expectation g-formula for causal inference from complex observational data. arXiv preprint arXiv:2410.21531. Cited by: §3.4.
[28] P. Renc, Y. Jia, A. E. Samir, J. Was, Q. Li, D. W. Bates, and A. Sitek (2024) Zero shot health trajectory prediction using transformer. NPJ digital medicine 7 (1), pp. 256. Cited by: §1.
[29] M. Reyna, C. Josef, R. Jeter, S. Shashikumar, B. Moody, M. B. Westover, A. Sharma, S. Nemati, and G. D. Clifford (2019) Early prediction of sepsis from clinical data: the PhysioNet/Computing in Cardiology Challenge 2019. PhysioNet. Note: RRID:SCR_007345 External Links: Document, Link Cited by: §A.3, item 3, §3.6.
[30] C. Shi, D. M. Blei, and V. Veitch (2019) Adapting neural networks for the estimation of treatment effects. arXiv preprint arXiv:1906.02120. Cited by: §3.4.
[31] A. Shmatko, A. W. Jung, K. Gaurav, S. Brunak, L. H. Mortensen, E. Birney, T. Fitzgerald, and M. Gerstung (2025) Learning the natural history of human disease with generative transformers. Nature 647 (8088), pp. 248–256. Cited by: §1.
[32] H. A. Simon (1955) A behavioral model of rational choice. Quarterly Journal of Economics 69 (1), pp. 99–118. Cited by: §A.2.
[33] C. Snell, I. Kostrikov, Y. Su, M. Yang, and S. Levine (2022) Offline rl for natural language generation with implicit language q learning. arXiv preprint arXiv:2206.11871. Cited by: §2, §3.4.
[34] S. Waxler, P. Blazek, D. White, D. Sneider, K. Chung, M. Nagarathnam, P. Williams, H. Voeller, K. Wong, M. Swanhorst, et al. (2025) Generative medical event models improve with scale. arXiv preprint arXiv:2508.12104. Cited by: §1.
[35] G. Y. Weng, B. Wang, and G. V. d. Broeck (2025) TRACE back from the future: a probabilistic reasoning approach to controllable language generation. arXiv preprint arXiv:2504.18535. Cited by: §2.
[36] H. Xiong, F. Wu, L. Deng, M. Su, Z. Shahn, and L. H. Lehman (2024) G-transformer: counterfactual outcome prediction under dynamic and time-varying treatment regimes. Proceedings of machine learning research 252, pp. https–proceedings. Cited by: §3.4.
[37] K. Yang and D. Klein (2021) FUDGE: controlled text generation with future discriminators. arXiv preprint arXiv:2104.05218. Cited by: §2.

Appendix A Appendix

All experiments were run on a computing cluster with a combination of NVIDIA H100 and H200 and RTX5000 for approximately 3,000 GPU hours. Training and inference runtime varies with model architecture, model size, input sequence length, and hardware configuration (Table A.1, A.2, and A.4).

A.1 Key-to-Door: Long-Term Credit Assignment

Model and Training Details:

Model specifications and hyperparameters are outlined in Fig. A.1. CAT and baseline models were trained on 10,000 random walk trajectories and evaluated on 1,000 random starts. Baseline models included random policy, behavioral cloning, percentile behavior cloning (trained on winning trajectories only), conservative Q-learning, and decision transformers.

Table A.1: Model configurations for the Key-to-Door task.

Model	Variant	Params	Layers	Dim	Heads	MLP Dim	Context	LR
Random Policy	Base	3M	8	128	8	512	114	3e-3
Behavioral Cloning	Base	3M	8	128	8	512	114	3e-3
Percentile Behavioral Cloning	Base	3M	8	128	8	512	114	3e-3
Conservative Q-Learning	Base	3M	8	128	8	512	114	3e-3
Decision Transformers	Base	3M	8	128	8	512	114	3e-3
CAT	Base	3M	8	128	8	512	114	3e-3

A.2 Language Modeling: Amazon Reviews

Data Set:

The Amazon Reviews data set (CC0: Public Domain) is a large-scale corpus of 574 million product reviews, each consisting of a product category, title, review text, and a 1–5 star rating [13]. Reviews were tokenized using a 32k BPE vocabulary and formatted as <|sos|><|category|><|sotitle|>title text<|sotext|>review text<|sor|><|*_R|>, where R denotes the rating. The data set was split 90:10 into training and validation sets. The distribution of ratings for 1–5 stars in the training set was 10.3%, 4.9%, 7.1%, 12.7%, and 65.0%, respectively, with a similar distribution observed in the validation set.

Model and Training Details:

We trained models ranging from 7-million to 1-billion parameters (Table A.2). For each model size, one standard decoder-only baseline and seven CAT variants with different values of $\lambda$ in Eq. 14 were trained. CAT was also extended to a pre-trained standard model by fine-tuning only the attribute head. For DExperts, a 1-billion parameter standard model was used as the base model and a $\lambda=0.15$ was used, as it outperformed the previously reported optimum of 0.2 on this dataset [1].

Table A.2: Model configurations for the language modeling task.

Model	Variant	Params	Layers	Dim	Heads	MLP Dim	Context	LR
GPT	Base	7M	2	96	2	384	512	3e-4
GPT	Base	72M	12	512	8	2048	512	3e-4
GPT	Base	270M	16	1024	16	4096	512	3e-4
GPT	Base	1B	32	1536	32	6144	512	2.5e-4
CTRL	Base	1B	32	1536	32	6144	512	2.5e-4
DExperts	1 star	1B	32	1536	32	6144	512	2.5e-4
DExperts	5 star	1B	32	1536	32	6144	512	2.5e-4
Director	Binary (1 star)	1B	32	1536	32	6144	512	2.5e-4
Director	Binary (5 star)	1B	32	1536	32	6144	512	2.5e-4
Director*	Base	1B	32	1536	32	6144	512	2.5e-4
CAT	Base	7M	2	96	2	384	512	3e-4
CAT	Base	72M	12	512	8	2048	512	3e-4
CAT	Base	270M	16	1024	16	4096	512	3e-4
CAT	Base	1B	32	1536	32	6144	512	2.5e-4
CAT	Fine-tuned	1B	32	1536	32	6144	512	5e-2
CAT	Attribute-only	1B	32	1536	32	6144	512	5e-2

Evaluation Details:

Next-token validation perplexity was evaluated over 2,000 iterations for the standard and CAT models across model sizes.

Counterfactual adjective substitution was evaluated on 1,000,000 validation reviews in which the true next token was good. Results are reported for all reviews and for the subset of 1,593 reviews with a preceding negation term or phrase: not, no, never, not really, isn’t, ain’t, wasn’t, weren’t. Fig. A.2 visualizes the conditional attribute estimates at each token position. This product is amazing, This product is good, This product is bad, and This product is horrible produce similar probability distributions up to the adjective token, after which the predicted rating probabilities diverge according to adjective sentiment. Positive adjectives increase the predicted probability of 5 star ratings, whereas negative adjectives increase the predicted probability of 1 star ratings. Stronger adjectives (amazing, horrible) produce larger shifts than weaker adjectives (good, bad). A similar pattern is observed under negation. The reviews produce similar distributions up to not, after which the predicted probability of a 5 star rating decreases and that of a 1 star rating increases, indicating that the model has learned to associate negation with more negative sentiment and lower ratings. Adjectives following negation produce effects consistent with those of negation shown in Table 2. This product is not good resembles This product is bad with an increased predicted probability of a 1 star rating, while not bad, not amazing, and not horrible are dominated by a 3 star rating, consistent with indicating okay.

For the steering experiment, validation reviews with true 3 star ratings were split after the first 50% of the text following the <|sotext|> token. The resulting first-half prompts were scored using an XGBoost rating classifier trained on 10,000,000 validation reviews with CountVectorizer features (Fig. A.3) [8]. Prompts classified as 3 stars were retained, from which 1,000 were randomly selected. Steering was performed using a satisficing criterion with top- $k$ decoding for CAT, Director, and Director* models. At each step, decoding was restricted to tokens with a next-token probability above $\epsilon$ and a conditional attribute probability above the specified attribute threshold [32]. DExperts scales the difference between expert and anti-expert logits with a control parameter $\alpha$ ; following prior work, we use $\alpha=0.2$ for steering [18]. Hyperparameters for all model types are listed in Table A.3. Each method generated 10 completions per prompt, and the resulting full text, consisting of the prompt and completion, was re-scored by XGBoost. Performance is reported as the proportion of full texts classified as either 1 or 5 stars. Fluency of the completion was measured by mean perplexity using a Hugging Face GPT-2 XL model and tokenizer [26], and diversity of the completion was measured using normalized unique $n$ -gram counts. Evaluation metrics closely follow those used for the toxicity task in [18].

Table A.3: Steering hyperparameters.

Method	$\boldsymbol{k}$	$\boldsymbol{\alpha}$	$\boldsymbol{\epsilon}$	Attribute Threshold	Sampling Method
CTRL	20	–	–	–	Top- $k$
DExperts	20	0.2	–	–	Top- $k$
Director	20	–	0.001	0.8	Top- $k$
Director*	20	–	0.001	0.8	Top- $k$
CAT	20	–	0.001	0.8	Top- $k$

A.3 PhysioNet Sepsis

Data Set:

The 2019 PhysioNet Challenge (CC-BY-4.0) provides a data set from 40,336 patients collected from the Intensive Care Units (ICUs) of three hospital systems as a benchmark for early sepsis identification [29]. This data set consists of patient demographic information and hourly measurements of vital signs and laboratory values. Each ICU course is annotated with whether sepsis occurred and its time of onset. Records were converted into sequences of discrete tokens. Patient demographics were included at the beginning of each sequence. Categorical measurements received a unique discrete token and each numeric measurement received binned value tokens. Values were discretized into deciles based on per-class distribution of values in the training set. A time token was inserted between each set of hourly measurements. Data was split 90:10 into training and validation sets, with a sepsis prevalence of 7.3% in training and 7.4% in validation. Model hyperparameters are listed in Table A.4.

Model and Training Details:

The optimal $\lambda$ used to balance the contribution of next-token and target attribute loss differed in the binary attribute prediction and regression target in the sepsis data set. We selected a $\lambda$ of 0.01 for binary sepsis risk prediction and 0.5 for predicting the regression target of the max heart rate in the next 6 hours. The large difference in these values highlights how predicting the next token correctly is essential for sepsis risk prediction, whereas a larger contribution of target regression loss is necessary to learn the more volatile value of heart rate.

Table A.4: Model configurations for the sepsis prediction task.

Model	Variant	Params	Layers	Dim	Heads	MLP Dim	Context	LR
GPT	Base	29M	8	512	8	2048	1024	5e-4
Director	Base	29M	8	512	8	2048	1024	5e-4
CAT	Binary	29M	8	512	8	2048	1024	5e-4
CAT	Regression	29M	8	512	8	2048	1024	5e-5