deepmodeling / deepmodeling/unimol_tools

[Code scan] Guarantee masked atoms or skip empty pretraining masked losses

Open
#28 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34
Forks
9
PR merge metrics
No merged PRs in 30d

Description

This issue is a result of a Codex global repository scan.

## Summary
Pretraining computes the number of masked atoms as int(sz * mask_prob + random). For small molecules this can be zero. The loss then computes mean token, coordinate, and distance losses over empty tensors, which can produce NaN and poison training.

## Code references
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/dataset.py#L267-L274
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/loss.py#L44-L68
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/loss.py#L79-L99

## Impact
A batch containing molecules with zero selected masked atoms can generate NaN loss values. Once NaNs enter gradients or checkpoints, the pretraining run may become unrecoverable.

## Suggested fix
When mask_prob > 0, clamp num_mask to at least one and at most the atom count. Alternatively, explicitly skip masked objectives for samples or batches with zero masked atoms and avoid mean reductions over empty tensors.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the referenced masking logic in unimol_tools/pretrain/dataset.py around lines 267-274, then trace the masked reductions in unimol_tools/pretrain/loss.py around lines 44-68 and 79-99. Check the behavior for small molecules that select zero atoms and verify the chosen handling prevents NaN losses while preserving valid pretraining objectives.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.