Skip to content

Fix masking distribution in training collators - #42

Open
OllieinCanada wants to merge 1 commit into
Adibvafa:mainfrom
OllieinCanada:fix/mlm-masking-distribution
Open

OllieinCanada wants to merge 1 commit into
Adibvafa:mainfrom
OllieinCanada:fix/mlm-masking-distribution

Conversation

@OllieinCanada

Copy link
Copy Markdown

Both training collators under-sample the random-replacement branch: after masking 80% of selected positions, they randomize only 10% of the remaining 20%, giving an expected 80/2/18 allocation.

Change that conditional probability to 0.5 in both scripts: half of the remaining 20% gives 10% of selected positions. Preserve the 15% selection rate, amino-acid-specific masks, original labels, organism metadata, attention masks, excluded positions, and random-token sampling bounds.

Thanks to @gui11aume for identifying and explaining the issue. Fixes #28.

Validation on Python 3.10.22 with CPU torch 2.8.0, using the real collators and fake tokenization:

  • python -m unittest tests.test_training_collators -v: 2 methods passed, covering both scripts in 4 subtests. All 4 subtests fail against the original probability. Controlled uniform draws depend on the supplied probabilities and cover a random token coinciding with the original.
  • coverage run -m unittest tests.test_training_collators tests.test_CodonData tests.test_CodonUtils tests.test_CodonJupyter: 22 tests passed, no skips.
  • pre-commit run --files pretrain.py finetune.py tests/test_training_collators.py: all applicable configured hooks passed.
  • make test_with_coverage with HF_HUB_OFFLINE=1: 22 tests ran with one existing prediction-class setup error because adibvafa/CodonTransformer model assets are not cached. No model download was attempted; the published-checkpoint integration tests remain unverified locally.

A seeded CPU experiment on 157,009 selected synthetic positions per script measured 80.0496/1.9024/18.0480% before and 80.0496/9.8918/10.0587% after. These are branch allocations, including random replacements matching the original token. Published checkpoints are unchanged; no accuracy or training-quality improvement is claimed.

This branch has not been deployed

No deployments
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.

Error in MLM logic

1 participant