mschoi/ai-rfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve geodesic training stability and update best eval baseline.

Browse Source 
Make wrapped torsion loss mandatory, add configurable train-loss early stopping controls, and log new architecture/loss attempts; latest geodesic run improves mean_rmsd_100 to 2.429895 for mainline integration.

Made-with: Cursor
This commit is contained in:
demian3b
2026-04-16 22:20:36 +09:00
parent 35d211f107
commit 3f04a380d8
4 changed files with 211 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@@ -69,3 +69,15 @@ This repository is intentionally pinned to CUDA 12.6 PyTorch wheels and matching
- 2026-04-16: Rollback (per `GUIDELINES.md`): restored `train.py`, `reports/latest_eval.json`, and `artifacts/latest_eval_best_model.pt` to last committed baseline after attempts K–N; `mean_rmsd_100` anchor unchanged at `2.461592` (`BEST_PRACTICE.json`). Objective-aligned early stopping remains disallowed for training control.
- 2026-04-16: Policy update: experiments run on **feature branches** with **one commit per attempt**; mean-RMSD pre-commit gate applies only on **`main`** (merge/cherry-pick integration). Re-triage failed stacks via **cherry-pick** / selective drops, not default full-tree rollback.
- 2026-04-16: Branch `attempt/gat-wrapped-torsion` (single commit batching three evals): Failed O — `gat` + `--torsion-wrapped-loss`, `mean_rmsd_100=2.691410`. Failed P — `gcn` + `--torsion-wrapped-loss`, `2.657594`. Failed Q — `gcn` + `--gcn-residual` (best on branch `2.514058`); all above main best `2.461592` — no merge to `main`.
- 2026-04-16: Branch `attempt/default-wrapped-clean-deps`: Made wrapped torsion loss default (`--torsion-wrapped-loss` via `BooleanOptionalAction`, default on) and added displacement-domain objective option. Dependency cleanup removed unused packages from `pyproject.toml`. Validation run (`python train.py --sdf reports/trajectories/trajectory_00.sdf --epochs 200 --batch-size 32 --eval-runs 100 --model-type gcn --hidden 256 --gcn-layers 6 --loss-domain displacement --seed 1`) reached `mean_rmsd_100=2.528226` (no improvement vs best `2.461592`), so branch not ready to merge.
- 2026-04-16: Branch `attempt/default-wrapped-clean-deps` update: removed `--torsion-wrapped-loss` CLI toggle and enforced wrapped torsion loss always-on in code. Failed R — stronger baseline (`sample.sdf`, `gcn hidden=512 layers=8`, displacement loss, `epochs=800`, `seed=1`) reached `mean_rmsd_100=2.512292`.
- 2026-04-16: Failed S — weighted config (`w_center=0.8, w_omega=2.0, w_torsion=3.0, grad_clip=0.8`, `epochs=1200`, `seed=1`) reached `mean_rmsd_100=2.507794` (better than R, still above best `2.461592`).
- 2026-04-16: Failed T — same weighted config with time-bias (`time_power=1.3`) reached `mean_rmsd_100=2.517704`; no branch promotion.
- 2026-04-16: Attempt U (recommended #1, residual GCN): `--gcn-residual` with weighted displacement setup (`epochs=1200`, `seed=1`) reached `mean_rmsd_100=2.463247` (close, but above best `2.461592`).
- 2026-04-16: Attempt V (recommended #2, SO(3) geodesic rotation loss): initial full-budget run was too slow, then reduced-budget run (`--rotation-loss geodesic --epochs=200 --batch-size=24 --seed=1`) improved to `mean_rmsd_100=2.429729` (new branch best at the time).
- 2026-04-16: Attempt W (recommended #3, split heads + normalization): `--channel-layernorm --head-mlp-layers 2` with weighted displacement setup (`epochs=1200`, `seed=1`) reached `mean_rmsd_100=2.634111` (degraded).
- 2026-04-16: Attempt X (geodesic refinement, longer budget): `--rotation-loss geodesic --epochs=400 --batch-size=24 --seed=2` showed NaN instability and reached `mean_rmsd_100=2.552385`.
- 2026-04-16: Attempt Y (geodesic seed sweep): `--rotation-loss geodesic --epochs=200 --batch-size=24 --seed=3` diverged to NaN early and reached `mean_rmsd_100=2.591940`.
- 2026-04-16: Attempt Z (geodesic stable rerun): same setup as V (`--rotation-loss geodesic --epochs=200 --batch-size=24 --seed=1`) improved further to `mean_rmsd_100=2.426296` (current best in this branch, better than anchor `2.461592`).
- 2026-04-16: Added train-loss-only early stopping controls (`--early-stop-patience`, `--early-stop-min-delta`, `--early-stop-check-every`, `--early-stop-warmup`) with `stop_reason`/`stop_epoch` reporting in logs and `reports/latest_eval.json`; objective-metric stopping remains disabled.
- 2026-04-16: Attempt AA (merge prep rerun on CUDA): repeated geodesic best-practice config (`--rotation-loss geodesic --epochs=200 --batch-size=24 --seed=1`) and measured `mean_rmsd_100=2.429895` (`num_runs=100`), still improving over main anchor `2.461592`.

34
pyproject.toml
View File

@@ -9,27 +9,8 @@ description = "RFM overfitting playground with strict performance gating."
readme = "README.md"
requires-python = ">=3.10,<3.11"
dependencies = [
"numpy>=1.26.4",
"rdkit>=2024.9.5",
"biopython>=1.85",
"hydra-core>=1.3.2",
"transformers>=4.48.0",
"graphein>=1.7.7",
"esm==3.2.0",
"e3nn>=0.5.0",
"jaxtyping>=0.3.2",
"mlcrate>=0.2.0",
"omegaconf>=2.3.0",
"mlflow>=2.0.0",
"tqdm>=4.65.0",
"accelerate>=1.9.0",
"trackio>=0.2.2",
"torchmetrics>=1.8.2",
"tmtools>=0.3.0",
"scikit-learn>=1.5.0",
"torch==2.7.1+cu126",
"torchaudio==2.7.1+cu126",
"torchvision==0.22.1+cu126",
"torch @ https://download.pytorch.org/whl/cu126/torch-2.7.1%2Bcu126-cp310-cp310-manylinux_2_28_x86_64.whl",
"pyg_lib @ https://data.pyg.org/whl/torch-2.7.0%2Bcu126/pyg_lib-0.4.0%2Bpt27cu126-cp310-cp310-linux_x86_64.whl",
"torch-scatter @ https://data.pyg.org/whl/torch-2.7.0%2Bcu126/torch_scatter-2.1.2%2Bpt27cu126-cp310-cp310-linux_x86_64.whl",
"torch-sparse @ https://data.pyg.org/whl/torch-2.7.0%2Bcu126/torch_sparse-0.6.18%2Bpt27cu126-cp310-cp310-linux_x86_64.whl",
@@ -45,16 +26,9 @@ dev = [
"pytest-cov>=6.1.1",
]
[tool.uv.sources]
torch = [{ index = "torch-cu126" }]
torchaudio = [{ index = "torch-cu126" }]
torchvision = [{ index = "torch-cu126" }]
[[tool.uv.index]]
name = "torch-cu126"
url = "https://download.pytorch.org/whl/cu126"
explicit = true
[tool.uv]
default-groups = ["dev"]
cache-keys = [{ file = "pyproject.toml" }]
[tool.setuptools]
packages = []

12
reports/latest_eval.json
View File

@@ -1,10 +1,12 @@
{
"mean_rmsd_100": 2.461592385172844,
"mean_rmsd_100": 2.4298952305316925,
"num_runs": 100,
"timestamp_utc": "2026-04-16T08:43:15.797310+00:00",
"command": "train.py --sdf /data/demian_dev/toy/sample.sdf --model-type gcn --epochs 140 --batch-size 96 --num-workers 8 --prefetch-factor 8 --hidden 512 --gcn-layers 8 --accum 2 --time-power 1.0 --weight-center 0.8 --weight-omega 2.0 --weight-torsion 3.0 --grad-clip 0.8 --lr 0.001 --seed 1 --eval-runs 100 --out-dir /tmp/ai_rfm_try_h_seed1",
"timestamp_utc": "2026-04-16T13:20:04.177521+00:00",
"command": "train.py --sdf /data/demian_dev/toy/sample.sdf --epochs 200 --batch-size 24 --model-type gcn --hidden 512 --gcn-layers 8 --eval-runs 100 --loss-domain displacement --weight-center 0.8 --weight-omega 2.0 --weight-torsion 3.0 --grad-clip 0.8 --rotation-loss geodesic --seed 1",
"notes": "Final test metric from 100 random-initialized rollouts to time=1.",
"best_train_mse": 5.079975128173828,
"best_train_mse": 3.4497900009155273,
"model_source": "best_train_checkpoint",
"checkpoint_path": "/data/demian_dev/toy/ai_rfm/artifacts/latest_eval_best_model.pt"
"checkpoint_path": "/data/demian_dev/toy/ai_rfm/artifacts/latest_eval_best_model.pt",
"stop_reason": "max_epochs",
"stop_epoch": 199
}

204
train.py
View File

@@ -203,6 +203,23 @@ def vel_torsion(current: torch.Tensor, target: torch.Tensor, t: torch.Tensor) ->
return wrap_angle(target - current) / (1 - t).clamp(min=1e-6)
def torsion_wrapped_mse_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
return torch.mean(wrap_angle(pred - target) ** 2)
def so3_geodesic_mse_loss(pred_omega: torch.Tensor, target_omega: torch.Tensor) -> torch.Tensor:
"""Geodesic MSE between batched axis-angle increments."""
if pred_omega.ndim != 2 or pred_omega.shape[-1] != 3:
raise ValueError("pred_omega must be [B,3]")
losses = []
for i in range(pred_omega.shape[0]):
r_pred = so3_exp(pred_omega[i])
r_tgt = so3_exp(target_omega[i])
delta = so3_log(r_pred.T @ r_tgt)
losses.append(torch.sum(delta * delta))
return torch.stack(losses).mean()
def skew(w: torch.Tensor) -> torch.Tensor:
z = torch.zeros((), device=w.device, dtype=w.dtype)
return torch.stack(
@@ -392,6 +409,9 @@ class RFMModel(nn.Module):
num_torsions: int,
hidden: int = 128,
num_layers: int = 4,
gcn_residual: bool = False,
channel_layernorm: bool = False,
head_mlp_layers: int = 1,
):
super().__init__()
self.model_type = model_type
@@ -399,13 +419,19 @@ class RFMModel(nn.Module):
self.num_torsions = num_torsions
self.hidden = hidden
self.num_layers = num_layers
self.gcn_residual = gcn_residual
self.channel_layernorm = channel_layernorm
self.act = nn.SiLU()
self.head_mlp_layers = max(1, int(head_mlp_layers))
if model_type == "gcn":
self.convs = nn.ModuleList()
self.convs.append(GCNConv(3, hidden))
for _ in range(num_layers - 1):
self.convs.append(GCNConv(hidden, hidden))
self.gcn_skip = nn.ModuleList([nn.Identity()])
for _ in range(num_layers - 1):
self.gcn_skip.append(nn.Identity())
trunk_dim = hidden
elif model_type == "mlp":
mlp_layers = []
@@ -421,17 +447,35 @@ class RFMModel(nn.Module):
self.time_feat_dim = 3
d = trunk_dim + self.time_feat_dim
self.trans = nn.Linear(d, 3)
self.rot_omega = nn.Linear(d, 3)
self.torsion = nn.Linear(d, num_torsions)
self.shared_norm = nn.LayerNorm(d) if channel_layernorm else nn.Identity()
def make_head(out_dim: int) -> nn.Sequential:
layers: list[nn.Module] = []
in_dim = d
for _ in range(self.head_mlp_layers - 1):
layers.append(nn.Linear(in_dim, d))
if channel_layernorm:
layers.append(nn.LayerNorm(d))
layers.append(nn.SiLU())
in_dim = d
layers.append(nn.Linear(in_dim, out_dim))
return nn.Sequential(*layers)
self.trans_head = make_head(3)
self.rot_head = make_head(3)
self.torsion_head = make_head(num_torsions)
def forward(self, x: Data, time: torch.Tensor) -> RigidTorsionVelocity:
if self.model_type == "gcn":
h = x.coords
edge_index = x.edge_index
batch = x.batch
for conv in self.convs:
h = self.act(conv(h, edge_index))
for idx, conv in enumerate(self.convs):
h_new = self.act(conv(h, edge_index))
if self.gcn_residual and h_new.shape == h.shape:
h = h_new + self.gcn_skip[idx](h)
else:
h = h_new
trunk = global_mean_pool(h, batch)
else:
bsz = time.shape[0]
@@ -441,10 +485,10 @@ class RFMModel(nn.Module):
t = time.to(dtype=trunk.dtype, device=trunk.device)
two_pi_t = 2.0 * torch.pi * t
time_feats = torch.cat([t, torch.sin(two_pi_t), torch.cos(two_pi_t)], dim=-1)
pooled = torch.cat([trunk, time_feats], dim=-1)
trans = self.trans(pooled)
rot_omega = self.rot_omega(pooled)
torsion = self.torsion(pooled)
pooled = self.shared_norm(torch.cat([trunk, time_feats], dim=-1))
trans = self.trans_head(pooled)
rot_omega = self.rot_head(pooled)
torsion = self.torsion_head(pooled)
return RigidTorsionVelocity(center=trans, rot_omega=rot_omega, torsion=torsion)
@@ -694,6 +738,60 @@ def main() -> int:
default=1.0,
help="Sample t as U(0,1)^time_power; >1 biases training toward t~0 states.",
)
parser.add_argument(
"--loss-domain",
type=str,
default="displacement",
choices=["velocity", "displacement"],
help="Train on raw velocity or displacement-to-target (v * (1 - t)).",
)
parser.add_argument(
"--rotation-loss",
type=str,
default="mse",
choices=["mse", "geodesic"],
help="Rotation-channel loss type.",
)
parser.add_argument(
"--gcn-residual",
action="store_true",
help="Enable residual skip in GCN trunk.",
)
parser.add_argument(
"--channel-layernorm",
action="store_true",
help="Apply layernorm to pooled trunk and head hidden layers.",
)
parser.add_argument(
"--head-mlp-layers",
type=int,
default=1,
help="Per-channel head depth (>=1).",
)
parser.add_argument(
"--early-stop-patience",
type=int,
default=0,
help="Train-loss early stop patience in checks (0 disables).",
)
parser.add_argument(
"--early-stop-min-delta",
type=float,
default=0.0,
help="Minimum train-loss improvement required to reset patience.",
)
parser.add_argument(
"--early-stop-check-every",
type=int,
default=1,
help="Evaluate early-stop condition every N epochs.",
)
parser.add_argument(
"--early-stop-warmup",
type=int,
default=0,
help="Skip early-stop checks for first N epochs.",
)
args = parser.parse_args()
torch.manual_seed(args.seed)
@@ -734,6 +832,9 @@ def main() -> int:
num_torsions=int(sample.bond_endpoints.shape[0]),
hidden=args.hidden,
num_layers=args.gcn_layers,
gcn_residual=args.gcn_residual,
channel_layernorm=args.channel_layernorm,
head_mlp_layers=args.head_mlp_layers,
).to(device)
if args.compile and hasattr(torch, "compile"):
model = torch.compile(model) # type: ignore[assignment]
@@ -750,6 +851,16 @@ def main() -> int:
f"center={args.weight_center} omega={args.weight_omega} torsion={args.weight_torsion}; "
f"grad_clip={args.grad_clip}"
)
print("loss domain=%s torsion_wrapped_loss=always_on" % args.loss_domain)
print(
f"rotation_loss={args.rotation_loss} gcn_residual={args.gcn_residual} "
f"channel_layernorm={args.channel_layernorm} head_mlp_layers={args.head_mlp_layers}"
)
print(
"early_stop: "
f"patience={args.early_stop_patience} min_delta={args.early_stop_min_delta} "
f"check_every={args.early_stop_check_every} warmup={args.early_stop_warmup}"
)
nw = args.num_workers
if nw < 0:
@@ -802,6 +913,10 @@ def main() -> int:
accum = max(1, args.accum)
best_train_mse = float("inf")
best_state: dict[str, torch.Tensor] | None = None
early_stop_best = float("inf")
early_stop_bad_checks = 0
stop_reason = "max_epochs"
stop_epoch = args.epochs - 1
for epoch in range(args.epochs):
optimizer.zero_grad(set_to_none=True)
@@ -815,14 +930,38 @@ def main() -> int:
tt = tt.to(device=device, dtype=torch.float32, non_blocking=True)
pred = model(batched, time_b)
time_remain = (1.0 - time_b).clamp(min=1e-6)
if args.loss_domain == "displacement":
pred_center = pred.center * time_remain
pred_omega = pred.rot_omega * time_remain
pred_torsion = pred.torsion * time_remain
tgt_center = tc * time_remain
tgt_omega = tw * time_remain
tgt_torsion = tt * time_remain
else:
pred_center = pred.center
pred_omega = pred.rot_omega
pred_torsion = pred.torsion
tgt_center = tc
tgt_omega = tw
tgt_torsion = tt
eps = 1e-8
scale_c = (tc * tc).mean().clamp(min=eps)
scale_w = (tw * tw).mean().clamp(min=eps)
scale_t = (tt * tt).mean().clamp(min=eps)
scale_c = (tgt_center * tgt_center).mean().clamp(min=eps)
scale_w = (tgt_omega * tgt_omega).mean().clamp(min=eps)
scale_t = (tgt_torsion * tgt_torsion).mean().clamp(min=eps)
if args.rotation_loss == "geodesic":
rot_loss = so3_geodesic_mse_loss(pred_omega, tgt_omega)
else:
rot_loss = F.mse_loss(pred_omega, tgt_omega)
mse = (
args.weight_center * (F.mse_loss(pred.center, tc) / scale_c)
+ args.weight_omega * (F.mse_loss(pred.rot_omega, tw) / scale_w)
+ args.weight_torsion * (F.mse_loss(pred.torsion, tt) / scale_t)
args.weight_center * (F.mse_loss(pred_center, tgt_center) / scale_c)
+ args.weight_omega * (rot_loss / scale_w)
+ args.weight_torsion
* (
torsion_wrapped_mse_loss(pred_torsion, tgt_torsion)
/ scale_t
)
)
(mse / accum).backward()
sum_mse = sum_mse + mse.detach()
@@ -838,12 +977,43 @@ def main() -> int:
core = model._orig_mod if hasattr(model, "_orig_mod") else model
best_state = {k: v.detach().cpu().clone() for k, v in core.state_dict().items()}
should_check_early_stop = (
args.early_stop_patience > 0
and epoch >= args.early_stop_warmup
and ((epoch - args.early_stop_warmup) % max(1, args.early_stop_check_every) == 0)
)
if should_check_early_stop:
improved = avg_mse_val < (early_stop_best - args.early_stop_min_delta)
if improved:
early_stop_best = avg_mse_val
early_stop_bad_checks = 0
else:
early_stop_bad_checks += 1
if early_stop_bad_checks >= args.early_stop_patience:
stop_reason = (
"early_stop_train_loss_plateau"
f"(patience={args.early_stop_patience},"
f"min_delta={args.early_stop_min_delta},"
f"check_every={max(1, args.early_stop_check_every)})"
)
stop_epoch = epoch
print(
f"early stopping at epoch {epoch}: "
f"bad_checks={early_stop_bad_checks}"
)
break
if epoch % log_every == 0 or epoch == args.epochs - 1:
print(f"epoch {epoch:6d} train_mse {avg_mse_val:.6f} best_train_mse {best_train_mse:.6f}")
stop_epoch = epoch
if device.type == "cuda":
torch.cuda.synchronize()
print(f"training done in {time.perf_counter() - t0:.1f}s")
print(
f"training done in {time.perf_counter() - t0:.1f}s "
f"(stop_reason={stop_reason}, stop_epoch={stop_epoch})"
)
ckpt_path = ""
if best_state is not None:
@@ -895,6 +1065,8 @@ def main() -> int:
"best_train_mse": float(best_train_mse),
"model_source": "best_train_checkpoint",
"checkpoint_path": ckpt_path,
"stop_reason": stop_reason,
"stop_epoch": int(stop_epoch),
}
with open(latest_path, "w", encoding="utf-8") as f:
json.dump(latest_report, f, indent=2)