PaperlayerView exact revision on arXiv

Source-linked AI summary

ATHENA: Knowledge-guided agentic neural architecture search for AutoFormer-based electronic health record modeling

Deyi Li, Qi Xu, Lingyao Li, Tiansheng Wang, Muxuan Liang, Mei Liu

arXiv:2608.21712v1cs.AIcs.MA

TL;DR

Transformer architectures for EHR prediction still require manual, site- and task-specific tuning, while prior NAS searches typically do not reuse knowledge across hospitals. ATHENA combines weight-sharing evaluation with cross-hospital architecture priors to guide agentic search, achieving the strongest overall AUPRC ranking across search budgets while using less computation than independent candidate pretraining.

  • Problem

    Transformer architectures for EHR modeling remain largely manually designed, and prior searches generally treat each task or hospital as an independent optimization problem without reusable architecture knowledge.

  • Method

    ATHENA combines a weight-sharing AutoFormer-style supernet with reusable cross-hospital architecture priors and multi-agent LLM search guided by target-site validation.

  • Results

    ATHENA achieved the strongest overall AUPRC ranking across search budgets, used less computation than independently pretrained candidates, and produced more consistent architecture selections across repeated searches.

  • Takeaways & Limitations

    Architecture-search experience can be reused across hospitals and clinical prediction tasks while preserving local, task- and site-specific model selection.

  • Takeaways & Limitations

    The evaluation covered only two target health systems and six prediction tasks, limiting conclusions about how broadly the observed transfer patterns generalize.

Abstract

from arXiv · show

Transformer-based models are widely used for clinical prediction from electronic health records (EHRs), yet their architectures still require substantial manual tuning, and the optimal configuration may vary across tasks and hospitals. Neural architecture search (NAS) automates architecture design, but conventional methods are computationally costly for Transformer-based EHR models. Recent large language model (LLM)-guided NAS methods reduce manual search design but typically conduct each search independently, without reusing architecture knowledge across hospitals. In this study, we propose ATHENA (Agentic Transfer across Hospitals for EHR Neural Architecture Search), a knowledge-guided agentic NAS framework for Transformer-based EHR modeling. ATHENA uses a weight-sharing supernet that is pretrained once per hospital, allowing candidate architectures to be instantiated as inherited subnetworks and evaluated through fine-tuning rather than independent pretraining. It also incorporates a two-layer cross-hospital architecture prior. The first layer retrieves high-performing architecture examples from source sites based on task descriptors, while the second estimates the effects of architectural components using SHapley Additive exPlanations (SHAP)-based meta-regression. These priors guide a multi-agent LLM search together with validation feedback from the target hospital. Across six clinical prediction tasks and two independent health systems, ATHENA matches or outperforms four NAS baselines in 9 of 12 hospital-task evaluations at a search budget of 30. It also shows more consistent architecture selection across repeated searches. ATHENA provides a practical approach for reducing manual architecture tuning in Transformer-based EHR modeling. Code is publicly available at https://github.com/GatorAIM/ATHENA.

1. Introduction

ATHENA addresses costly, manually designed Transformer architectures for EHR modeling by reusing pretrained weight-sharing supernets and cross-hospital architecture knowledge. Across six tasks and two health systems, it accelerates evaluation and improves or matches NAS baselines while producing more consistent selections.

  • Motivation: Transformer architectures for EHR modeling remain largely manual, and their optimal configurations vary across prediction tasks and health systems.Differences in cohort composition, clinical workflows, and outcome characteristics may favor different architectural configurations.
  • Motivation: Repeated LLM-guided NAS searches generally operate independently, forcing each new search to rediscover useful architectural patterns.Previous information is not explicitly retained as a reusable prior across related clinical tasks or health systems.
  • Approach: ATHENA uses an AutoFormer-style supernet pretrained once per hospital and reused during search for weight-sharing evaluation and cross-hospital architecture transfer.The framework addresses manual tuning and knowledge-reuse limitations through weight-sharing evaluation and a cross-hospital architecture prior.
  • Results: 15.3-fold speedup: ATHENA accelerates architecture evaluation across six clinical prediction tasks and two independent health systems compared with conventional independent pretraining and fine-tuning.With a limited search budget, ATHENA matches or outperforms conventional and LLM-guided NAS baselines and selects architectures more consistently across repeated searches.

2. Related work

Transformer-based EHR models have advanced through richer temporal modeling, pretraining, and contextualized representation learning, but their architectures remain largely handcrafted. NAS and LLM-guided NAS offer automation, yet LLM-driven NAS has not been systematically explored for longitudinal EHR modeling and existing methods typically search independently.

  • Transformer-based EHR modeling: Transformers are widely used for EHR modeling because they capture dependencies across longitudinal clinical events.BEHRT tokenizes medical codes and models longitudinal histories with Transformer encoders, while Med-BERT demonstrates large-scale pretraining on millions of patient records.
  • Transformer-based EHR modeling: Recent EHR Transformers incorporate temporal-aware embeddings, auxiliary objectives, hierarchical architectures, generative pretraining, autoregressive pretraining, and transferable patient representations.Examples include CEHR-BERT, Hi-BEHRT, TransformEHR, Foresight, and CLMBR.
  • Motivation for NAS: Despite these advances, Transformer-based EHR architectures remain largely handcrafted and are typically selected through manual tuning or limited hyperparameter search.This limitation motivates NAS for systematic exploration of Transformer architectures in downstream clinical prediction.
  • Existing NAS methods: NAS reduces architecture-design costs through defined search spaces, search strategies, performance estimation, gradient-based optimization, and weight-sharing supernets.Transformer NAS methods include Evolved Transformer, HAT, NAS-BERT, AutoBERT-Zero, and Primer, while healthcare applications have focused predominantly on medical imaging, multimodal learning, and selected EHR fusion or multi-task settings.
  • LLM-guided NAS: LLM-driven NAS can generate and refine candidate architectures, but healthcare applications remain limited and longitudinal EHR modeling has not been systematically explored.Existing methods typically treat each search independently, while healthcare examples have focused primarily on medical imaging.

3. Methods · 3.1. Overview · 3.2. Problem formulation

ATHENA is a two-stage, knowledge-guided NAS framework that builds cross-hospital architecture priors from source hospitals and uses them with multi-agent search and target-hospital feedback. It formulates EHR inputs as chronological token sequences and searches constrained AutoFormer-style architectures under a parameter and evaluation budget.

  • 3.1. Overview: ATHENA excludes the held-out target hospital from all metadata used to construct its cross-hospital architecture prior.
  • 3.1. Overview: In Stage 1, ATHENA pretrains one hospital-specific AutoFormer-style supernet per source hospital with masked language modeling and evaluates fixed subnet architectures on source tasks.
  • 3.1. Overview: Pooled source records support two priors: Layer 1 retrieves task-matched high-performing architectures, while Layer 2 summarizes architecture–performance associations across hospitals.
  • 3.1. Overview: In Stage 2, Proposal, Critic, and Experiment Agents generate, refine, evaluate, and adaptively balance exploration with exploitation using priors and target-validation feedback.
  • 3.2. Problem formulation: EHR examples comprise chronologically ordered visits containing diagnosis, medication, laboratory, and procedure tokens, with no assumed within-visit event order.
  • 3.2. Problem formulation: Only visits available at each prediction example’s index time are retained, then flattened with [CLS], token-type embeddings, and visit-index embeddings.
  • 3.2. Problem formulation: Each candidate architecture varies embedding dimension, encoder-layer count, MLP expansion ratio, and attention-head count, subject to candidate-set, parameter-count, and head-divisibility constraints.The parameter-count limit is denoted Pmax, and divisibility ensures an integral attention-head dimension.
  • 3.2. Problem formulation: Target search records validation metric vectors and an ordered evaluation history, ranks metrics with average ties, and deterministically selects the earliest architecture when composite ranks tie.Binary-task F1, AUROC, and AUPRC use positive-class definitions; multilabel metrics are macro-averaged.

3.3. Transformer subnet architecture · 3.4. AutoFormer-style supernet training and subnet evaluation

ATHENA defines Transformer subnetworks by varying embedding dimension, depth, FFN expansion, and attention heads, then uses an AutoFormer-style shared-weight supernet to pretrain once per hospital and evaluate inherited subnetworks through fine-tuning. Candidates are sampled during masked-language-model pretraining, fine-tuned with fixed inherited slices, and ranked using target-validation metrics.

  • 3.3. Transformer subnet architecture: Each architecture is parameterized by embedding dimension d, encoder depth L, FFN expansion γ, and attention-head count h.The architecture is denoted a = (d, L, γ, h).
  • 3.3. Transformer subnet architecture: The input combines a prepended [CLS] token with the flattened EHR token sequence and sums code, token-type, and visit-index embeddings.These embeddings provide the token representation before Transformer encoding.
  • 3.3. Transformer subnet architecture: ATHENA uses a pre-normalization Transformer encoder with MHA, FFN, LN, dropout, and drop-path applied within residual branches.The FFN hidden width is γd, and MHA uses h heads.
  • 3.3. Transformer subnet architecture: The final [CLS] hidden state serves as the patient-level representation for downstream prediction.This representation is taken for the target admission under consideration.
  • 3.4. AutoFormer-style supernet training and subnet evaluation: Instead of training every architecture independently, ATHENA uses an AutoFormer-style shared-weight supernet spanning A₀ and the componentwise maximal search configuration.A subnet activates the first d embedding channels, L encoder blocks, and γd FFN units, while partitioning fixed query/key/value representations into h heads.
  • 3.4. AutoFormer-style supernet training and subnet evaluation: One supernet is pretrained per hospital, with each masked-language-model minibatch sampling a uniformly random architecture and updating only its active parameter slices.During downstream fine-tuning, the architecture is fixed and the same inherited slices are updated each epoch.
  • 3.4. AutoFormer-style supernet training and subnet evaluation: Masked-language-model pretraining corrupts selected non-special tokens by masking, same-modality random replacement, or leaving them unchanged.The masked positions are denoted by Ω, and the loss uses the corrupted-token representation with active subnet weights.
  • 3.4. AutoFormer-style supernet training and subnet evaluation: Each candidate inherits encoder weights from the supernet checkpoint, is fine-tuned on the target training split, and contributes target-validation metrics to the search archive.These validation records support subsequent candidate evaluation.

3.5. Cross-hospital prior construction

ATHENA constructs a two-layer cross-hospital architecture prior from source-hospital metadata. The priors provide task-matched retrievals and directional architectural evidence to guide agents without imposing hard constraints.

  • Layer 1 retrieval: Layer 1 retrieves the top K legal architectures from task-matched source hospitals using composite source-performance rankings.When exact task metadata are unavailable, ATHENA uses dataset-profile similarity to select a source hospital and surrogate task.
  • Task descriptors: The target task descriptor uses data-dependent features computed only from the target hospital’s downstream training split.These features include normalized output classes, label entropy, positive-label prevalence, and normalized prediction horizon.
  • Layer 2 meta-analysis: Layer 2 pools architecture evaluations across source hospitals and fits task-specific surrogate models with hospital identity retained for mixed-effects analysis.Discrete architecture features are represented numerically for the XGBoost surrogate.
  • Layer 2 meta-analysis: SHAP contributions yield preferred, discouraged, or inconclusive feature levels and supported interaction rules based on confidence intervals relative to zero.Feature importance, level labels, confidence intervals, and interaction rules form the Layer 2 prior and are presented as non-causal, directional evidence.

3.6. Agentic NAS

ATHENA uses coordinated LLM agents and deterministic controller operations to propose, critique, revise, evaluate, and select architectures for each hospital-task pair. Its ordered search transcript carries accumulated search experience, while a structured validation archive remains the authoritative record of empirical evidence.

  • Agent coordination: ATHENA coordinates a Proposal Agent, a Critic Agent, and an Experiment Agent comprising a deterministic evaluator and an LLM-based strategy module.LLM calls handle proposal, critique, revision, and strategy selection; controller operations handle legality checks, subnet fine-tuning, metric calculation, ranking, and final selection.
  • Search context: For each target hospital-task pair (𝐻, 𝑡), ATHENA provides the task descriptor, 𝐾 retrieved architectures, and the Layer 2 prior as fixed context.
  • Search state: ATHENA maintains an ordered search memory containing proposals, critiques, revisions, evaluation feedback, and strategy decisions, alongside a structured archive containing evaluated architectures and validation metrics.The memory supplies LLM context, whereas the archive supports deterministic ranking and selection; the memory is not a separately learned or retrieval-based long-term-memory system.
  • Search strategy: Exploration targets under-sampled architectural choices, whereas exploitation concentrates proposals around configurations that performed well on target validation.
  • Proposal and critique: The Proposal Agent generates candidates with rationales, while the Critic checks legality and novelty and permits rejected non-duplicate candidates to undergo revision for at most 𝑅 passes.Parameter-limit violations and duplicates are hard rejections, whereas disagreement with a Layer 2 preference is only a soft concern.
  • Termination: Search terminates when the budget 𝐵 is exhausted or after 𝐹max consecutive empty rounds, returning the current validation-selected leader; runs with no evaluated candidate are unsuccessful.A successful evaluation round resets the consecutive-failure counter, and an empty round consumes no evaluation budget.

4. Experiment configuration

Experiments use OneFlorida+ and MIMIC-IV EHR data, with cross-hospital prior construction separated from held-out targets. Evaluation covers six clinical prediction tasks using a shared 256-architecture Transformer search space, standardized splits, and comparisons with classical, LLM-based, and ablated ATHENA variants.

  • Datasets and cohort design: Experiments use four OneFlorida+ prior-source sites, one internal held-out target site, and MIMIC-IV as an external held-out target.Records from both target cohorts are excluded from the cross-hospital prior, while both datasets follow the same EHR tokenization, task-construction, and split conventions.
  • Datasets and cohort design: Each hospital’s unlabeled EHR corpus is separated from downstream labeled patients, with the pretraining pool split 90%/10% for masked-language-model supernet pretraining and validation.Downstream training supports subnet fine-tuning, validation guides architecture search and model selection, and test data are evaluated once after search termination.
  • Downstream tasks and splits: Evaluation covers six tasks: Mortality, Stay > 7d, Readmission (3M), Phenotype (6M), Phenotype (12M), and same-visit Drug Recommendation.The phenotype tasks are multilabel predictions over 18 classes; binary tasks use 20%/40%/40% training/validation/test splits, while phenotype and drug recommendation use 40%/30%/30% splits.
  • Search space: The shared search space varies embedding dimension, depth, attention heads, and MLP expansion ratio, yielding 256 candidate architectures from their Cartesian product.Embedding dimension takes values in {32, 64, 128, 256}; each other parameter takes values in {1, 2, 4, 8}.
  • Baselines and ablations: ATHENA is compared with Random Search, EA, GENIUS, and CoLLM-NAS, alongside L1-only, Leave-One-Task-Out Retrieval, and Cold Start ablations.The ablations respectively remove Layer 2, exclude task-matched retrieval records, or remove both cross-hospital priors.

5. Results

ATHENA’s supernet and cross-hospital priors enabled efficient, high-fidelity architecture search across six clinical tasks and two health systems. It achieved leading performance and more consistent selections without relying on the largest models, while both prior layers contributed to performance.

  • Supernet validation: Positive Spearman correlations across every task and metric showed that inherited subnetworks preserved independently trained architecture rankings.Correlations reached 0.940–0.951 for Phenotype (6M)/(12M) and 0.944–0.958 for Drug Recommendation.
  • Supernet validation: 15.3× speedup at 30 evaluations on MIMIC-IV reduced estimated cost from 900 to 59 GPU-minutes; OneFlorida+ reached 10.7× and 960 to 90 GPU-minutes.At five evaluations, speedups were 4.0× on OneFlorida+ and 4.4× on MIMIC-IV.
  • Search performance: ATHENA ranked first in 9 of 12 comparisons at both 20 and 30 evaluations, with average ranks of 1.29 and 1.42.At budget 30, ATHENA also ranked first in 7 of 12 AUROC comparisons and achieved the best average rank of 1.71.
  • Search performance: ATHENA’s gains generally emerged within the first few evaluations and remained leading or competitive, especially for Phenotype (6M), Phenotype (12M), and Drug Recommendation.Later evaluations produced smaller, task-dependent improvements, while Mortality, Stay > 7d, and Readmission (3M) converged more closely.
  • Architecture analysis: ATHENA averaged 1.55 million selected parameters on OneFlorida+, versus 1.90 million for GENIUS, and achieved the highest MIMIC-IV modal selection rate at 43.3%.On OneFlorida+, ATHENA and EA shared the highest mean modal selection rate at 40.0%.
  • Ablation analysis: Removing Layer 2 reduced AUPRC on several tasks, while removing both prior layers caused the clearest degradation for Phenotype (6M) and Phenotype (12M) on both targets.LOTO Retrieval remained comparable to exact retrieval and was occasionally higher, supporting related-task records as fallback guidance.

6. Discussion

ATHENA reuses cross-hospital architecture knowledge while preserving target-site validation and weight-sharing evaluation for Transformer-based EHR neural architecture search. Across six clinical prediction tasks and two target health systems, it showed strong performance, computational advantages, and more consistent architecture selection, while broader generalization remains untested.

  • Framework: ATHENA combines weight-sharing evaluation with reusable cross-hospital architecture knowledge for Transformer-based EHR modeling.Candidate architectures are evaluated as inherited subnetworks rather than through independent pretraining.
  • Cross-hospital priors: The retrieval prior supplies architecture examples, while the architecture-effect prior estimates task-specific effects of individual design choices across hospitals.Removing the architecture-effect prior reduced performance on several tasks, and removing both priors produced larger declines, particularly for multilabel outcomes.
  • Cross-hospital priors: LOTO Retrieval remained competitive without exact task matches, indicating that related prediction problems can provide useful architectural guidance for new tasks.The observed task-specific architecture effects also argue against a single universally optimal architecture.
  • External evaluation: On MIMIC-IV, source-hospital architecture knowledge remained informative despite different cohort characteristics and clinical data distributions.Source knowledge guided candidate generation, while final selection continued to depend on target-site validation rather than directly transferring a source architecture.
  • Limitations: The evaluation covered two target health systems and six downstream prediction tasks, but additional institutions, specialties, coding systems, populations, and prediction horizons are needed to assess generalization.MIMIC-IV represents a single academic medical center and does not capture the heterogeneity of potential real-world deployment environments.
  • Conclusion: ATHENA reframes architecture search as reusing search knowledge across hospitals and clinical prediction tasks while retaining task- and site-specific adaptation.Its framework combines retrieved architecture examples, architecture-effect guidance, target-site validation, and weight-sharing evaluation under limited search budgets.

CRediT authorship contribution statement

The CRediT statement assigns core conceptual, analytical, methodological, software, and drafting contributions to Deyi Li and Qi Xu, with additional validation and leadership roles distributed among the coauthors.

  • Deyi Li contributed to conceptualization, data curation, formal analysis, investigation, methodology, software, visualization, and original-draft writing.
  • Qi Xu contributed to conceptualization, data curation, formal analysis, investigation, methodology, and original-draft writing, while Lingyao Li, Tiansheng Wang, and Muxuan Liang contributed validation and original-draft writing.
  • Mei Liu contributed to conceptualization, funding acquisition, project administration, resources, supervision, validation, and original-draft writing.

Data availability

Patient-level OneFlorida+ data are not publicly available because of privacy, institutional, and data-use restrictions, but eligible researchers may request access through the network’s Front Door or Coordinating Center.

  • Data availability: OneFlorida+ data access requires eligibility, network and institutional review, applicable data-use agreements, and ethical or institutional review board approvals.Requests may be submitted through the OneFlorida+ Clinical Research Network Front Door or Coordinating Center.

Supplementary materials for:

The supplementary materials define cohort eligibility and the shared legal architecture/search space used across methods. They also specify ATHENA’s cross-hospital prior, agentic-search limits, and candidate fine-tuning and evaluation procedures.

  • Cohort criteria: Admissions required diagnosis, medication, procedure, and laboratory records, included only adults aged ≥18 years, and excluded patients with more than eight distinct admissions.Ages ≥89 years were coded as 90 and discretized into 20 equal-width bins; no minimum event count or admissions per patient was otherwise required.
  • Architecture constraints: The legal architecture space imposed Pmax = 4,000,000, used maximal configuration (d, L, γ, h) = (256, 8, 8, 8), and fixed query/key/value width at 256.All methods used the same hospital-specific pretrained supernet.
  • Cross-hospital prior: ATHENA’s cross-hospital prior used four OneFlorida+ source sites, 100 architectures per source-hospital–task pair, and retrieved K = 5 source architectures.Source metrics averaged the three highest source-validation-AUPRC checkpoints, while target cohorts were excluded from prior construction; Layer 2 effects used normal-approximation 95% CIs and required at least three records per interaction cell.
  • Agentic search: The agentic search allowed at most R = 3 critique passes per round, stopped after Fmax = 3 consecutive rounds without an accepted legal candidate, and used maximum budget B = 30.Results were summarized after 5, 20, and 30 evaluations.
  • Candidate evaluation: Candidate fine-tuning ran for at most 30 epochs with validation-AUPRC patience 5, search-time metrics averaged the three best validation-AUPRC epochs, and final weights came from the single best epoch.Attention dropout, residual dropout, and drop-path rate were each 0.1.
Loading 2608.21712v1…