Papers
Topics
Authors
Recent
Search
2000 character limit reached

A Nationwide Japanese Medical Claims Foundation Model: Balancing Model Scaling and Task-Specific Computational Efficiency

Published 24 Apr 2026 in cs.LG | (2604.22348v1)

Abstract: Clinical risk prediction using longitudinal medical data supports individualized care. Self-supervised foundation models have emerged as a promising approach for leveraging large-scale unlabeled healthcare records. In natural language processing, scaling laws suggest that larger models achieve predictably lower pretraining losses, supporting the foundation model paradigm. However, for structured medical data, characterized by a limited vocabulary and sparse observations, whether increasing model size consistently improves downstream predictions is unclear, as most studies evaluate only a single model scale. In this study, we evaluated the relationship between model scale and downstream task performance for structured medical foundation models. Using a random sample (2.3 million patients, 32 hospitals) from a nationwide 519-hospital Japanese claims database, we pretrained encoder-only Transformers at five scales (2.2M-101M parameters) for disease incidence and medication prediction. Downstream performance saturated at task-dependent thresholds: disease prediction benefited from larger models (32M-101M), whereas medication prediction saturated at 11M, reducing pretraining time by 178 h. Across all tasks, the best-performing model consistently outperformed a Light Gradient Boosting Machine baseline in the area under the precision-recall curve. These findings indicate that, unlike the monotonically decreasing pretraining loss, the optimal model size varied depending on task characteristics. This task-dependent saturation provides practical guidance for balancing predictive performance and computational cost in structured medical foundation models.

Summary

  • The paper demonstrates that increasing model size improves pretraining loss but leads to saturation in downstream performance for specific tasks.
  • It employs encoder-only Transformers on extensive Japanese claims data with combined categorical and continuous embeddings for clinical events.
  • Empirical results highlight that larger models benefit disease prediction while medication prediction saturates at smaller scales, optimizing computational efficiency.

Task-Dependent Model Scaling for Structured Medical Foundation Models in Japanese Claims Data

Introduction

This study introduces a systematic analysis of model scaling in the context of self-supervised foundation models for longitudinal structured medical data, specifically focusing on Japanese nationwide claims records. Unlike prior work emphasizing free-text or single-scale architectures, this paper explores the performance-computation tradeoffs in encoder-only Transformers (2.2M to 101M parameters) pre-trained via masked language modeling and fine-tuned on both disease and medication prediction tasks. A key claim, strongly supported by results, is the task-dependent saturation in downstream discriminative performance, in which increasing model size continues to improve pretraining loss, but larger models do not always translate to better downstream predictive utility. This challenges the paradigm of monotonic scaling gains, providing empirical guidance for optimal model selection under practical computational constraints.

Dataset and Model Architecture

The study leverages a robust dataset of 2.3 million patients from 32 hospitals, sampled from a comprehensive claims/DPC database covering 519 Japanese hospitals. Each patient record is represented as a sequence of clinical events (diagnosis or medication codes paired with age in days and sex), encoded using specialized input embeddings that combine discrete code representations with a Piecewise Linear Encoding (PLE) of continuous age. The models adopt encoder-only Transformer architectures, varying hidden size for five model scales, and are pre-trained using masked language modeling (MLM) objectives. Figure 1

Figure 1: Overview of the pretraining (MLM on clinical sequences with codes and age) and fine-tuning (disease and medication prediction) paradigm, including labeled event rates.

The input structure utilizes a combined categorical and continuous embedding approach: Figure 2

Figure 2: Construction of patient token sequences, showing [CLS], sex, and chronological event tokens with joint MLM/Cross-entropy/MSE loss for codes, sex, and age.

Label frequencies are explicitly reported, with disease targets (hypertension: 13.8%, chronic kidney disease: 2.63%) and medication targets (amlodipine: 14.9%, pregabalin: 6.93%). Positive event sparsity underscores the need for evaluation under significant class imbalance.

Scaling Laws: Pretraining Loss and Computational Cost

The experiments verify that pretraining loss (on MLM) decreases monotonically with larger models, consistent with scaling laws observed in NLP. Reported losses for the largest (101M) model are significantly lower than smaller scales. However, exponential increases in computational cost are observed—from 14.4 hours (2.2M) up to 232.2 hours (101M) on high-performance GPU clusters. Figure 3

Figure 3: Relationship between model scale (parameters, GFLOPs) and pretraining loss.

Fine-Tuning Performance: Task-Dependent Saturation

The core contribution is the demonstration that downstream predictive performance plateaus at task-dependent model sizes:

  • For disease prediction (primary hypertension, CKD), best results (AUPRC) are obtained by large models (32M, 101M).
  • For medication prediction (amlodipine, pregabalin), maximal AUPRC is achieved by the 11M model, with larger models offering no further gains despite quadrupling pretraining time.

This effect is robust across multiple dataset realizations and varying label counts (100, 500, 1,000 fine-tuning samples). Figure 4

Figure 4: Benefit of pretraining versus from-scratch initialization at the 11M scale across tasks and label sizes, measured by AUPRC.

Comparison to LightGBM (a strong tabular baseline) further highlights that the task-optimal pretrained Transformer consistently outperforms LGBM; however, the scale at which the superiority emerges differs between tasks, indicating the necessity of data/task-centric model size selection. Figure 5

Figure 5: AUPRC for five pretrained Transformers (2.2M–101M) and LGBM across all tasks; disease prediction benefits from large models, medication prediction saturates much earlier.

Implications for Model Scaling in Clinical Prediction

These findings provide the first robust evidence in structured healthcare data that scaling encoder-only Transformer models offers diminishing returns dependent on downstream task semantics:

  • Disease onset prediction tasks, requiring more complex longitudinal dependencies, justify substantial parameter growth.
  • Medication prediction tasks, mediated by regular clinical guidelines, saturate at much smaller model scales.

The results directly contradict the universal "bigger is always better" axiom seen in recent free-text and autoregressive event modeling (e.g., Waxler et al. [curiosity_2025]), and underscore that representation capacity must be balanced against the information density and regularity of target prediction tasks. For discriminative tasks on structured medical data, over-scaling leads to wasted computational resources.

Furthermore, the study demonstrates that pretraining yields practical gains over from-scratch initialization, especially notable in low-label regimes and under severe outcome imbalance, with a much stronger improvement in AUPRC than AUROC—a critical distinction in clinical prediction scenarios with rare events.

Practical Recommendations and Limitations

For real-world clinical machine learning deployments with finite resources, these results advocate for explicit determination of task-dependent model capacity ceilings. It is not sufficient to blindly scale up models; empirical saturation analysis can yield large reductions in pretraining time (76% decrease for medication prediction) without any sacrifice in predictive performance. This is particularly salient for institutions aiming to deploy deep models at scale but constrained by compute budgets.

Key limitations include the absence of laboratory/vital data and free-text, restriction to disease and medication prediction (limiting generalizability to other endpoints), lack of cross-dataset external validation, and potential data incompleteness inherent to hospital-based claims. Expansion to multimodal datasets (see [Moor_FMgeneralist_2023]) and additional task types are directions for extension.

Conclusion

This study systematically maps the interaction of model size, computational cost, and downstream task performance for structured medical foundation models in Japanese claims data (2604.22348). The demonstrated task-dependent saturation of predictive performance establishes that both resource efficiency and clinical benefit can be optimized by aligning model scale with the complexity of the prediction target. These results supply actionable guidance for constructing, deploying, and scaling encoder-based foundation models in clinical informatics—moving beyond one-size-fits-all scaling to data- and task-informed model selection.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.