misato-opentslm-v1b
OpenTSLM-SoftPrompt fine-tuned on MISATO molecular-dynamics trajectories for protein–ligand binding-affinity prediction (pK).
Variant v1b (hybrid): adds a 2-layer MLP regression head on the LLM's last input-position hidden state. Joint loss L = L_LM + λ · MSE(pK_pred, pK_true), λ = 0.5.
Inputs
Four precomputed per-frame channels over a 100-frame window of a 10 ns MD trajectory (fixed order):
| Channel | Source |
|---|
rmsd_ligand | MISATO frames_rmsd_ligand |
interaction_energy | MISATO frames_interaction_energy |
distance | MISATO frames_distance |
bSASA | MISATO frames_bSASA (clipped to [0, 2500] Ų) |
Per-channel z-normalisation using train-split mean/std (see norm_stats.json in the source preprocessing).
Architecture
- Base LLM:
meta-llama/Llama-3.2-1B (frozen) + LoRA r=32 on q/k/v/o + MLP projections
- Trajectory encoder:
TransformerCNNEncoder (Conv1d patch=4 → 6-layer Transformer, embed dim 128) — trainable
- Projector:
MLPProjector to LLM hidden size 2048 — trainable
- Regression head: 2-layer MLP on last-token hidden state — trainable
- Warm-start:
OpenTSLM/llama-3.2-1b-tsqa-sp (5 stages of TSQA curriculum pretraining)
Training
| |
|---|
| Hardware | NVIDIA A100-SXM4-40GB (Lambda Labs) |
| Batch size | 4 |
| Epochs | 5 |
| LR (LoRA / head) | 1e-4 |
| LR (encoder / projector) | 2e-4 / 1e-4 |
| Optimizer | AdamW, weight decay 1e-2, linear warmup 3% |
| Grad clip | 1.0 |
| Wall time | ~728 s/epoch (~60 min total) |
| Train set | 13,758 systems (MISATO MD splits) |
| Val / Test set | 1,595 / 1,612 systems |
Labels: pK derived from misato-affinity/data/affinity_data.csv via pK = 9 − log10(Kd|Ki|IC50 nM), priority Kd > Ki > IC50.
Results
The regression head (head) bypasses Llama's tokenization of float strings; the string_parse numbers come from parsing Answer: X.XX out of the generated rationale.
Val (1595 systems):
| Epoch | RMSE (head) | Pearson R (head) | MAE (head) | RMSE (string) | R (string) |
|---|
| 1 | 1.690 | 0.305 | 1.353 | 1.887 | 0.294 |
| 2 | 1.775 | 0.332 | 1.423 | 2.016 | 0.316 |
| 3 | 1.803 | 0.339 | 1.462 | 1.995 | 0.303 |
| 4 | 1.692 | 0.346 | 1.355 | 1.888 | 0.321 |
| 5 | 1.725 | 0.361 | 1.381 | 1.984 | 0.327 |
Test (1612 systems, held out):
| Epoch | RMSE (head) | Pearson R (head) | MAE (head) | RMSE (string) | R (string) |
|---|
| 1 | 1.595 | 0.329 | 1.299 | 1.799 | 0.304 |
| 2 | 1.709 | 0.338 | 1.406 | 1.920 | 0.336 |
| 3 | 1.720 | 0.351 | 1.420 | 1.918 | 0.320 |
| 4 | 1.597 | 0.371 | 1.311 | 1.803 | 0.328 |
| 5 | 1.651 | 0.370 | 1.357 | 1.912 | 0.335 |
Recommended checkpoint: ckpt_ep4.pt — best test Pearson (0.371) above the mlp_engineered 4-channel ceiling (0.36), and test RMSE essentially tied with epoch 1.
vs baselines
| Test RMSE | Test Pearson |
|---|
predict_train_mean | 1.93 | — |
ols_means | 1.78 | 0.27 |
mlp_engineered | 1.68 | 0.36 |
| v1b ep4 (this model, head) | 1.597 | 0.371 |
First model in this project to beat mlp_engineered on every test metric.
Zero-shot baseline (without fine-tuning)
The same architecture loaded from OpenTSLM/llama-3.2-1b-tsqa-sp with no MISATO fine-tuning was evaluated on the same val + test splits:
| Val (n=1595) | Test (n=1612) |
|---|
Generations with Answer: X.XX | 0 | 0 |
| Generations with any decimal | 0 | 2 |
| RMSE / Pearson R | NaN | NaN |
The pretrained model produces ~30 characters of arbitrary LaTeX/code fragments per system — no structured rationale, no parseable answer. All of this model's predictive power was acquired during the 5-epoch MISATO fine-tuning. Full breakdown with sample outputs in zero_shot_eval.html and pretrained_generations.jsonl.
Files
| File | Size | Purpose |
|---|
ckpt_ep1.pt – ckpt_ep5.pt | ~57 MB each | Per-epoch checkpoints |
ckpt_final.pt | ~57 MB | Same as ep5 |
history.jsonl | 3.6 KB | Per-epoch val + test metrics |
dashboard.html | ~14 KB | Self-contained interactive training dashboard |
zero_shot_eval.html | ~10 KB | Full report on pretrained warm-start (no fine-tuning) |
pretrained_baseline.json | 700 B | Zero-shot metrics summary |
pretrained_generations.jsonl | ~1 MB | All 3207 raw zero-shot generations |
v1b_run.log | 3.2 MB | Full training log |
README.md | this file | |
Each checkpoint contains LoRA adapter weights + the (trainable) encoder, projector, and regression head. The frozen Llama-3.2-1B base must be loaded separately from meta-llama/Llama-3.2-1B.
Honest limits
- 10 ns trajectory ceiling. Experimental Kd/Ki/IC50 reflect thermodynamic equilibrium; 10 ns samples a tiny corner. Hard Pearson R ceiling from physics.
- Train→test pK shift. Train mean 6.59, test 5.55. Calibration on val recommended for downstream use.
- 4-channel summary inputs. This model sees
rmsd_ligand / interaction_energy / distance / bSASA summary channels, not raw 3D coordinates. The data ceiling for 4-channel features is ~0.36 Pearson on this split (see mlp_engineered baseline).
- MISATO selection bias. Only stable-enough-to-simulate systems are in the dataset.
Citation
If you use this model, please cite OpenTSLM:
Stanford BDHG / ETH Zurich. OpenTSLM: Open Time Series Language Model — Curriculum Learning for Multimodal Time Series Understanding with Large Language Models. arXiv:2510.02410.
