# 🚀 NormWear2.0: A World Model for Multivariate Physiological Signals.
[](https://github.com/Mobile-Sensing-and-UbiComp-Laboratory/NormWear2)
[](https://arxiv.org/)
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.
## 📈 Usage
Full code base available at:
```sh
# Clone the repository
git clone git@github.com:Mobile-Sensing-and-UbiComp-Laboratory/NormWear2.git
```
### Inference Example
```python
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](https://huggingface.co/mosaic-laboratory/normwear2/blob/main/demo_code.ipynb) containing several example of making forecasting inference with and without the proposed latent bayesian mechanism.
> [!NOTE]
> The training and downstream evaluation generally follows the same pipeline as [NormWear](https://github.com/Mobile-Sensing-and-UbiComp-Laboratory/NormWear). The python libraries dependencies are specified in `NormWear/dependencies.txt`, with an example bash script `NormWear/config_env.sh`.
## 🔥 Pre-training
To run the pretraining, run the following command:
```sh
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
```
| Parameter | Type | Example Value | Description |
| --------------------- | ---------: | --------------------: | ------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `--nproc_per_node` | `