Skip to content

Fix atom backward mask for attributed edge removal#152

Merged
bengioe merged 1 commit into
recursionpharma:trunkfrom
tgritsaev:fix-qm9-attributed-edge-backward-mask
May 21, 2026
Merged

Fix atom backward mask for attributed edge removal#152
bengioe merged 1 commit into
recursionpharma:trunkfrom
tgritsaev:fix-qm9-attributed-edge-backward-mask

Conversation

@tgritsaev

@tgritsaev tgritsaev commented May 1, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR fixes a mismatch between the atom-level molecular environment's backward
action mask and the parent-counting logic used by the graph-building
environment.

Previously, MolBuildingEnvContext.graph_to_Data allowed RemoveEdge for any
non-bridge edge. This included edges with bond attributes such as
type=BondType.DOUBLE. However, GraphBuildingEnv.parents() and
count_backward_transitions() only allow direct edge removal when the edge has
no attributes; attributed edges should first remove the edge attribute and only
then remove the plain edge.

Why this matters

When sampling uniform backward trajectories from QM9 molecules, the old mask
could sample RemoveEdge directly on a typed bond. The reverse forward action is
then only AddEdge, which does not recreate the original bond attribute. This
can produce inconsistent trajectories where applying the stored forward action
does not recover the stored child graph.

This mainly affects training with uniform backward and results in unstable training.

Change

  • Require direct RemoveEdge to be both non-bridge and attribute-free.
  • Add a regression test with a three-atom cycle containing one typed bond.
  • The test verifies that direct removal of the typed bond is masked out, while
    removing the bond attribute remains allowed, and that the backward mask count
    agrees with count_backward_transitions.

Validation

Locally verified the regression logic with an RDKit-enabled environment:

remove_edge_mask 0.0
remove_edge_attr_mask 1.0
mask_count 3.0
count_backward_transitions 3

Also verified on a QM9 backward-sampling diagnostic that the fix removes the
inconsistent attributed-edge transitions:

before: bad_trajs 21 / 256, bad_steps 21 / 3002
after:  bad_trajs 0, bad_steps 0, count_mismatches 0

@tgritsaev tgritsaev requested a review from bengioe as a code owner May 1, 2026 19:56

@bengioe bengioe left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch!
Thank you for the contribution.

@bengioe bengioe merged commit 9ca98c5 into recursionpharma:trunk May 21, 2026
0 of 4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants