Longitudinal Targeted Minimum Loss-based Estimation with Temporal-Difference Heterogeneous Transformer

We consider the general longitudinal setting involving repeated measurements of a set of variables for a group of $n$ patients over a period of time. In particular, our observed data contains $n$ independent and identically distributed copies of random vector

with baseline covariates $W$, time-dependent covariates $L_{t}$, treatments $A_{t}$, and outcome $Y_{t}$. We use $P_{0}$ to denote the true probability distribution of $O$ that generated the data, and $P_{0}$ is in some statistical model $\mathcal{M}$. Stop** time $T$ is a random variable (e.g. time of death in the case of survival analysis) and we use $\tau$ to denote the maximum time. We make the remark that in real-world data, patients are often subject to censoring. For a formulation of the data structure involving censoring nodes, see Appendix H.

To define the target parameter, we introduce a structural causal model (SCM). In brief, SCM assumes each observed random variable $X$ is generated from the parent nodes $pa(X)$ and the external noise $U_{X}$ by a production function $f_{X}$ as $X=f_{X}(pa(X),U_{X})$. By abusing notation, we also denote the induced probability measure of $X$ by the same symbol $f_{X}$. See Appendix C.1 for details.

Our target parameter is the counterfactual mean of the final outcome $Y$ under a user-specified dynamic treatment policy $g=[g_{t}]_{t=1}^{\tau}$ where $g_{t}$ is a probability measure on the treatment space conditioned on the whole history, $pa(A_{t})=(L_{1:t},A_{1:t-1},Y_{1:t-1})$ up until $A_{t}$ (not including $A_{t}$). Specifically, our target parameter is given by

which is the mean of the counterfactual outcome produced by replacing $\pi$, defined as the observed treatment policy from the data, with $g$ in the structural causal model.

Given we have an estimator $\hat{P}_{n}$ of the data-generating distribution $P_{0}$, a natural estimator of the target functional is the plug-in estimator $\hat{\psi}_{n}=\psi(\hat{P}_{n})$. Under a regularity condition, $\psi$ admits the following first-order expansion

where $D^{\star}$ is the efficient influence function of $\psi$, and $R_{2}(\hat{P}_{n},P_{0})$ is the exact remainder. Influence functions quantifies the amount of changes of an estimator under small perturbations of the sample. The efficient influence function is the influence function with minimal variance. The idea of TMLE is to eliminate the empirical analogue of the first term of the right hand side by fluctuating $\hat{P}_{n}$ to find a distribution $\hat{P}^{\star}_{n}$ with $P_{n}D^{\star}(\hat{P}^{\star}_{n})=0$, where $Pf=\int fdP$ is a shorthand for the expectation of a measurable function $f$ with respect to a probability measure $P$. Our problem is to obtain an initial estimate $\hat{P}_{n}$ with a potentially large scale and high dimensional longitudinal data, and correct bias of the plug-in estimator $\psi(\hat{P}_{n})$ by fluctuating $\hat{P}_{n}$.

To learn the initial model $\hat{V}^{g}_{1}$, we use temporal-difference loss as the objective to learn underlying models $\hat{Q}^{g}_{t}$ for $t=1,\ldots,\tau$ via stochastic gradient descent (SGD). The principle of temporal difference learning (Sutton, 1988; Mnih et al., 2013) is to supervise $\hat{Q}^{g}_{t}$ to obey the temporal equality of $Q^{g}_{t}$:

for $t=1,\ldots,\tau-1$ and $Q^{g}_{\tau}=\mathbb{E}_{p(Y_{\tau}\mid pa(Y_{\tau}))}Y_{\tau}$. The temporal difference loss on a sample trajectory is thus given by $\mathcal{L}^{Q}_{t}=\mathcal{L}(\hat{Q}^{g}_{t},\hat{V}^{g}_{t+1})$ for $t=1,\ldots,T-1$, where $\hat{V}^{g}_{t+1}(pa(A_{t+1}))=\mathbb{E}_{A_{t+1}\sim g_{t+1}(pa(A_{t+1}))}\big{[}\hat{Q}^{g}_{t+1}(pa(A_{t+1}),A_{t+1})\big{]}$ can be computed by Monte-Carlo estimation if $A$ is continuous, and $\mathcal{L}^{Q}_{T}=\mathcal{L}(\hat{Q}^{g}_{T},Y_{T})$. In the case of survival analysis, the components for $t=1,\ldots,\tau-1$ are defined as $\mathcal{L}^{Q}_{t}=\mathcal{L}_{\mathrm{bce}}(\hat{Q}^{g}_{t},\hat{V}^{g}_{t+1})$, where $\mathcal{L}_{\mathrm{bce}}(\hat{y},y)=-\big{[}y\log\hat{y}+(1-y)\log(1-\hat{y})\big{]}$ is the binary cross entropy loss. To yield the updated model $\hat{Q}^{g}_{t,\varepsilon}$, we need to adjust $\hat{Q}^{g}_{t}$ after model training factoring in the estimating model $\hat{G}_{t}$ for the propensity score

which we will describe in detail in the next section. Hence, the loss function also needs to include $\mathcal{L}^{G}_{t}=\mathcal{L}(\hat{G}_{t},A_{t})$ and is thus given by

where $\alpha$ is a hyperparameter that controls the weights of losses. See Algorithm 1 for the optimization workflow. Convergence of the algorithm can be found in Appendix D.

For the estimation of $Q^{g}_{t}$ and $G_{t}$, we propose a unified model architecture to simultaneously optimize deep neural networks $\hat{Q}^{g}_{t}$ and $\hat{G}_{t}$ in an efficient, non-sequential manner by adapting a decoder-only Transformer (Vaswani et al., 2017; Brown et al., 2020) to longitudinal data with heterogeneous tokens. An overview of the model architecture is given in Figure 1. For each sampled sequence in the training set, we feed each token in the sequence to a linear embedding layer according to its variable type. In the case of Figure 1, there are four different embedding layers $e_{W}$, $e_{L}$, $e_{A}$, and $e_{Y}$. Each embedding layer has the same number of output dimensions. Then, each embedding is integrated with its positional encoding $E_{0},\ldots,E_{\tau}$ and type encoding $E_{W}$, $E_{L}$, $E_{A}$, and $E_{Y}$ that represent its timestamp and variable type information through an aggregation function (e.g. sum, concat):

for $\bullet_{t}\in(W_{0},L_{1},A_{1},Y_{1},\ldots,L_{\tau},A_{\tau},Y_{\tau})$ where we used concat as aggr in the experiments in this work. Note that we include type embedding $E_{\bullet}$ because $e_{\bullet}$ need not necessarily be type-specific linear layers. For more efficient and parallelizable embedding operation, we can pad each variable to the same number of dimensions before feeding into the same embedding layer $e_{\bullet}=e$. Then, the embedded sequence is fed into the transformer and produce $\hat{G}$ and $\hat{Q}^{g}$ through output heads $f_{G}$ and $f_{Q}$ at each position that corresponds to token type $L$ and $A$ respectively:

In practice, we can use a joint output layer $f$ for $f_{G}$ and $f_{Q}$ for more efficient and parallelizable output generation, where the output number of dimensions is the sum of the number of dimensions $\mathrm{dim}(A)$ for treatment $A$ and $\mathrm{dim}(Y)$ for outcome $Y$. Then, we compute softmax probabilities masking out the last $\mathrm{dim}(Y)$ dimensions for $\hat{G}_{t}$ and first $\mathrm{dim}(A)$ dimensions for $\hat{Q}^{g}_{t}$.

Our proposed architecture does not entail concatenation of variables at the same timestamp or sequential decoding of outputs following the transformer embedding block like prior work Melnychuk et al. (2022), which 1) allows us to handle different types of and different number of variables at different timestamps (e.g. starting from $W_{0}$, ending at $L_{8}$, while $A_{3}$ and $Y_{3}$ are missing), and 2) is fully parallelizable when we use padding instead of learnable linear map** for the embedding layer $e_{\bullet}$ and use the joint output layer $f$.