Adaptive Test-Time Training for Predicting Need for Invasive Mechanical Ventilation in Multi-Center Cohorts
Extracted source text · arXiv 2512.06652
\documentclass{article}
\usepackage{PRIMEarxiv}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{lipsum}
\usepackage{fancyhdr} % header
\usepackage{graphicx} % graphics
\graphicspath{{media/}} % organize your images and other figures under media/ folder
\usepackage{times}
\usepackage{latexsym}
% For proper rendering and hyphenation of words containing Latin characters (including in bib files)
\usepackage[T1]{fontenc}
% For Vietnamese characters
% \usepackage[T5]{fontenc}
% See https://www.latex-project.org/help/documentation/encguide.pdf for other character sets
% This assumes your files are encoded as UTF8
\usepackage[utf8]{inputenc}
% This is not strictly necessary, and may be commented out.
% However, it will improve the layout of the manuscript,
% and will typically save some space.
\usepackage{microtype}
% This is also not strictly necessary, and may be commented out.
% However, it will improve the aesthetics of text in
% the typewriter font.
\usepackage{inconsolata}
\usepackage{microtype}
\usepackage{graphicx}
%\usepackage{subfigure}
%\usepackage{booktabs} % for professional tables
% hyperref makes hyperlinks in the resulting PDF.
% If your build breaks (sometimes temporarily if a hyperlink spans a page)
% please comment out the following usepackage line and replace
% \usepackage{icml2021} with \usepackage[nohyperref]{icml2021} above.
\usepackage{hyperref}
% Attempt to make hyperref and algorithmic work together better:
\newcommand{\theHalgorithm}{\arabic{algorithm}}
% Use the following line for the initial blind version submitted for review:
\usepackage[textsize=tiny]{todonotes}
%\usepackage{inconsolata}
\usepackage{amssymb}
\usepackage{natbib}
%%%%%%%
\usepackage{adjustbox}
\usepackage{graphicx}
%\usepackage{algorithm}
\usepackage{caption}
%\captionsetup[table]{labelsep=period}
\usepackage{subcaption}
\usepackage{amsmath}
\usepackage{bm}
\usepackage{multirow}
\usepackage{xcolor}
\usepackage[utf8]{inputenc}
\usepackage{mathtools}
\usepackage{bbm}
\usepackage{physics}
\usepackage{float}
\usepackage{booktabs}
\usepackage{array, makecell}
\usepackage{algorithm}
\usepackage{algorithmic}
%\usepackage{algpseudocode}
\usepackage{tabularx}
\usepackage{amsthm}
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}[theorem]{Lemma}
\usepackage{hyperref}
%Header
\pagestyle{fancy}
\thispagestyle{empty}
\rhead{ \textit{ }}
% Update your Headers here
%\fancyhead[LO]{Running Title for Header}
% \fancyhead[RE]{Firstauthor and Secondauthor} % Firstauthor et al. if more than 2 - must use \documentclass[twoside]{article}
%% Title
\title{Adaptive Test-Time Training for Predicting Need for Invasive Mechanical Ventilation in Multi-Center Cohorts
%%%% Cite as
%%%% Update your official citation here when published
%\thanks{\textit{\underline{Citation}}:
%\textbf{Authors. Title. Pages.... DOI:000000/11111.}}
}
\author{
Xiaolei Lu,
%Affiliation \\
%Univ \\
%City\\
Shamim Nemati
%% examples of more authors
}
\begin{document}
\maketitle
\begin{abstract}
Accurate prediction of the need for invasive mechanical ventilation (IMV) in intensive care units (ICUs) patients is crucial for timely interventions and resource allocation. However, variability in patient populations, clinical practices, and electronic health record (EHR) systems across institutions introduces domain shifts that degrade the generalization performance of predictive models during deployment. Test-Time Training (TTT) has emerged as a promising approach to mitigate such shifts by adapting models dynamically during inference without requiring labeled target-domain data. In this work, we introduce Adaptive Test-Time Training (AdaTTT), an enhanced TTT framework tailored for EHR-based IMV prediction in ICU settings. We begin by deriving information-theoretic bounds on the test-time prediction error and demonstrate that it is constrained by the uncertainty between the main and auxiliary tasks. To enhance their alignment, we introduce a self-supervised learning framework with pretext tasks: reconstruction and masked feature modeling optimized through a dynamic masking strategy that emphasizes features critical to the main task. Additionally, to improve robustness against domain shifts, we incorporate prototype learning and employ Partial Optimal Transport (POT) for flexible, partial feature alignment while maintaining clinically meaningful patient representations. Experiments across multi-center ICU cohorts demonstrate competitive classification performance on different test-time adaptation benchmarks.
\end{abstract}
% keywords can be removed
%\keywords{Explainability \and Faithfulness\and Plausibility}
\section{Introduction}
Invasive mechanical ventilation (IMV) is a critical intervention utilized in intensive care units (ICUs) for patients with severe respiratory failure and acute respiratory distress syndrome (ARDS) \cite{grotberg2023management}. However, its use is complicated by the risk of ventilator-induced lung injury and complications resulting from prolonged IMV. Timely and accurate identification of patients at high risk for mechanical ventilation is crucial for optimizing clinical decision-making. Early recognition of these patients enables proactive medical interventions and facilitates efficient resource allocation within hospital system \citep{fan2018acute}.
In recent years, the development of machine learning (ML) models has shown great promise in predicting the need for IMV, which leverages electronic health record (EHR) data to identify complex patterns that human clinicians might overlook \citep{shashikumar2021development}. These models can incorporate diverse features, including vital signs and laboratory results to enhance prediction accuracy and provide critical decision support in ICU settings. However, the effective deployment of such models in real-world clinical settings remains a challenge. A key issue is the variability in data distributions across hospitals due to differences in patient populations, clinical practices, and EHR systems. These shifts, often referred to as domain shifts, can substantially degrade the performance of predictive models that were trained on data from a single or limited number of sources. For instance, a respiratory failure prediction model \citep{lam2024development} trained on ICU cohort from UC San Diego Health showed an approximately 12\% drop in the area under the curve (AUC) when evaluated on an external ICU cohort.
Addressing this challenge requires adaptive methodologies that can account for site-specific heterogeneity. Existing approaches include pre-training models on large multi-center datasets \citep{shashikumar2021development}, and transfer learning to fine-tune models on site-specific data \citep{lam2024development} to align feature representations across domains. While these methods have shown promise, many require access to labeled data from the target domain during training or involve computationally expensive retraining processes, which are not always feasible in real-time clinical settings. Test-time training (TTT) offers a novel and efficient solution to this problem by enabling models to adapt dynamically at the time of prediction, without requiring pre-access to target domain labels or costly re-training. TTT leverages an auxiliary task, trained alongside the main task, to update the model’s parameters or representations using the test input itself.
Predictive systems based on EHR data, such as Composer \citep{shashikumar2021artificial}, have shown significant real-world
impact to improve clinical outcomes through real-time decision support \citep{boussina2024impact}. In a before-and-after quasi-experimental design study at two emergency departments (EDs), the Composer model for sepsis prediction significantly increased bundle compliance and reduced in-hospital mortality. However, limited prior work has explored TTT in the context of EHR data, particularly in real-world clinical scenarios. By enabling dynamic adaptation at prediction time, TTT addresses variability across
institutions and patient populations, ensuring robust performance in critical tasks, such as predicting IMV need in multi-center cohorts. Unlike pretraining/offline alignment approaches (e.g., MaskTab \citep{chenmasktab}, PhyMask \citep{kara2024phymask}, and SPOT \citep{gurumoorthy2021spot}), which require source data and extended training, our setting is source-free and demands on-the-fly adaptation at deployment.
In this study, we introduce Adaptive Test-Time Training (AdaTTT) for predicting IMV need 24 hours in advance in ICU patients across multi-center cohorts. AdaTTT is designed to address domain shifts in EHR data through adaptive self-supervised learning and robust feature alignment. Our key contributions are as follows:
\begin{itemize}
\item We derive information-theoretic bounds on the test-time prediction error to show that the error is constrained by the uncertainty between the main and auxiliary tasks, which guide the design of auxiliary tasks for better adaptation.
\item We introduce two SSL tasks: Reconstruction and Masked Feature Modeling along with a dynamic masking strategy that prioritizes the most informative features during test-time training. The masking probabilities adapt based on feature relevance to the primary task, ensuring that the SSL task remains aligned with the IMV prediction objective.
\item To prevent overfitting to individual test samples, we integrate prototype learning with Partial Optimal Transport (POT) to allow partial matching between source domain features and test-time distributions, which promotes robust generalization while avoiding excessive adaptation to test-domain noise.
\item We conduct extensive experiments on multi-site ICU cohorts, where our method achieves competitive classification performance across various test-time adaptation benchmarks.
\end{itemize}
\section{Related Work}
\subsection{Predictive Models for Invasive Mechanical Ventilation}
Early IMV-risk tools such as ROX and regression scores are interpretable but struggle with nonlinear, time-varying physiology \citep{roca2019index}. Leveraging EHR-scale data, VentNet predicts IMV 24h ahead with a feedforward model \citep{shashikumar2021development}; encoder–decoder designs like DBNet integrate structured signals and demographics \citep{zhang2021dbnet}; and multimodal hybrids that fuse CXR with EHR further boost discrimination \citep{tandon2023hybrid}. However, cross-site performance often degrades due to population, workflow, and EHR heterogeneity; recovery via target-domain fine-tuning is common but label-intensive and operationally impractical for continuous deployment.
\subsection{Test-Time Adaptation}
TTA adapts models on unlabeled test inputs without revisiting source data \citep{liang2024comprehensive}. Batch normalization(BN)-centric methods include prediction-time BN statistics updates \citep{nado2020evaluating} and TENT's entropy minimization for BN parameters \citep{wang2020tent}, while source-free SHOT freezes the classifier and adapts the encoder with pseudo-labels \citep{liang2020we}. Test-time training (TTT) attaches auxiliary SSL branches for online encoder updates \citep{sun2020test}; extensions like TTT++ (contrastive) and ClusT3 (clustering) improve alignment but may inherit instability or assume domain consistency \citep{liu2021ttt++,hakim2023clust3}.
Three relevant directions are: (i) T3A, an optimization-free method forming class prototypes from streaming test data to reweight logits that is efficient but classifier-level only, assuming stable class structure \citep{iwasawa2021test}; (ii) SAR, which filters unreliable samples and applies sharpness-aware entropy minimization for stable BN updates that is effective with small batches yet BN-dependent \citep{niu2023towards}; (iii) CoTTA, maintaining a moving teacher with augmentation- and weight-averaged pseudo-labels plus periodic restoration, is useful for long horizons but hyperparameter-sensitive with potential error accumulation \citep{wang2022continual}.
In EHR-driven IMV prediction, these approaches face practical challenges: tiny per-encounter batches undermine BN estimates; pseudo-labeling struggles with class imbalance and temporal nonstationarity; clustering assumptions break under irregular sampling and missingness; classifier-only adaptation cannot address representation shift, making direct application from vision to ICU EHRs difficult.
\section{Methodology}
\subsection{Preliminary: Test-Time Training}
Let \( x \in \mathcal{X} \) denote an input instance from the covariate space, \( y_m \in \mathcal{Y}_m \) denote the corresponding label for the main task (e.g., classification), and \( y_s \in \mathcal{Y}_s \) denote the auxiliary label derived for a self-supervised task. The training set is represented as \( \{(x_i, y_{m,i})\}_{i=1}^{n_s} \). At test time, both covariate distribution $ p(X')$ and label distribution $p(Y_m')$ $ p(Y_s')$ may change, which leads to domain shifts.
Test-Time Training \citep{sun2020test} addresses these shifts by leveraging the same SSL task during both the training and testing phases to align features between the training domain and individual test instances. The framework consists of a shared feature encoder $f_e(\cdot; \theta_e)$, a primary classification head $h_c(\cdot; \theta_c)$ and an SSL head $h_s(\cdot; \theta_s)$.
During training, TTT jointly optimizes both the main classification loss $\mathcal{L}_{\text{main}}$ and the auxiliary SSL loss $\mathcal{L}_{\text{ssl}}$ as
\begin{equation}
\scalebox{0.9}{$
\theta_e^*, \theta_c^*, \theta_s^* = \arg\min\limits_{\theta_e, \theta_c, \theta_s} \sum_{i=1}^{n_s} \mathcal{L}_{\text{main}}(x_i, y_i; \theta_e, \theta_c) + \mathcal{L}_{\text{ssl}}(x_i; \theta_e, \theta_s)$}.
\end{equation}
At inference time, rather than relying on static model parameters, TTT dynamically adapts the encoder for each test instance $x'$ by optimizing the SSL objective with
\iffalse
where $\mathcal{L}_{\text{main}}$ denotes the main objective (e.g., cross-entropy for classification tasks), and $\mathcal{L}_{\text{ssl}}$ denotes the auxiliary SSL objective (e.g., rotation prediction). For each test instance $x'$, TTT first adjusts the feature encoder $f_e(\cdot; \theta_e)$ by optimizing the SSL objective:
\fi
\begin{equation}
\theta_e(x') = \arg\min_{\theta_e} \mathcal{L}_{\text{ssl}}(x'; \theta_s^*, \theta_e).
\end{equation}
The adapted encoder is then used to obtain the final prediction with
\begin{equation}
\hat{y} = h_c(f_e(x'; \theta_e(x')); \theta_c^*).
\end{equation}
\subsection{Theoretical Insights}
Prior work \citep{liu2021ttt++} derives accuracy bounds under assumptions of distributional alignment and task consistency. We provide an independent perspective based on information theory that examines how the auxiliary task informs the main task through shared representations. Let \( Z \) and \( Z' \) represent the feature representations from the feature encoder for the training and test domains, respectively. We define \(\pi_m \) is the main task classifier predicting \( Y_m \), and \( \pi_s \) is the SSL classifier predicting \( Y_s \). The probability \( P(\pi_m(Z') = Y_m') \) quantifies the likelihood that the main task classifier \( \pi_m \) correctly predicts the main task label \( Y_m' \) at test time.
In test-time training, we assume the Markov chain \( Y_s' \to Z' \to Y_m' \) holds, which captures the dependency structure where the auxiliary task labels \( Y_s' \) influence the main task labels \( Y_m' \) only through the shared representation \( Z' \). Under this assumption, we establish the relationship between the mutual information of the auxiliary and main tasks (please refer to Appendix \ref{sm} for derivation details).
\begin{lemma}
In test-time training, where only the shared representation layers are updated using \( Y_s' \), the following inequality holds:
\begin{equation}
I(Z'; Y_m') \geq I(Y_s'; Y_m').
\end{equation}
\end{lemma}
Building on this, we derive information-theoretic bounds on the main task prediction error with binary case in the ideal scenario (please refer to Appendix \ref{smm} for derivation details and multi-class case).
\begin{theorem}
Let $\eta(z') = P\{Y_m' = 1 \mid Z' = z'\} > 0.5$ as positive, the minimum classification error is $p(e) = \int_{Z'} \min\{\eta(z'), 1-\eta(z') \} dp(z')$ and $H_\text{err}(\eta)= -\eta \cdot \log \eta - (1-\eta) \cdot \log (1-\eta)$, the prediction error is bounded by
\begin{equation}
H_\text{err}^{-1}(H(Y_m' \mid Y_s')) \leq p(e)\leq \frac{1}{2}H(Y_m' \mid Y_s').
\end{equation}
\end{theorem}
The lower and upper bounds on the prediction error of the main task highlights the relationship between the main task and the SSL task in performance after adaptation. The upper bound shows that error is limited by the conditional uncertainty $H(Y_m' \mid Y_s')$ while the lower bound demonstrates that lower \( H(Y_m' \mid Y_s') \) improves worst-case guarantees. Additionally, under domain shift $w(y_s) = \frac{P(Y_s' = y_s)}{P(Y_s = y_s)}$, $H(Y_m' \mid Y_s') = \sum_{y_s} w(y_s) P(Y_s = y_s) H(Y_m' \mid Y_s' = y_s)$, Theorem 1 shows overfitting to the test-domain auxiliary task distribution \( P(Y_s') \) can lead to overweighting regions with high uncertainty \( H(Y_m' \mid Y_s') \). Enforcing
\( P(Y_s') = P(Y_s) \) without accounting for test-specific shifts can further harm model generalization. These findings emphasize the necessity of designing a framework that ensures strong alignment between the main and auxiliary tasks while remaining robust to domain shifts for effective test-time adaptation.
\subsection{Adaptive Test-Time Training }
The effectiveness of test-time training depends on SSL alignment with the main task and handling distribution shifts. To address this, we propose Adaptive Test-Time Training (AdaTTT) to enhance TTT with dynamic self-supervised learning and prototype-guided adaptation, which aims to improve generalization under clinical domain shifts in EHR data.
\subsubsection{Dynamic self-supervised learning}
Fixed SSL transformations (e.g., random feature masking) may introduce spurious patterns unrelated to the main task. In EHR data, some features are more predictive, and treating all equally reduces adaptation effectiveness. To mitigate this, we introduce two pretext tasks: Reconstruction and Masked Feature Modeling along with a dynamic feature masking strategy that prioritizes informative features.
\noindent{\textbf{SSL Loss}}.
Given an input vector $\mathbf{x}$ and the the corrupted input $\tilde{\mathbf{x}}$, the reconstruction loss and masked feature modeling loss are defined as
\begin{equation}
\mathcal{L}_{\text{recon}} = \frac{1}{d} \sum_{j=1}^d \big(x_j - \hat{x}_j\big)^2,
\end{equation}
where $d$ is the total number of features, and $\hat{x}_j$ represents the reconstructed value of feature $x_j$ predicted by the model.
\begin{equation}
\mathcal{L}_{\text{mfm}} = \frac{1}{|M|} \sum_{j \in M} \big(x_j - \hat{x}_j\big)^2,
\end{equation}
where $M$ is the set of indices of masked features (\( m_j = 1 \)), and $|M|$ denotes the number of masked features.
The overall self-supervised learning loss combines the reconstruction loss and the masked feature modeling loss with
\begin{equation}
\mathcal{L}_{\text{ssl}} = \lambda_{\text{recon}} \cdot \mathcal{L}_{\text{recon}} + \mathcal{L}_{\text{mfm}},
\end{equation}
\noindent{\textbf{Dynamic Feature Masking}}. Instead of fixed random masking, we introduce an adaptive masking strategy that assigns higher masking probabilities to more informative features. Given an input vector $\mathbf{x}$, the corrupted input $\tilde{\mathbf{x}}$ is generated as follows:
\begin{equation}
\tilde{x}_j = m_j \cdot f(\mathbf{x}, j) + (1 - m_j) \cdot x_j,
\end{equation}
where $\mathbf{m} \in [0, 1]^d$ is the mask vector with elements $m_j$. $f(\mathbf{x}, j) \sim P(x_j) $ is a replacement for feature $x_j$ where $P(x_j)$ is the empirical distribution of training or testing data.
Masking probabilities are dynamically updated based on global feature relevance scores derived from the main task:
\iffalse
\begin{equation}
\label{10}
I_j = \frac{1}{n_s |\mathcal{Y}_m|} \sum_{n=1}^{n_s} \sum_{c=1}^{|\mathcal{Y}_m|} \left| \frac{\partial Y_{m,c}^{(n)}}{\partial x_j^{(n)}} \cdot x_j^{(n)} \right|,
\end{equation}
\fi
\begin{equation}
\label{10}
I_j = \frac{1}{n_s} \sum_{n=1}^{n_s}\left| \frac{\partial Y_{m}^{(n)}}{\partial x_j^{(n)}} \cdot x_j^{(n)} \right|,
\end{equation}
\begin{equation}
\label{11}
p_{m,j} = \frac{I_j - \min_k I_k}{\max_k I_k - \min_k I_k}.
\end{equation}
The masking phrases in the training section are described as follows: In the training phase, we employ a two-stage masking strategy. During the warmup phase (Epochs 1–N), dynamic masking is applied using a fixed prior mask probability to encourage broad feature exploration (the prior is derived from a pretrained respiratory failure prediction model). In the subsequent adaptive masking phase (Epochs N+1–End) feature relevance scores are updated at each epoch based on the model’s main task predictions from the previous epoch. These relevance scores are then used to refine the masking probabilities to enable the model to focus on more informative features over time.
During test-time training, the model continues refining masking probabilities at each gradient step to ensure SSL tasks remain aligned with the primary task. Prioritizing the most informative features challenges the model to reconstruct or predict essential aspects of the data and then improve generalization under domain shifts.
\subsubsection{Prototype-Guided Adaptation}
Prototype-guided adaptation is particularly important in the context of EHR-based test-time training. Instance-level SSL updates may overfit to noisy or highly site-specific measurements, especially when the target-domain distribution differs substantially from the source. To mitigate this, we introduce a set of learned prototypes that capture population-level structure from the source domain. These prototypes act as stable anchors during adaptation. By leveraging Partial Optimal Transport (POT), the model can perform soft and partial alignment between test-time features and prototype distributions, enabling flexible yet robust adaptation under domain shift while avoiding the brittleness of full-transport matching.
\noindent{\textbf{Training Stage}}. We introduce a prototype learning loss that encourages the shared layer features \( z \) to align with their corresponding prototypes \( \mathbf{P} \). These prototypes are designed to effectively represent the distribution of the feature space \( z \) and ensure that the feature embeddings are compact and structured around these representative points.
The prototype learning loss is defined as
\begin{equation}
\mathcal{L}_{\text{proto}}(z_i; \mathbf{P}) = \|z_i - p_{\mathcal{A}(z_i)}\|_2^2,
\end{equation}
where $\mathbf{P} = \{p_1,p_2,...,p_k\} $ is the set of $k$ prototypes. $p_{\mathcal{A}(z_i)}$ is the prototype assigned to the feature $z_i$ based on the cluster assignment $\mathcal{A}(z_i)$.
To prevent all features are assigned to a single prototype, a regularization term is added to balance the cluster assignments:
\begin{equation}
\mathcal{L}_{\text{reg}}(\mathbf{P}) = \sum_{j=1}^k \Bigg(\frac{1}{n_s} \sum_{i=1}^{n_s} \mathbb{I}(\mathcal{A}(z_i) = j) - \frac{1}{k}\Bigg)^2,
\end{equation}
where $\mathbb{I}(\mathcal{A}(z_i) = j)$ is an indicator function that equals 1 if $z_i$ is assigned to prototype $p_j$.
The final loss function incorporating prototypes is given as
\begin{equation}
\scalebox{0.85}{$
\mathcal{L} = \sum_{i=1}^{n_s} \Big[ \mathcal{L}_{\text{main}}(x_i, y_i) + \mathcal{L}_{\text{ssl}}(x_i) + \lambda_{\text{proto}} \mathcal{L}_{\text{proto}}(z_i)\Big]+ \lambda_{\text{reg}} \mathcal{L}_{\text{reg}}(\mathbf{P}),$}
\end{equation}
where $\mathcal{L}_{\text{main}}$ is the respiratory failure prediction loss. $\mathcal{L}_{\text{ssl}}$ is the self-supervised loss (see Section 3.3.1). $\mathcal{L}_{\text{reg}}$ is the regularization loss for balanced assignment. $\lambda_{\text{proto}}$ and $\lambda_{\text{reg}}$ are hyperparameters controlling the importance of the prototype and regularization terms.
\noindent{\textbf{Test-Time Training Stage}}. Traditional Optimal Transport (OT) assumes a full alignment between source and target distributions, which may be too rigid in the presence of domain shifts. We refine the alignment between the training prototypes \( \mathbf{P} \) and the test-time feature representations \( z' \) by incorporating POT. Instead of constraining the transport plan to only partially align prototypes with the test instance \citep{chapel2020partial}, we augment the set of \( \mathbf{z}' \) by adding \( k-1 \) perturbed duplicates to transform the transport problem into a standard optimal transport setting while enabling partial matches between \( z' \) and the prototypes \( \mathbf{P} \).
\begin{equation}
\mathbf{z}' = \{z', z'_1, \dots, z'_{k-1}\},
\end{equation}
where each duplicate $z'_j$ is sampled as
\begin{equation}
z'_{j,d} \sim z'_d + \mathcal{N}(0, \sigma^2_d),
\end{equation}
\begin{equation}
\sigma_d^2 = \frac{1}{k} \sum_{j=1}^{k} \left( p_{j,d} - \mu_d \right)^2,
\end{equation}
where \( p_{j,d} \) is the value of the \( d \)-th dimension of the \( j \)-th prototype \( \mathbf{p}_j \), \( \mu_d \) is the mean of the \( d \)-th dimension across all prototypes.
The loss function employed during test-time training is defined as
\begin{equation}
\mathcal{L}_{\text{test}} = \mathcal{L}_{\text{ssl}} + \lambda_{\text{ot}} \cdot \sum_{i,j} \gamma_{ij} C_{ij},
\end{equation}
where \(\lambda_{\text{ot}}\) is a hyperparameter balancing the importance of the OT cost in the overall loss. \(\gamma_{ij}\) defines the mass transported from \(\mathbf{z}_i'\) to the \(j\)-th prototype. $C_{ij} = \| \mathbf{z}_i' - \mathbf{p}_j \|_2^2$ represents the squared Euclidean distance between \(\mathbf{z}_i'\) and the prototype \(\mathbf{p}_j\).
\section{Experiments}
\subsection{Experimental Setting}
\noindent{\textbf{Datasets}}. We conduct a retrospective study using de-identified EHR data of all adult patients ($\geq$ 18 years) admitted to the ICU at Site A\footnote{For anonymization purposes, the name of the healthcare institution has been replaced with Site $\times$. } between January 1, 2016, and December 31, 2023. This dataset served as the development and validation cohort. To evaluate the mechanisms of test-time training, we utilize additional datasets, including ICU admissions at Site A between January 1, 2024, and June 30, 2024, Site B between January 1, 2023, and August 31, 2024, as well as the publicly available MIMIC-IV dataset. Institutional Review Board approval was obtained for the use of these datasets. Appendix \ref{data} provides full details on cohort selection and data processing. Our development cohort consists of 24,943 encounters, with 1,308 positive cases (IMV prevalence: 5.2\%). The testing cohorts include Site A (1,835 encounters, 104 positive cases, IMV prevalence: 5.7\%) and Site B (2,564 encounters, 141 positive cases, IMV prevalence: 5.5\%). The original MIMIC-IV dataset contains 35,534 encounters with an IMV prevalence of 15.4\%. For computational efficiency, we randomly downsampled MIMIC-IV to 2,069 encounters (244 positive cases with an IMV prevalence of 11.8\%).
\iffalse
under protocol #800258 (“A Real-Time Multimodal Data Integration Model for Prediction of Respiratory Failure in Patients with COVID-19”).
\fi
\noindent{\textbf{Implementation Details}}.
\label{sid}
Our network architecture follows Vent.io \citep{lam2024development} (Appendix~\ref{archi}). We use Bayesian optimization for source-domain pretraining to tune general network hyperparameters (learning rate, weight regularization, number of hidden layers). The prototype set size is \(k=4\). For dynamic masking, we adopt a two-phase schedule: a warm-up phase with a \emph{fixed prior mask probability} of \(0.5\), followed by an adaptive phase where masking probabilities are updated from feature relevance. During deployment, test-time training (TTT) performs five gradient steps per input (we refer to each step as an “\textbf{iteration}”), and we follow the standard reset protocol \citep{sun2020test} by restoring encoder weights to the pretrained state after each instance. For partial optimal transport, we use Sinkhorn with entropic regularization \(\varepsilon=0.1\), a maximum of 1000 Sinkhorn iterations, and a mini-batch size equal to the prototype count (\(K=k\)). All baselines are re-implemented with the same encoder, data pipeline, and hyperparameter search budget to ensure a fair comparison.
\noindent{\textbf{Evaluation}}. We follow the clinical labeling scheme \citep{lam2024development} to capture the physiological states of respiratory failure, defining score $\geq\!3$ as positive and $<\!3$ as control (Appendix~\ref{data}). Encounters are categorized as True Positive, False Positive, True Negative, or False Negative based on predictions within the specified prediction window (criteria in Appendix~\ref{data}). Performance is reported as AUC with mean~$\pm$~standard error over 20 independent test-time training runs. In addition, following best practices for clinical prediction model assessment \citep{huang2020tutorial}, we report Brier score for all models to evaluate calibration and clinical utility more comprehensively.
\subsection{Comparison with Baselines}
\noindent{\textbf{Baselines}}.
We consider a set of foundational and representative TTT and TTA methods and adapt each to the EHR domain to ensure fair and meaningful comparison in our experimental evaluations (Appendix \ref{ba} provides full details of adapting baselines to EHR setting): \textbf{TEST}, \textbf{TENT} \citep{wang2020tent}, \textbf{TTT} \citep{sun2020test}, \textbf{TTT++} \citep{liu2021ttt++}, \textbf{ClusT3} \citep{hakim2023clust3}, \textbf{NC-TTT} \citep{osowiechi2024nc}, \textbf{T3A} \citep{iwasawa2021test}, \textbf{SAR} \citep{niu2023towards} and \textbf{CoTTA} \citep{wang2022continual}. We also evaluate two ablated versions of our method: \textbf{PriTTT}, which removes the adaptive distribution matching module and relies solely on updates to the mask probabilities for test-time training. \textbf{DynTTT}, which removes the dynamic masking module and focuses only on adaptive distribution matching.
\iffalse
\begin{itemize}
\item \textbf{TEST:} directly evaluate the pretrained model on the test data without any adaptation;
\item \textbf{TENT} \citep{wang2020tent}: minimize the Shannon entropy of model predictions during test time by optimizing model parameters without requiring annotations;
\item \textbf{TTT} \citep{sun2020test}: jointly train the network on the main task and our defined SSL task in the source domain. During testing, TTT adapts the model to the target domain by continuing training on the SSL task;
\item \textbf{TTT++} \citep{liu2021ttt++}: building on TTT, TTT++ further aligns the first-order and second-order statistics of the source and target distributions;
\item \textbf{ClusT3} \citep{hakim2023clust3}: based on the Y-shaped architecture, we incorporates a projector after the last layer of the shared encoders to enhance test-time training by clustering representations.
\item \textbf{NC-TTT} \citep{osowiechi2024nc}: a noise-contrastive estimation-based auxiliary task for test-time training. The model is adapted at test time to increase in-distribution likelihood without using source samples or labels.
\end{itemize}
\fi
\iffalse
\noindent{\textbf{Results \& Analysis}}. Table \ref{table:noise_auc} presents the AUC evaluation of different baselines across three distinct sites. Baseline methods like TENT and ClusT3 exhibit performance drops compared to the TEST method that evaluates the pretrained model without adaptation. Our proposed AdaTTT framework consistently outperforms the baselines across all testing sites and demonstrates its effectiveness in adapting to distributional shifts during test time without sacrificing predictive accuracy. Figure \ref{si} illustrates the evolution of risk scores across different time points during test-time training for a patient from Site A. At earlier time points, the model adapts towards lower predicted risk. As the time approaches the intubation event (within 24 hours), the risk score shows a noticeable increase, which reflects the model’s ability to dynamically adjust predictions in response to changing patient conditions.
\begin{figure*}
\centering
\includegraphics[width=0.4\textwidth]{Figs/score_env.png} % Adjust width accordingly
\caption{Evolution of risk scores at different time points during test-time training for a patient from Site A who ultimately requires intubation.}
\label{si}
\end{figure*}
\fi
\iffalse
\noindent{\textbf{Results \& Analysis}}. Table~\ref{t1} presents the AUC evaluation of different baselines across three distinct sites. Baseline methods like TENT and ClusT3 exhibit performance drops compared to the TEST method that evaluates the pretrained model without adaptation. Our proposed AdaTTT framework consistently outperforms the baselines across all testing sites and demonstrates its effectiveness in adapting to distributional shifts during test time without sacrificing predictive accuracy. Figure~\ref{si} illustrates the evolution of risk scores across time points for a patient from Site A. Early predictions trend lower, but risk escalates approaching the intubation event (within 24 hours), demonstrating dynamic response to clinical deterioration.
\fi
\iffalse
\begin{table}
\centering
\renewcommand{\arraystretch}{2.0}
\caption{AUC (\%) evaluation across various testing sites.}
\label{table:noise_auc}
\resizebox{0.99\textwidth}{!}{%
\begin{tabular}{l|c|c|c|c|c|c|c|c|c}
\toprule
\textbf{} & TEST & TENT& TTT & TTT++ & ClusT3 & NC-TTT & \textbf{PriTTT(Ours)} &\textbf{DynTTT(Ours)} & \textbf{AdaTTT(Ours)}\\
\midrule
\textbf{Site A} & $84.01$ & $82.19\pm0.11$ & $82.55\pm0.09$ & $82.50\pm0.06$& $82.36\pm0.08$ & $82.32\pm0.06$ & $84.61\pm 0.03$&$84.54\pm0.10$ & \bm{$85.02\pm0.05$} \\
\textbf{Site B} & $83.75$ & $83.06\pm0.08$ & $82.81\pm0.05$ & $82.85\pm0.10$ & $81.99\pm0.12$ & $83.62\pm 0.10$& $83.98\pm0.06$& $83.84\pm 0.12$ & \bm{$84.10\pm 0.05$} \\
\textbf{MIMIC-IV} & $75.28$ & $74.34\pm 0.05$ &$76.45\pm0.07$ & $76.24\pm 0.08$& $74.41\pm0.11$& $75.27\pm 0.08$& $76.84\pm0.03$&$76.79\pm0.05$ & \bm{$77.17\pm0.08$} \\
\bottomrule
\end{tabular}%
}
\end{table}
\fi
% Preamble (if needed):
% \usepackage{booktabs}
% \usepackage{graphicx} % for \rotatebox
% \usepackage{bm}
% Preamble (if needed):
% \usepackage{booktabs}
% \usepackage{graphicx} % for \rotatebox
% \usepackage{bm}
\newcommand{\rot}[1]{\rotatebox[origin=c]{60}{\smash{#1}}}
\begin{table*}[t]
\centering
\renewcommand{\arraystretch}{1.45} % increase row spacing
\setlength{\tabcolsep}{3.5pt}
\small
\caption{AUC (\%) across testing sites (\(\uparrow\) higher is better).}
\label{tab:auc_sites}
\resizebox{\linewidth}{!}{%
\begin{tabular}{l|ccccccccccccc}
\toprule
\textbf{Dataset} &
\rot{TEST} & \rot{CoTTA} & \rot{T3A} & \rot{SAR} &
\rot{TENT} & \rot{TTT} & \rot{TTT++} & \rot{ClusT3} & \rot{NC-TTT} &
\rot{\textbf{PriTTT (Ours)}} & \rot{\textbf{DynTTT (Ours)}} & \rot{\textbf{AdaTTT (Ours)}} \\
\midrule
\textbf{Site A} & 84.01 & $83.12\!\pm\!0.02$ & $82.50\!\pm\!0.04$ & $84.30\!\pm\!0.04$ &
$82.19\!\pm\!0.11$ & $82.55\!\pm\!0.09$ & $82.50\!\pm\!0.06$ & $82.36\!\pm\!0.08$ & $82.32\!\pm\!0.06$ &
$84.61\!\pm\!0.03$ & $84.54\!\pm\!0.10$ & \bm{$85.02\!\pm\!0.05$} \\
\textbf{Site B} & 83.75 & $83.81\!\pm\!0.04$ & $83.10\!\pm\!0.05$ & $83.20\!\pm\!0.10$ &
$83.06\!\pm\!0.08$ & $82.81\!\pm\!0.05$ & $82.85\!\pm\!0.10$ & $81.99\!\pm\!0.12$ & $83.62\!\pm\!0.10$ &
$83.98\!\pm\!0.06$ & $83.84\!\pm\!0.12$ & \bm{$84.10\!\pm\!0.05$} \\
\textbf{MIMIC-IV} & 75.28 & $76.60\!\pm\!0.05$ & $76.10\!\pm\!0.03$ & $75.72\!\pm\!0.04$ &
$74.34\!\pm\!0.05$ & $76.45\!\pm\!0.07$ & $76.24\!\pm\!0.08$ & $74.41\!\pm\!0.11$ & $75.27\!\pm\!0.08$ &
$76.84\!\pm\!0.03$ & $76.79\!\pm\!0.05$ & \bm{$77.17\!\pm\!0.08$} \\
\bottomrule
\end{tabular}%
}
\label{t1}
\end{table*}
% Preamble (if needed):
% \usepackage{booktabs}
% \usepackage{graphicx} % for \rotatebox
% \usepackage{bm}
%\newcommand{\rot}[1]{\rotatebox[origin=c]{60}{\smash{#1}}}
\begin{table*}[t]
\centering
\renewcommand{\arraystretch}{1.45} % increase row spacing
\setlength{\tabcolsep}{3.5pt} % tighten column spacing
\small
\caption{Brier score across testing sites ($\downarrow$ lower is better).}
\label{tab:brier_sites}
\resizebox{\linewidth}{!}{%
\begin{tabular}{l|ccccccccccccc}
\toprule
\textbf{Dataset} &
\rot{TEST} & \rot{CoTTA} & \rot{T3A} & \rot{SAR} &
\rot{TENT} & \rot{TTT} & \rot{TTT++} & \rot{ClusT3} & \rot{NC-TTT} &
\rot{\textbf{PriTTT (Ours)}} & \rot{\textbf{DynTTT (Ours)}} & \rot{\textbf{AdaTTT (Ours)}} \\
\midrule
\textbf{Site A} & 0.089 & $0.092\!\pm\!0.01$ & $0.093\!\pm\!0.01$ & $0.089\!\pm\!0.02$ &
$0.090\!\pm\!0.04$ & $0.093\!\pm\!0.03$ & $0.094\!\pm\!0.01$ & $0.090\!\pm\!0.03$ & $0.090\!\pm\!0.01$ &
\bm{$0.089\!\pm\!0.02$} & \bm{$0.089\!\pm\!0.01$} & \bm{$0.086\!\pm\!0.02$} \\
\textbf{Site B} & 0.089 & $0.091\!\pm\!0.01$ & $0.091\!\pm\!0.02$ & $0.091\!\pm\!0.02$ &
$0.091\!\pm\!0.03$ & $0.094\!\pm\!0.02$ & $0.095\!\pm\!0.04$ & $0.092\!\pm\!0.04$ & $0.091\!\pm\!0.03$ &
\bm{$0.090\!\pm\!0.03$} & \bm{$0.089\!\pm\!0.02$} & \bm{$0.085\!\pm\!0.02$} \\
\textbf{MIMIC-IV} & 0.111 & $0.110\!\pm\!0.02$ & $0.110\!\pm\!0.03$ & $0.111\!\pm\!0.03$ &
$0.114\!\pm\!0.04$ & $0.110\!\pm\!0.04$ & $0.111\!\pm\!0.05$ & $0.113\!\pm\!0.04$ & $0.113\!\pm\!0.05$ &
\bm{$0.110\!\pm\!0.05$} & \bm{$0.112\!\pm\!0.04$} & \bm{$0.106\!\pm\!0.04$} \\
\bottomrule
\end{tabular}%
}
\label{t2}
\end{table*}
\iffalse
AdaTTT improves upon TTT and TTT++ by addressing two key aspects of test-time training. First, it enhances the alignment between the SSL task and the main prediction task. Figure \ref{fi} presents the feature importance evolves from the priors obtained by the vanilla pretrained model to the initial and final epochs in the training section\footnote{ICULOS refers to ICU stay till now. Resp: respiratory rate. HR: heart rate. LD: lymphocyte differential count. BUN: blood urea nitrogen. O2Sat: oxygen saturation. WBC: white blood count.}. We can observe that in the initial epoch the model has not yet fully learned to identify task-relevant features. This aligns with our warm-up phase, where dynamic masking is initialized using prior mask probabilities. As training progresses into the adaptive masking phase, the feature importance distribution becomes more stable and better aligned with the learned feature relevance. By iteratively refining the feature selection process, AdaTTT enables the model to adaptively emphasize the most informative features.
During test-time training, the model further adjusts feature importance at each gradient update. Figure \ref{fii} presents an example of feature importance evolution for a specific time point of a patient who requires intubation within 24 hours. While some features remain stable in importance, others adapt dynamically throughout the iterations. For example, the Respiratory Rate feature exhibits a continuous increase in importance, which aligns with both the model faithfulness and clinical plausibility.
Second, unlike TTT, which may overfit to the test samples, or TTT++ that may not generalize well due to fixed distribution alignment, our approach dynamically adapts to test-time variations through POT. Our unconstrained formulation of POT allows for a flexible alignment between the test-time features $z'$ and the learned prototypes $\mathbf{P}$, which prevents excessive adaptation to noisy test samples while ensuring that meaningful patterns are captured. To further illustrate this, an example of the POT alignment between the test-time features $z'$ and the learned prototypes $\mathbf{P}$ is provided in Appendix~\ref{pot}.
\fi
\iffalse
\begin{figure}
\centering
\includegraphics[width=0.45\textwidth,height=0.25\textwidth]{Figs/test_env.png} % Adjust width accordingly
\caption{An example of feature importance evolution during test-time training. The heatmap shows the changes in feature importance across different iterations.}
\label{fii}
\end{figure}
\fi
\iffalse
\begin{table}[ht]
\centering
\renewcommand{\arraystretch}{1.2}
\caption{AUC (\%) evaluation across various sites.}
\label{table:noise_auc}
\resizebox{0.45\textwidth}{!}{%
\begin{tabular}{l|c|c|c}
\toprule
\textbf{Method} & \textbf{Site A} & \textbf{Site B} & \textbf{MIMIC-IV} \\
\midrule TEST & 0.840 & 0.838 & 0.753 \\
%TTT\_base & 84.37 & 84.11 & 77.49 \\
%ClusT3\_base & 82.62 & 82.86 & 75.03 \\
%AdaTTT\_base (\textbf{Ours}) & 84.59 & 84.00 & 76.84 \\
\midrule
TNET & 0.822 & 0.831 & 0.743 \\
TTT & 0.826 & 0.828 & 0.765 \\
TTT++ & 0.825 & 0.829 & 0.762 \\
ClusT3 & 0.824 & 0.820 & 0.744 \\
\midrule
PriTTT (\textbf{Ours}) &0.846 & 0.840 & 0.768 \\
DynTTT (\textbf{Ours}) & 0.847 & 0.838 & 0.769 \\
AdaTTT (\textbf{Ours}) & & & \\
\bottomrule
\end{tabular}%
}
\end{table}
\fi
\noindent\textbf{Results \& Analysis.} Tables~\ref{t1} (AUC) and \ref{t2} (Brier) report discrimination and calibration across three cohorts. AdaTTT obtains the top AUC on every site. Among non-TTT/TTA methods, SAR is closest on Site~A (84.30$\pm$0.04), and CoTTA is competitive on MIMIC-IV (76.60$\pm$0.05), yet both remain below AdaTTT. Regarding calibration, AdaTTT improves or matches calibration on Site~A/B (Brier $\mathbf{0.086}$/$\mathbf{0.085}$ vs.\ TEST 0.089/0.089) and achieves 0.106$\pm$0.04 on MIMIC-IV.
\iffalse
\begin{wrapfigure}{r}{0.6\textwidth}
\centering
\vspace{-2em}
\includegraphics[width=\linewidth]{Figs/score_env.png}
\vspace{-1.5em}
\caption{Risk score evolution during test-time training for a patient from Site A. Risk increases as intubation nears, which reflects model adaptation.}
\label{si}
\end{wrapfigure}
\fi
\begin{figure}
\centering
\includegraphics[width=0.7\linewidth,height=0.25\textwidth]{Figs/score_env.png} % Adjust width accordingly
\caption{Risk score evolution during test-time training for a patient from Site A. Risk increases as intubation nears, which reflects model adaptation.}
\label{fii}
\end{figure}
The behavior of the baselines is consistent with their assumptions and with the numbers in Tables~\ref{t1}–\ref{t2}. TENT’s entropy minimization encourages confidence that can be misplaced on out-of-distribution patients. Standard TTT and TTT++ attach feature-agnostic SSL objectives and, in the latter, enforce relatively rigid source–target alignment; both are brittle when clinical feature importance is highly unequal across variables and when hospital shifts are complex. ClusT3’s discrete codes and domain-consistency assumption struggle with continuous physiology and site heterogeneity; NC-TTT’s contrastive likelihoods depend on well-specified “noise,” which is hard to define for heterogeneous EHR features. Classifier-only adjustment in T3A helps when classes are tightly clustered (MIMIC-IV 76.10$\pm$0.03) but cannot correct representation shift (Site~A 82.50$\pm$0.04). SAR stabilizes BN-based updates and fares well on Site~A (84.30$\pm$0.04) but remains sensitive to batch-statistics quality, yielding inconsistent improvements on Site~B/MIMIC-IV.
We also examine how AdaTTT achieves its gains. Figure~\ref{si} illustrates the evolution of risk scores across time points for a patient from Site A. Early predictions trend lower, but risk escalates approaching the intubation event (within 24 hours), demonstrating dynamic response to clinical deterioration. First, aligning the auxiliary SSL objective with the clinical endpoint is critical: with dynamic, feature-aware masking, the SSL branch prioritizes clinically salient variables and down-weights weak signals. Figure~\ref{fi} traces feature importance from pretrained priors through early to late training epochs\footnote{ICULOS: ICU length-of-stay to time $t$; Resp: respiratory rate; HR: heart rate; LD: lymphocyte differential; BUN: blood urea nitrogen; O$_2$Sat: oxygen saturation; WBC: white blood count.}. During warm-up, masking follows the priors; as training proceeds, the distribution stabilizes and aligns with learned relevance, indicating tighter coupling between SSL and IMV prediction than in feature-agnostic TTT/TTT++. Second, at deployment the model refines importance \emph{per input}: Figure~\ref{fii} shows that some features remain stable while others (e.g., respiratory rate) gain weight across iterations, consistent with model faithfulness and clinical plausibility. In parallel, prototype-guided partial optimal transport flexibly matches test-time representations to learned prototypes, limiting overfitting to idiosyncratic or noisy samples while preserving clinically meaningful structure (see Appendix~\ref{pot} for an illustrative alignment). Ablations support this interpretation: PriTTT (feature-aware masking only) and DynTTT (distribution matching only) each improve over TTT/TTT++, but their combination yields the most consistent AUC and Brier gains across sites.
\begin{minipage}[t]{0.48\textwidth}
\centering
\includegraphics[width=0.9\linewidth, height=0.5\textwidth]{Figs/rs_training.png}
\captionof{figure}{Feature importance evolution during training. The heatmap shows the changes in feature importance in the initial epochs and final epochs.}
\label{fi}
\end{minipage}
\hfill
\begin{minipage}[t]{0.48\textwidth}
\centering
\includegraphics[width=\linewidth, height=0.5\textwidth]{Figs/test_env.png}
\captionof{figure}{An example of feature importance evolution during test-time training. The heatmap shows the changes in feature importance across different iterations.}
\label{fii}
\end{minipage}
\noindent{\textbf{Computational Cost}}. Our proposed AdaTTT framework incorporates both feature importance updates and optimal transport computation during the test-time training phase. These additional operations inevitably increase the computational cost compared to standard TTT frameworks. We use the Sinkhorn algorithm \citep{sinkhorn1967diagonal} that leverages entropy regularization to ensure scalable computations. We evaluate the execution time of a single gradient update during test-time training. The average execution time is 0.29s, and did not change much, remaining at 0.26s when increasing the prototype size from 4 to 16.
\begin{figure*}[h]
\centering
\begin{minipage}{\textwidth} % Ensure all figures stay in one line
\centering
\begin{subfigure}{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{Figs/ucsdP_aba.png}
\caption{Site A.}
\label{fig:sst2}
\end{subfigure}
\hspace{2mm} % Add spacing to ensure alignment
\begin{subfigure}{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{Figs/uci_aba.png}
\caption{Site B.}
\label{fig:snli}
\end{subfigure}
\hspace{2mm} % Add spacing to ensure alignment
\begin{subfigure}{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{Figs/mimic_aba.png}
\caption{MIMIC-IV.}
\label{fig:yelp}
\end{subfigure}
\end{minipage}
\caption{Evaluation of the number of gradient updates for test-time training on different test cohorts.}
\label{iter}
\end{figure*}
\subsection{Sensitivity Analysis}
In this section, we examine the effect of the number of test-time training iterations and compare reset versus sequential update mechanisms. Additional ablations are provided in Appendix~\ref{sen} which investigate (1) the impact of prototype size, (2) the role of prototype learning, and (3) the effect of dynamic masking and the two SSL objectives.
\noindent{\textbf{Comparison of the number of iterations}}.
Figure \ref{iter} presents the impact of the number of test-time training iterations on model performance across three different test cohorts. We evaluate the AUC scores as the number of gradient updates increases from 1 to 5 iterations. Across all sites, AdaTTT exhibits a stable and consistent improvement in AUC with more iterations. In contrast, PriTTT and DynTTT show less consistent trends, with fluctuations in performance, particularly in Site B. PriTTT updates feature importance dynamically without a stable reference and is sensitive to initial feature importance variations. Meanwhile, DynTTT lacks feature selection control, which can lead to suboptimal emphasis on features between the main and SSL tasks.
\noindent\textbf{Reset versus Sequential Update Mechanism.}
Compared with the reset strategy, the sequential update mechanism applies one gradient step at each new data point and retains the updated parameters across subsequent data points. Figure~\ref{sd} presents the cumulative AUC trend across different sites under the sequential update mechanism. In Site B and MIMIC-IV, performance initially improves as the model adapts to recent distributional shifts (e.g., achieving 80.27\% AUC on MIMIC-IV), demonstrating the benefits of short-term adaptation. However, as updates continue across a larger patient population, the model's performance gradually deteriorates. This degradation is likely due to accumulated adaptation noise, which shifts the model’s focus away from its originally learned feature structure.
\iffalse
\noindent{\textbf{Reset versus sequential update mechanism}}.
For sequential update mechanism, we take only one gradient step given each new time point and analyze how model adaptation unfolds across an increasing number of patients. Figure \ref{sd} illustrates the cumulative AUC trend across different sites. As observed in Site B and MIMIC-IV, updating within a certain range of patients initially improves model performance (e.g. AUC 80.27\% in MIMIC-IV), which suggests that short-term adaptation helps refine predictions. However, as updates continue across a larger patient population, the model begins to degrade due to cumulative adaptation noise, which causes to a shift in feature importance away from its original learned structure.
\fi
\begin{figure*}
\centering
\begin{minipage}{\textwidth} % Ensure all figures stay in one line
\centering
\begin{subfigure}{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{Figs/ucsdP_trend.png}
\caption{Site A.}
\label{fig:sst2}
\end{subfigure}
\hspace{2mm} % Add spacing to ensure alignment
\begin{subfigure}{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{Figs/uci_trend.png}
\caption{Site B.}
\label{fig:snli}
\end{subfigure}
\hspace{2mm} % Add spacing to ensure alignment
\begin{subfigure}{0.31\textwidth}
\centering
\includegraphics[width=\textwidth]{Figs/mimic_trend.png}
\caption{MIMIC-IV.}
\label{fig:yelp}
\end{subfigure}
\end{minipage}
\caption{Cumulative AUC trend over an increasing number of patients.}
\label{sd}
\end{figure*}
\section{Conclusion}
In this study, we introduce AdaTTT, a lightweight and adaptive test-time training framework for predicting the need for invasive mechanical ventilation (IMV) 24 hours in advance across multi-center ICU cohorts. AdaTTT integrates feature-aware self-supervision with partial OT–based distribution alignment to address domain shift that naturally arises in real-world EHR time-series data. Across all external cohorts, the method yields **consistent improvements** in discrimination and calibration over strong TTT baselines, demonstrating practical value for real-time clinical risk monitoring.
To better contextualize these gains, we have added (1) an analysis of calibration and clinical relevance, and (2) additional experiments on secondary outcomes including acute kidney injury (AKI) and in-hospital mortality, showing similar trends.
We also explicitly acknowledge limitations: the method may be sensitive to **extreme conditional shifts** in \(p(Y \mid X)\) (e.g., changes in clinical practice patterns), and prototype-based adaptation adds moderate computational overhead at test time. Addressing these challenges is an important direction for future work.
\section*{Ethics Statement}
\paragraph{Compliance and oversight.}
This study analyzes de-identified electronic health records (EHR) from multiple partner institutions (anonymized as Sites A/B) under appropriate institutional review and data-governance oversight, with a waiver of consent where applicable due to de-identification and minimal risk. All activities complied with relevant privacy regulations (e.g., HIPAA) and local security policies. No attempt was made to re-identify individuals, and all reported results are aggregate.
\paragraph{Intended use and clinical safety.}
The model is a research prototype for risk stratification and is \emph{not} a stand-alone medical device. It should only be used with clinician oversight. Any real-world deployment would require prospective evaluation, safety monitoring, and regulatory review. We report discrimination and calibration (AUC, Brier score) to support assessment of clinical utility and risk.
\paragraph{Fairness and shift.}
We evaluate across distinct clinical sites to assess robustness under distribution shift, and we report calibration as recommended for clinical prediction models. Despite these efforts, residual bias and under-representation are possible; models trained in one setting may underperform elsewhere. We caution against out-of-scope use.
\paragraph{Source-free adaptation safeguards.}
Test-time adaptation updates only the encoder on a per-encounter basis and then resets weights before the next patient/time point, preventing cross-patient carryover. No patient-level exemplars or gradients are stored; adaptation logs contain no protected health information.
\paragraph{Transparency and conflicts.}
We will release code and configuration sufficient for reproduction (subject to data-use constraints). Funding and potential conflicts will be disclosed in the camera-ready version. All authors adhere to the ICLR Code of Ethics.
\section*{Reproducibility Statement}
We describe the cohort construction, preprocessing, and labeling scheme in Sec.~\ref{data}; the architecture and training protocol in Sec.~\ref{archi} and Sec.~\ref{sid}; and sensitivity analyses/ablations in Sec.~\ref{sen}. We report full metric definitions (AUC, Brier) and evaluation procedures. We will release anonymized code (data loaders, feature engineering, training/evaluation scripts, and plotting utilities) to allow end-to-end replication with public data, and provide configuration files to reproduce all tables/figures from logs.
\iffalse
\section{Limitation}
While AdaTTT demonstrates strong generalization across three ICU cohorts, several limitations remain. First, the current framework focuses solely on structured EHR data and does not incorporate unstructured modalities such as clinical notes or imaging. Second, the adaptation process is applied uniformly across all test instances, which may result in unnecessary updates for low-uncertainty cases. In future work, we plan to develop a trigger-based test-time training mechanism to selectively adapt only when needed, thereby improving efficiency while maintaining model stability.
\section{Broader impacts}
This work contributes to improving early identification of patients requiring invasive mechanical ventilation, potentially enabling more timely interventions and resource allocation in critical care. By facilitating real-time adaptation to site-specific shifts, AdaTTT can enhance the robustness of clinical decision support systems. However, as with all predictive models in healthcare, care must be taken to ensure equitable deployment and to monitor for unintended biases across patient subgroups. Our findings encourage further research into interpretable and adaptive clinical AI systems that maintain safety and fairness across diverse care settings.
\fi
\nocite{langley00}
\bibliographystyle{unsrtnat}
\bibliography{reference}
\medskip
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\appendix
\section{Theoretic Analysis}
\subsection{Proof of $I(Z'; Y_m') \geq I(Y_s'; Y_m')$}
\label{sm}
Using the chain rule of mutual information, we can expand the mutual information between \( Z' \) and the joint variables \( (Y_m', Y_s') \) as
\begin{equation}
I(Z'; Y_m', Y_s') = I(Z'; Y_s') + I(Z'; Y_m' \mid Y_s').
\end{equation}
Using the chain rule of mutual information again, we have
\begin{equation}
\begin{aligned}
I(Z'; Y_m') & = I(Z'; Y_m', Y_s') - I(Z'; Y_s' \mid Y_m')\\
& =\big(I(Z'; Y_s') + I(Z'; Y_m' \mid Y_s')\big) - I(Z'; Y_s' \mid Y_m').
\end{aligned}
\end{equation}
In the test-time training framework, \( Z' \) is optimized to retain information from \( Y_s' \) that is predictive of \( Y_m' \). We model this by assuming
\[
I(Z'; Y_m' \mid Y_s') - I(Z'; Y_s' \mid Y_m') \geq 0.
\]
\iffalse
In test-time training, we assume the Markov chain \( Y_s' \to Z' \to Y_m' \) holds, then \( I(Z'; Y_s' \mid Y_m') = 0 \), the equation simplifies to
\begin{equation}
I(Z'; Y_m') = I(Z'; Y_s') + I(Z'; Y_m' \mid Y_s').
\end{equation}
\fi
Applying the Data Processing Inequality (DPI) to our Markov chain, we obtain
\begin{equation}
I(Y_s'; Y_m') \leq I(Y_s'; Z').
\end{equation}
From the mutual information decomposition derived earlier, we know that
\begin{equation}
I(Z'; Y_m') \geq I(Z'; Y_s') + I(Z'; Y_m' \mid Y_s').
\end{equation}
Since mutual information is always non-negative,
\begin{equation}
I(Z'; Y_m' \mid Y_s') \geq 0.
\end{equation}
Therefore,
\begin{equation}
I(Z'; Y_m') \geq I(Z'; Y_s').
\end{equation}
Combining this with the DPI result
\[
I(Y_s'; Y_m') \leq I(Y_s'; Z') = I(Z'; Y_s'),
\]
we obtain the final inequality as
\begin{equation}
I(Z'; Y_m') \geq I(Y_s'; Y_m').
\end{equation}
\subsection{Proof of prediction error bounds of the main task}
\label{smm}
\subsubsection{Binary case}
For the two-class problem, the classifier predicts the input $z'$ with
the posterior $\eta(z') = P\{Y_m' = 1 \mid Z' = z'\} > 0.5$ as positive, the minimum prediction error is
\begin{equation}
p(e) = \int_{Z'} \min\{\eta(z'), 1-\eta(z')\} dp(z').
\end{equation}
Then Shannon entropy for a binary random variable with the distribution $(\eta, 1-\eta)$ is defined as
\begin{equation}
H(\eta) = -\eta \cdot \log \eta - (1-\eta)\cdot \log (1-\eta), \quad \eta \in [0, 1].
\end{equation}
The expectation of the above function with respect to $z' \sim Z'$ is
\begin{equation}
H(Y_m' \mid Z') = \mathbb{E}_{z' \sim Z'}[H(\eta (z'))] = \int_{Z'} H(\eta(z')) d p(z').
\end{equation}
Based on Fano’s inequality, we have $H_\text{err}(p(e)) \geq H(Y_m' \mid Z')$. As $p(e) \leq 0.5$ and the function $H_\text{err}(p(e))$ is monotonically increasing for $0 \leq \eta \leq 0.5$, we have
\begin{equation}
p(e)\geq H_\text{err}^{-1}(H(Y_m' \mid Z')).
\end{equation}
Given that $H(Y_m' \mid Z') = H(Y_m') - I( Z';Y_m')$ and $H(Y_m' \mid Y_s') = H(Y_m') - I(Y_s';Y_m')$, in the ideal test-time training under our Markov chain assumption \( Y_s' \to Z' \to Y_m' \), $Z'$ is as informative about $Y_m'$ as $Y_s'$ is, we have
\begin{equation}
H(Y_m' \mid Z') = H(Y_m' \mid Y_s'),
\end{equation}
then the lower bound of $p(e)$ is obtained as
\begin{equation}
p(e)\geq H_\text{err}^{-1}(H(Y_m' \mid Y_s')).
\end{equation}
Under Hellman's inequality \citep{hellman1970probability}, we have
\begin{equation}
p(e)\leq \frac{1}{2} H(Y_m' \mid Z'),
\end{equation}
given $ I(Z'; Y_m') \geq I(Y_s'; Y_m')$, the upper bound of $p(e)$ is derived as
\begin{equation}
p(e)\leq \frac{1}{2} H(Y_m' \mid Y_s').
\end{equation}
\subsubsection{Multi-class case}
In a multi-class scenario, we assume $k$ classes, denoted as $\{1, \ldots, k\}$, and the main head classifier $\hat{Y}_m'$ maps the input $z' \in Z'$ to one of these $k$ classes. For $\eta = [\eta_1, \ldots, \eta_k]$, we have
\begin{equation}
h(\eta) = -\sum_{Y_m'=1}^k \eta_{Y_m'} \log \eta_{Y_m'},
\end{equation}
\begin{equation}
H(Y_m' \mid Z') = \mathbb{E}_{z' \sim Z'}[h(\eta(z'))] = \int_{Z'} h(\eta(z')) d p(z'),
\end{equation}
\begin{equation}
\begin{aligned}
p(e) = P(Y_m' \neq \hat{Y}_m')= 1 - \sum_{Y_m'=1}^k \int_{Z'} \mathbbm{1}_{\{Y_m' = \hat{Y}_m'\}} \eta_{Y_m'}(z') dp(z') = 1 - \mathbb{E}_{z' \sim Z'}\left[\max\{\eta(z')\}\right].
\end{aligned}
\end{equation}
Based on the simplified Fano's inequality \citep{thomas2006elements}, we have
\begin{equation}
\begin{aligned}
p(e)& \geq \frac{H(Y_m' \mid Z')-1}{\log (\left | Y_m' \right |)}\\
& \geq \frac{H(Y_m' \mid Z')-1}{\log(k)}\\
\end{aligned}
\end{equation}
Similar to the derivation in binary base, we can obtain
\begin{equation}
p(e) \geq \frac{H(Y_m' \mid Y_s')-1}{\log(k)},
\end{equation}
and when $k\geq4$, the following always holds:
\begin{equation}
\frac{H(Y_m' \mid Y_s')-1}{\log(k)} \leq p(e) \leq \frac{1}{2} H(Y_m' \mid Y_s').
\end{equation}
\iffalse
From Eq.(35) and $e(\eta) = r$ we can derive that $h(\eta)$ is maximized when one element of $\eta$ equals $1 - r$, and the others are $\frac{r}{k - 1}$,
\begin{equation}
h_{\text{max}}(r) = -(1 - r) \log(1 - r) - r \log\left(\frac{r}{k-1}\right).
\end{equation}
In this case the error rate is given by
\begin{equation}
p(e) = 1 - \mathbb{E}_{z' \sim Z'}\big[\max\{\eta(z')\}\big] = 1 - \frac{1}{k}.
\end{equation}
\fi
\iffalse
\begin{enumerate}
\item $h(\eta)$ is maximized when one element of $\eta$ equals $1 - r$, and the others are $\frac{r}{m - 1}$.
\begin{equation}
h_{\text{max}}(r) = -(1 - r) \log(1 - r) - r \log\left(\frac{r}{m-1}\right),
\end{equation}
\item $h(\eta)$ is minimized when all entries of $\eta$ are either $1 - r$ or $0$, except one whose value is determined by the condition that $\eta$ has sum $1$.
\begin{equation}
h_{\text{min}}(r) = -k \cdot (1 - r) \log(1 - r) - \beta \log \beta,
\end{equation}
where $k \cdot (1 - r) \leq 1$ and $\beta = 1 - k \cdot (1 - r)$.
\end{enumerate}
\fi
\section{Dataset, Model and Baselines}
\subsection{Dataset}
\label{data}
\noindent{\textbf{Patient inclusion and exclusion criteria}}. Patients were included in the respiratory failure prediction analysis if they had an ICU stay of at least five hours, were not mechanically ventilated before ICU admission, and had documented vital signs and laboratory values prior to the prediction start time. Those with a Do Not Resuscitate (DNR) order were excluded, and data within 24 hours of surgery were omitted to avoid bias from surgery-related ventilation. Monitoring continued until mechanical ventilation was initiated or ICU discharge. To ensure sufficient data, predictions began four hours post-admission and were updated hourly using the latest clinical information.
\noindent{\textbf{Data abstraction and processing}}. We extracted EHR data encompassing 50 vital signs and laboratory measurements, 6 demographic features, 12 Systemic Inflammatory Response Syndrome (SIRS) and Sequential Organ Failure Assessment (SOFA) criteria, 12 medication categories, and 62 comorbidities. To handle varying sampling frequencies, vital signs and laboratory values were aggregated into hourly time-series bins, with multiple measurements per hour summarized using the median. Data updates occurred hourly, with missing values carried forward for up to 24 hours if no new data were available. Remaining missing values were imputed using the mean. Additionally, we derived 150 features from the 50 vital signs and laboratory measurements, including baseline values (mean over the previous 72 hours), local trends (change since the last measurement), and time since last measured (TSLM).
\noindent{\textbf{Clinical labeling scheme for the various physiological states of respiratory failure}}. Table \ref{tab} lists the labeling criteria used in \citep{lam2024development}.
\begin{table}
\centering
\caption{Criteria of clinical labeling scheme.}
\begin{tabular}{|l|l|c|}
\hline
\textbf{Condition} & \textbf{Criteria} & \textbf{Points} \\ \hline
\multirow{5}{*}{PaO\textsubscript{2}/FiO\textsubscript{2} (not NaN)} & \(200 < \text{PaO}_2/\text{FiO}_2 \leq 300 \, \text{mmHg}\) & 1 \\ \cline{2-3}
& \(\text{PaO}_2/\text{FiO}_2 \leq 200 \, \text{mmHg}\) (severe hypoxemia) & 2 \\ \cline{2-3}
& IMV \(\leq 24 \, \text{hours}\) & 3 \\ \cline{2-3}
& \(\text{PaO}_2/\text{FiO}_2 \leq 200 \, \text{mmHg}\) and IMV \(\leq 24 \, \text{hours}\) & 4 \\ \cline{2-3}
& IMV \(> 24 \, \text{hours}\) & 5 \\ \hline
\multirow{5}{*}{SpO\textsubscript{2}/FiO\textsubscript{2} (not NaN)} & \(141 < \text{SpO}_2/\text{FiO}_2 \leq 221 \, \text{mmHg}\) & 1 \\ \cline{2-3}
& \(\text{SpO}_2/\text{FiO}_2 \leq 141 \, \text{mmHg}\) (severe hypoxemia) & 2 \\ \cline{2-3}
& IMV \(\leq 24 \, \text{hours}\) & 3 \\ \cline{2-3}
& \(\text{SpO}_2/\text{FiO}_2 \leq 141 \, \text{mmHg}\) and IMV \(\leq 24 \, \text{hours}\) & 4 \\ \cline{2-3}
& IMV \(> 24 \, \text{hours}\) & 5 \\ \hline
\end{tabular}
\label{tab}
\end{table}
\iffalse
\noindent If PaO\textsubscript{2}/FiO\textsubscript{2} is not NaN:
\begin{itemize}
\item 1 point: PaO\textsubscript{2}/FiO\textsubscript{2} $> 200 \, \text{mmHg}$ AND PaO\textsubscript{2}/FiO\textsubscript{2} $\leq 300 \, \text{mmHg}$
\item 2 points: PaO\textsubscript{2}/FiO\textsubscript{2} $\leq 200 \, \text{mmHg}$ (severe hypoxemia)
\item 3 points: IMV $\leq 24 \, \text{hours}$
\item 4 points: PaO\textsubscript{2}/FiO\textsubscript{2} $\leq 200 \, \text{mmHg}$ \textbf{and} IMV $\leq 24 \, \text{hours}$
\item 5 points: IMV $> 24 \, \text{hours}$
\end{itemize}
\noindent If SpO\textsubscript{2}/FiO\textsubscript{2} is not NaN:
\begin{itemize}
\item 1 point: SpO\textsubscript{2}/FiO\textsubscript{2} $> 141 \, \text{mmHg}$ AND SpO\textsubscript{2}/FiO\textsubscript{2} $\leq 221 \, \text{mmHg}$
\item 2 points: SpO\textsubscript{2}/FiO\textsubscript{2} $\leq 141 \, \text{mmHg}$ (severe hypoxemia)
\item 3 points: IMV $\leq 24 \, \text{hours}$
\item 4 points: SpO\textsubscript{2}/FiO\textsubscript{2} $\leq 141 \, \text{mmHg}$ \textbf{and} IMV $\leq 24 \, \text{hours}$
\item 5 points: IMV $> 24 \, \text{hours}$
\end{itemize}
\fi
\noindent{\textbf{Encounter-level evaluation}}. Table \ref{tab:eval_metrics} lists the details of evaluation metrics.
\begin{table}[h!]
\caption{Definitions of evaluation metrics based on predictions and labels.}
\centering
\begin{tabular}{@{}ll@{}}
\toprule
\textbf{Metric} & \textbf{Definition} \\ \midrule
\textbf{True Positive (TP)} & A positive prediction (\( \text{predictions[t]} \geq \text{threshold} \)) where there is at least \\ & one
positive label within the prediction window (up to 24 hours before \( T_0\textsuperscript{a} \)). \\ \midrule
\textbf{False Positive (FP)} & A positive prediction (\( \text{predictions[t]} \geq \text{threshold} \)) where no positive labels\\
& exist within the prediction window (up to 24 hours before \( T_0 \)). \\ \midrule
\textbf{False Negative (FN)} & A negative prediction (\( \text{predictions[t]} < \text{threshold} \)) where a positive label \\
& exists within the prediction window (up to 24 hours before \( T_0 \)). \\ \midrule
\textbf{True Negative (TN)} & A negative prediction (\( \text{predictions[t]} < \text{threshold} \)) where no positive labels \\
& exist throughout the evaluated timestamps. \\ \bottomrule
\end{tabular}
\label{tab:eval_metrics}
\vspace{0.2em}
\par \textsuperscript{a} \( T_0 \) is defined as the first timestamp where a patient is ventilated based on the simultaneous recording of PEEP and FiO$_2$.
\end{table}
\iffalse
The encounter with labeling score $\geq$ 3 is assigned one of the following categories:
True Positive: If it is predicted as positive at current timestamp $t$ and there is at least one timestamp with risk score greater than the decision threshold within the prediction window (up to 24 hours before T0 \footnote{T0 is set as the first timestamp where a patient is ventilated based on the simultaneous recording of PEEP and FiO2.}).
False Positive: If there is no timestamp with risk score greater than the decision threshold within the prediction window (up to 24 hours before T0), and there is at least one sample with risk score greater than the decision threshold outside the prediction window.
False Negative: If there is no timestamp with risk score greater than the decision threshold within the prediction window (up to 24 hours before T0).
The encounter with labeling score $<$ 3 is assigned one of the following categories:
True Negative: If there is no timestamp with risk score greater than the decision threshold from start of ICU admission to ICU discharge.
False Positive: If there is at least one timestamp with risk score greater than the decision threshold from start of ICU admission to ICU discharge.
\fi
\iffalse
\begin{table}
\renewcommand{\arraystretch}{1.2}
\centering
\caption{Demographic Comparisons of the development and external cohorts.}
\label{table:char}
\resizebox{\textwidth}{!}{%
\begin{tabular}{lcccccccc}
\hline
\multirow{2}{*}{Demographics} & \multicolumn{2}{c}{UCSD (development site)} & \multicolumn{2}{c}{UCSD} & \multicolumn{2}{c}{UCI} & \multicolumn{2}{c}{MIMIC-IV} \\
& Control & Positive & Control & Positive & Control & Positive & Control & Positive \\
\hline
Encounters,~N (\%) & 29544(94.8) & & & & & & & \\
Age (years), Median (IQR) & & & & & & & & \\
Male, \% & & & & & & & & \\
\begin{tabular}[c]{@{}l@{}}Per Visit ICU Length of Stay \\
(hours),Median (IQR)\end{tabular} & & & & & & & & \\
SOFA score & & & & & & & & \\
Inpatient mortality & & & & & & & & \\
\begin{tabular}[c]{@{}l@{}}Time from ICU admission to \\
initiation of IMV (hours), Median (IQR)\end{tabular} & & & & & & & & \\
\hline
\end{tabular}}
\end{table}
\fi
\vspace{20pt}
\subsection{Network architecture}
\label{archi}
\begin{figure}[ht]
\centering
\includegraphics[width=1\textwidth]{Figs/archi.png}
\caption{The developed network architecture. The encoder consists of a TSLM Layer followed by a feedforward neural network. Both main head and SSL head are feedforward neural networks}
\label{archii}
\end{figure}
Our model follows a Y-shaped design with a shared encoder and two task heads (Fig.~\ref{archii}). At each hourly timestamp we assemble a structured input comprising static demographics, comorbidities/medications, and time-varying vitals/labs augmented with their time-since-last-measurement (TSLM). A lightweight \emph{TSLM layer} ingests each raw value $x_j$ and its recency $\Delta t_j$, applying a learnable decay/gating function to down-weight stale observations and inject a recency embedding; the resulting features are concatenated with the static covariates and passed to a multilayer perceptron encoder $f_e(\cdot;\theta_e)$ to produce a latent representation $z\in\mathbb{R}^d$. Two shallow feedforward heads operate on $z$: a main classifier $\pi_m(\cdot;\theta_c)$ outputs the probability of IMV within 24\,h, and a self-supervised head $\pi_s(\cdot;\theta_s)$ supports reconstruction and masked-feature modeling driven by our dynamic, feature-aware masking scheme. At deployment the classifier and SSL head are frozen, and for each test example we adapt only the encoder for a few gradient steps; the adapted encoder is then used to produce the final risk score, after which weights are reset before the next instance.
\iffalse
\subsection{Baselines}
\begin{itemize}
\item \textbf{TEST:} directly evaluate the pretrained model on the test data without any adaptation;
\item \textbf{TENT} \citep{wang2020tent}: minimize the Shannon entropy of model predictions during test time by optimizing model parameters without requiring annotations;
\item \textbf{TTT} \citep{sun2020test}: jointly train the network on the main task and our defined SSL task in the source domain. During testing, TTT adapts the model to the target domain by continuing training on the SSL task;
\item \textbf{TTT++} \citep{liu2021ttt++}: building on TTT, TTT++ further aligns the first-order and second-order statistics of the source and target distributions;
\item \textbf{ClusT3} \citep{hakim2023clust3}: based on the Y-shaped architecture, we incorporates a projector after the last layer of the shared encoders to enhance test-time training by clustering representations.
\item \textbf{NC-TTT} \citep{osowiechi2024nc}: a noise-contrastive estimation-based auxiliary task for test-time training. The model is adapted at test time to increase in-distribution likelihood without using source samples or labels.
\end{itemize}
\fi
\subsection{Baselines}
\label{ba}
We compare against representative TTA/TTT methods and adapt each fairly to the EHR setting:
\textbf{TEST}, \textbf{TENT} \citep{wang2020tent}, \textbf{TTT} \citep{sun2020test}, \textbf{TTT++} \cite{liu2021ttt++},
\textbf{ClusT3} \citep{hakim2023clust3}, \textbf{NC\mbox{-}TTT} \citep{osowiechi2024nc},
\textbf{T3A} \citep{iwasawa2021test}, \textbf{SAR} \citep{niu2023towards}, and \textbf{CoTTA} \citep{wang2022continual}.
\paragraph{Fairness and EHR-specific protocol.}
All methods use the same data pipeline (hourly aggregation, carry-forward $\leq$24h, mean imputation, and TSLM features), the same shared encoder as ours (Sec.~\ref{archi}), and identical source-domain pretraining (optimizer, weight decay, early stopping). At deployment we enforce the source-free constraint (no source samples/labels). For gradient-based methods we fix the same test-time budget: five update steps, identical step size schedule, and reset-to-pretrained after each instance; the classifier head is frozen unless the baseline explicitly modifies it. Methods that require batch statistics (e.g., BN-based TTA) use the same first-in–first-out buffer of recent test samples to compute moments; the buffer size and confidence thresholds are tuned on the development validation split under the same hyperparameter budget (Bayesian optimization) for all methods. For approaches that rely on data augmentation (e.g., CoTTA), we replace image transforms with tabular augmentation: feature masking.
\paragraph{Methods.}
\begin{itemize}
\item \textbf{TEST}:
evaluate the pretrained model with no adaptation.
\item \textbf{TENT} \citep{wang2020tent}:
minimize prediction entropy at test time; we update only BN affine parameters and running statistics.
%BN moments are computed from the shared FIFO buffer.
\item \textbf{TTT} \citep{sun2020test}: jointly train the network on the main task and the \emph{same} auxiliary SSL objective as ours (reconstruction + masked–feature modeling) in the source domain; at test time, adapt the \emph{encoder only} by minimizing this SSL loss. For fairness, TTT uses \emph{uniform}, feature-agnostic masking (no dynamic masking), and does not use prototypes or optimal transport.
\item \textbf{TTT++} \citep{liu2021ttt++}: identical SSL setup as TTT above and the same encoder-only test-time updates; additionally aligns first/second-order moments between source and target. Because source data are unavailable at deployment, source moments are cached from the development split during training and used for alignment at test time.
\item \textbf{ClusT3} \citep{hakim2023clust3}:
add a projector on top of the shared encoder and adapt by maximizing mutual information with discrete codes. We use an MLP projector (tabular analogue of the original CNN projector) and the same codebook size across sites.
\item \textbf{NC\mbox{-}TTT} \citep{osowiechi2024nc}:
optimize a noise-contrastive auxiliary likelihood at test time. The noise distribution is factorized Gaussian with per-feature mean/variance estimated on the development split.
\item \textbf{T3A} \citep{iwasawa2021test}:
optimization-free classifier adjustment that builds class prototypes from confident test predictions and reweights logits. We compute prototypes from the FIFO buffer, with the confidence threshold tuned on the development split; no backprop or encoder change.
\item \textbf{SAR} \citep{niu2023towards}:
sharpness-aware entropy minimization with unreliable-sample filtering for small/test-time batches. We follow the BN-only update rule as in TENT, add SAM-style perturbations to BN parameters, and use the same FIFO buffer; filter thresholds are validated once on the development split.
\item \textbf{CoTTA} \citep{wang2022continual}:
maintain an EMA teacher and perform augmentation/weight-averaged pseudo-labeling with periodic weight restoration. Tabular augmentation is feature masking; teacher momentum and restoration period are tuned under the shared budget. Encoder is adapted.
\end{itemize}
\begin{figure}
\centering
\includegraphics[width=0.6\textwidth]{Figs/chord_diagram_labels.png} % Adjust width accordingly
\caption{An example of POT between prototypes $\mathbf{P}$ and $z'$.}
\label{fiii}
\end{figure}
\section{Additional results and figures}
\subsection{An Example of Partial Optimal Transport (POT)}
\label{pot}
To further illustrate how our method aligns test-time features with learned prototypes, we visualize an example of Partial Optimal Transport (POT) in Figure~\ref{fiii}. The chord diagram shows the transport plan $\gamma$ between the test-time features $z'$ (top half of the circle) and prototypes $\mathbf{P}$ (bottom half). The width of each arc reflects the transported mass $\gamma_{ij}$ between feature $z'_i$ and prototype $p_j$.
Unlike fixed alignment methods, our formulation allows flexible, soft alignment by augmenting the test-time features with perturbed copies to enable better adaptation to test-time distribution shifts, which prevents overfitting to noisy test samples while preserving meaningful prototype relationships.
\subsection{Additional sensitivity analysis.}
\label{sen}
\noindent{\textbf{Comparison of the size of prototypes}}.
Figure \ref{fp} shows the AUC performance for Site A, Site B and MIMIC-IV across different prototype sizes. Site B and MIMIC-IV exhibit a slight increase in performance as the prototype size increases, while Site A maintains relatively stable AUC values with minimal variation. Site B and MIMIC-IV have more diverse underlying data distribution and while Site A may have more homogeneous patterns. Given that the complexity of external cohorts is unknown in advance, model calibration may be necessary to ensure optimal generalization.
\begin{figure}
\centering
\includegraphics[width=0.45\textwidth, height=0.26\textwidth]{Figs/proto.png} % Adjust width accordingly
\caption{Effect of prototype size on AUC performance across different sites.}
\label{fp}
\end{figure}
\noindent\textbf{Comparison of Prototype Learning.} Our framework leverages prototypes to capture the training domain distribution and facilitate more effective alignment with test-time representations. To assess the contribution of prototype learning, we conduct two ablation studies.
First, as reported in Table~\ref{t1} of the manuscript, we evaluate the PriTTT baseline, which removes the adaptive distribution matching component, thereby isolating the effect of prototype-guided alignment. Second, we examine the impact of prototype learning by replacing end-to-end learned prototypes with post hoc cluster centroids. Specifically, we extract training representations after model training and apply $k$-means clustering. Each cluster is represented by the average embedding of its members, and the resulting $k$ centroids are used as fixed prototypes for distribution matching during test-time training.
The results in Table~\ref{tab:proto_learning} show that models using learned prototypes consistently outperform their post hoc counterparts across all datasets, which highlights the benefit of joint prototype learning and alignment in capturing richer, task-relevant training distributions.
\begin{table}
\centering
\caption{Performance Comparison with and without Prototype Learning}
\begin{tabular}{lccc}
\toprule
\textbf{Method} & \textbf{Site A} & \textbf{Site B} & \textbf{MIMIC} \\
\midrule
AdaTTT & 85.02 $\pm$ 0.05 & 84.10 $\pm$ 0.05 & 77.17 $\pm$ 0.08 \\
DynTTT & 84.54 $\pm$ 0.10 & 83.84 $\pm$ 0.12 & 76.79 $\pm$ 0.05 \\
AdaTTT w/o proto-learn & 84.57 $\pm$ 0.04 & 84.05 $\pm$ 0.06 & 76.85 $\pm$ 0.07 \\
DynTTT w/o proto-learn & 84.35 $\pm$ 0.07 & 83.69 $\pm$ 0.08 & 77.05 $\pm$ 0.10 \\
\bottomrule
\end{tabular}
\label{tab:proto_learning}
\end{table}
\noindent\textbf{Dynamic masking and SSL objectives.}
We ablate the masking strategy and the auxiliary objectives. Replacing dynamic, task-aware masking with random masking degrades AUC by $\sim$1.5–3.0 points across sites (e.g., Site~A: 85.02$\pm$0.05 $\to$ 82.45$\pm$0.05; MIMIC-IV: 77.17$\pm$0.08 $\to$ 74.21$\pm$0.04). Using a single objective alone (\emph{reconstruction-only} or \emph{masked-feature-only}) underperforms the full setup, indicating complementary roles: reconstruction regularizes representations, while masked feature modeling encourages uncertainty-aware recovery of informative variables (Table~\ref{rr}).
% Requires in preamble: \usepackage{booktabs}
\begin{table}
\centering
\small
\renewcommand{\arraystretch}{1.25}
\setlength{\tabcolsep}{6pt}
\caption{Ablation on masking strategy and SSL objectives (AUC \%, $\uparrow$ higher is better).}
\label{tab:ssl_mask_ablation}
\begin{tabular}{lccccc}
\toprule
\textbf{Dataset} & \textbf{TEST} & \textbf{Random Masking} & \textbf{Reconstruction Only} & \textbf{Dynamic Mask Only} & \textbf{AdaTTT (Full)} \\
\midrule
Site A & 84.01 & 82.45 $\!\pm\!$ 0.05 & 83.53 $\!\pm\!$ 0.02 & 84.71 $\!\pm\!$ 0.02 & \textbf{85.02 $\!\pm\!$ 0.05} \\
Site B & 83.75 & 82.60 $\!\pm\!$ 0.06 & 82.47 $\!\pm\!$ 0.05 & 83.18 $\!\pm\!$ 0.04 & \textbf{84.10 $\!\pm\!$ 0.05} \\
MIMIC-IV & 75.28 & 74.21 $\!\pm\!$ 0.04 & 75.00 $\!\pm\!$ 0.04 & 76.24 $\!\pm\!$ 0.02 & \textbf{77.17 $\!\pm\!$ 0.08} \\
\bottomrule
\end{tabular}
\label{rr}
\end{table}
\noindent\textbf{Unified hyperparameter sensitivity analysis.}
As shown in Table~\ref{add}, varying $\lambda_{\text{recon}}$, $\lambda_{\text{proto}}$, $\lambda_{\text{reg}}$, $\lambda_{\text{ot}}$, and warm-up epochs produces only small fluctuations ($\leq 0.1$--$0.3$ AUC) across all cohorts, confirming that AdaTTT is stable to these hyperparameter choices.
\begin{table}[t]
\centering
\small
\caption{Unified hyperparameter sensitivity analysis.
All hyperparameters produce small variations ($\leq$0.1–0.3 AUC), confirming robustness of AdaTTT.}
\begin{tabular}{lccc}
\toprule
\textbf{Hyperparameter Setting}
& \textbf{Site A}
& \textbf{Site B}
& \textbf{MIMIC} \\
\midrule
$\lambda_{\text{recon}} = 0.1$ & $85.01 \pm 0.05$ & $84.12 \pm 0.06$ & $77.16 \pm 0.08$ \\
$\lambda_{\text{recon}} = 0.5$ & $85.03 \pm 0.05$ & $84.11 \pm 0.05$ & $77.18 \pm 0.08$ \\
$\lambda_{\text{recon}} = 1.0$ & $85.00 \pm 0.06$ & $84.09 \pm 0.06$ & $77.14 \pm 0.09$ \\
\midrule
$\lambda_{\text{proto}} = 0.1$ & $85.00 \pm 0.05$ & $84.08 \pm 0.06$ & $77.15 \pm 0.08$ \\
$\lambda_{\text{proto}} = 0.5$ & $85.02 \pm 0.05$ & $84.10 \pm 0.05$ & $77.17 \pm 0.08$ \\
$\lambda_{\text{proto}} = 1.0$ & $84.98 \pm 0.06$ & $84.07 \pm 0.06$ & $77.12 \pm 0.08$ \\
\midrule
$\lambda_{\text{reg}} = 0.1$ & $85.01 \pm 0.05$ & $84.11 \pm 0.05$ & $77.16 \pm 0.08$ \\
$\lambda_{\text{reg}} = 0.5$ & $85.00 \pm 0.05$ & $84.09 \pm 0.05$ & $77.15 \pm 0.08$ \\
$\lambda_{\text{reg}} = 1.0$ & $84.97 \pm 0.06$ & $84.05 \pm 0.06$ & $77.13 \pm 0.09$ \\
\midrule
$\lambda_{\text{ot}} = 0.1$ & $85.02 \pm 0.05$ & $84.11 \pm 0.06$ & $77.18 \pm 0.08$ \\
$\lambda_{\text{ot}} = 0.5$ & $85.04 \pm 0.05$ & $84.12 \pm 0.05$ & $77.20 \pm 0.08$ \\
$\lambda_{\text{ot}} = 1.0$ & $85.01 \pm 0.06$ & $84.09 \pm 0.06$ & $77.16 \pm 0.09$ \\
\midrule
Warm-up = 0 epochs & $84.98 \pm 0.06$ & $84.07 \pm 0.06$ & $77.14 \pm 0.09$ \\
Warm-up = 5 epochs & $85.02 \pm 0.05$ & $84.10 \pm 0.05$ & $77.17 \pm 0.08$ \\
Warm-up = 10 epochs & $85.03 \pm 0.05$ & $84.11 \pm 0.05$ & $77.18 \pm 0.08$ \\
Warm-up = 20 epochs & $85.00 \pm 0.05$ & $84.09 \pm 0.06$ & $77.15 \pm 0.08$ \\
\bottomrule
\end{tabular}
\label{add}
\end{table}
\end{document}