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

1. 10 ns trajectory ceiling. Experimental Kd/Ki/IC50 reflect thermodynamic equilibrium; 10 ns samples a tiny corner. Hard Pearson R ceiling from physics.
2. Train→test pK shift. Train mean 6.59, test 5.55. Calibration on val recommended for downstream use.
3. 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).
4. 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.