gridfm / gridfm/gridfm-graphkit

LayeredWeightedPhysicsLoss reads layer_residuals off the model instead of receiving it

Open
#101 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
105
Forks
36
Avg merge
1d 8h
Merged PRs (30d)
9

Description

Low priority / cleanup. Pre-existing design issue, no known incorrect results today. Surfaced while reviewing #87.

GNS_heterogeneous.forward writes self.layer_residuals as a side effect, and LayeredWeightedPhysicsLoss (training/loss.py:363) reads model.layer_residuals afterwards. The data dependency is invisible in both signatures — the loss just receives model=self.model.

Concrete downsides:

  • Unenforced ordering. Requires forward() before loss_fn() on the same instance. If the dict is empty, L = 0 produces no loop iterations, total_loss stays a Python float, and it fails on total_loss.item() — an obscure error far from the cause.
  • Not reuse-safe. One dict per module instance: two forward passes before a loss call silently overwrite, so gradient accumulation or evaluating two batches before reducing would read the wrong batch's residuals and return a plausible number rather than erroring.
  • Hidden model/loss coupling. Only works with models exposing layer_residuals, i.e. GNS_heterogeneous; AttributeError on GRIT. A config pairing GRIT with this loss fails at runtime, not at validation.

Working as intended today: the stored tensors keep their autograd graph (no .detach()), so gradients flow correctly.

Suggested fix: return residuals from forward alongside the predictions (as #87 now does for embeddings) and pass them into the loss explicitly, so the dependency is in the signature and the model= kwarg is no longer needed for this.

Deferred from #87 because the fix touches the shared BaseLoss.forward signature across all losses.

Refs #87

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with GNS_heterogeneous.forward and LayeredWeightedPhysicsLoss in training/loss.py:363, then inspect the shared BaseLoss.forward signature and the changes referenced in #87. Trace how residuals are produced and consumed, including other losses affected by the shared signature. Done means residuals are passed explicitly, the model= dependency is removed, and the loss no longer relies on stored per-instance state.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.