normwear2

<div align="center"> <img src="./normwear2_logo.png" alt="Overview" width="500"/>

🚀 NormWear2.0: A World Model for Multivariate Physiological Signals.

</div>

This is the official implementation of the paper "Toward World Modeling of Physiological Signals with Chaos-Balancing and Latent Dynamical Modeling".

✨ Introduction

Physiological time-series signals reflect complex, multi-scale dynamical processes of the human body. Existing modeling studies primarily focus on static downstream tasks such as classification, event forecasting, or short-horizon next-step prediction, while long-horizon signal-level forecasting and our understanding toward inherently temporal and predictive nature of physiological signals remains underexplored. We introduce NormWear-2, a world model that encodes multimodal physiological signals and intervention variables into a shared latent space and models their temporal evolution through latent dynamics. Our approach combines inference from prior pre-trained knowledge (intuition) with instant unsupervised adaptation for latent transition inference (insight), enabling forecasting across multiple temporal scales, from raw waveform prediction to minute-, hour-, and day-level physiological pattern evolution. During the pretraining phase, we show that effective dynamical modeling is influenced by the diversity of dynamical regimes present in pretraining data. To this end, we propose chaos-theory inspired metrics to quantify time-series dynamics and construct dynamically balanced pretraining corpora. Pretraining on such data leads to more robust latent representations, improving forecasting quality and and maintaining a competitive downstream performance. We experiments across multiple real-world physiological datasets, with diverse temporal resolutions and intervention information at levels of daily-life, point-of-care, and clinical practice. Our results demonstrate that combining dynamics-aware pretraining with latent dynamical modeling consistently improves signal-level forecasting quality in time, frequency, and representation domains, providing a step toward general-purpose world models for physiological signals.

<img src="./method_all_rearrange.png" alt="Overview" width="1000"/>

📈 Usage

Inference Example

import torch
from transformers import AutoModel

import torch
torch.uint64 = torch.int64
torch.uint32 = torch.int32
torch.uint16 = torch.int16

from transformers import AutoModel

model = AutoModel.from_pretrained(
    "mosaic-laboratory/normwear2",
    trust_remote_code=True
)
model.eval()
print("Load Succes!")



# ----- Example forecast I/O -------------------------
# synthetic data
x = torch.rand(2, 256, 3) # bs, L, nvar

# config
context_length = 128
pred_length = 64

# model forward
with torch.no_grad():
    base_out = model.predict(x[:, :context_length, :], pred_length)
print(f"{base_out.shape=}") # base_out.shape=torch.Size([2, 64, 3])

We provide a Demo-code containing several example of making forecasting inference with and without the proposed latent bayesian mechanism.

[!NOTE] The training and downstream evaluation pipelines are inherited from NormWear. The python libraries dependencies are specified in NormWear/dependencies.txt, with an example bash script NormWear/config_env.sh. The full codebase is currently undergoing internal code review and will be published simultaneously with the paper.

🔥 Pre-training

To run the pretraining, run the following command:

torchrun --nproc_per_node=4 \
    -m model_and_pretrain.pretrain_main --num_workers 2 \
    --batch_size 4 --accum_iter 4 \
    --epochs 100 --save_every_epoch 10 \
    --blr 2.5e-4 --min_lr 1e-5 --weight_decay 1e-2 \
    --clip_grad 1.0 --ddp 1 \
    --output_dir train_results/ckpts \
    --log_dir train_results/logs \
    --remark normwear2 \
    --mlp_ratio 4.0 --embed_dim 768 \
    --num_heads 12 --depth 12 --fuse_freq 2 \
    --decoder_embed_dim 512 --decoder_num_head 8 \
    --decoder_depth 2 \
    --window_size 4096 --nvar 10 \
    --use_casual 1 --use_cls 0 --token_level_fuse 1 \
    --jepa 0 

📝 Citation

If you find NormWear-2 model useful for your research, please consider citing the associated preprint.

@misc{luo2026worldmodelingphysiologicalsignals,
      title={Toward World Modeling of Physiological Signals with Chaos-Theoretic Balancing and Latent Dynamics}, 
      author={Yunfei Luo and Xi Chen and Yuliang Chen and Lanshuang Zhang and Md Mofijul Islam and Siwei Zhao and Peter Kotanko and Subhasis Dasgupta and Andrew Campbell and Rakesh Malhotra and Tauhidur Rahman},
      year={2026},
      eprint={2605.15465},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.15465}, 
}

📃 License

This project is licensed under the Apache License 2.0.