lightly-ai / lightly-ai/lightly-train

[Bug] DenseCL updates its key encoder and queues between accumulated microbatches

Open
#973 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.7k
Forks
116
Avg merge
2d 21h
Merged PRs (30d)
6

Description

### Describe the Bug

With gradient accumulation, DenseCL updates the key encoder and both queues on every microbatch. Later microbatches in the same optimizer step therefore see different keys and negatives than earlier ones.

One batch of 128 images and four accumulated batches of 32 do not give the same update, even with identical images, views, initial checkpoint and resolved learning rate. This is the same kind of issue as the DINO teacher/center update in #968, but DenseCL also has two FIFO queues to keep consistent.

### Steps to Reproduce

Run through `lightly_train.pretrain` with:

- 128 ordered, deterministic 32×32 RGB images and fixed two-view transforms;
- the same warm checkpoint and queue initialization;
- memory bank size 256, constant key momentum 0.9;
- SGD without momentum or weight decay, one update, resolved learning rate 0.01;
- either `batch_size=128, accumulate_grad_batches=1` or `batch_size=32, accumulate_grad_batches=4`.

Match the resolved learning rate, not just the configured base rate, since `pretrain` scales it by batch size. Compare the key encoder, both queues and the accumulated student gradient at the optimizer step.

| One batch of 128 vs. four of 32 | Relative L2 difference |
| --- | ---: |
| Current code, student gradient | `0.6732714` |
| Once-per-window update, student gradient | `1.45739e-7` |

With the once-per-window update the key encoder and both queues also match exactly. Freezing the key encoder and both queues removes the difference, which confirms they are the cause. Reproduced on a second host.

### Environment Details

- OS: Linux
- Python version: 3.11
- Frameworks/Libraries (with versions): LightlyTrain 0.17.0 at [c1e2b0dc](https://github.com/lightly-ai/lightly-train/blob/c1e2b0dcabfacc3bd3574204c643bd43e6aa6157/src/lightly_train/_methods/densecl/densecl.py), Lightly `e7c81fa1`, PyTorch 2.13.0+cu130
- How did you install the package: from source at the commit above

### Additional Context

Expected behavior: one optimizer step should use one key-encoder state and the same queue snapshots for every microbatch, then append all of the window's keys once, in input order.

The fix I tested keeps the four microbatches, runs the key EMA once, reads the old queues throughout backward, and commits all 128 keys at the window boundary. This is a one-update test; it says nothing about downstream accuracy, partial accumulation windows or distributed key gathering.

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 src/lightly_train/_methods/densecl/densecl.py and reproduce the issue through lightly_train.pretrain using the deterministic two-configuration setup described. Done means one optimizer step produces matching student gradients, key encoder state, and both queues between one batch of 128 and four accumulated batches of 32, while preserving input order; distributed gathering and partial windows are outside the reported test.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.