lightly-ai / lightly-ai/lightly-train
[Bug] DINO teacher EMA and center update once per microbatch with gradient accumulation
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 116
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 6
Description
### Summary
With gradient accumulation enabled, DINO updates its teacher EMA and its center on every microbatch rather than once per optimizer step. With four microbatches per step, both updates run four times, and later microbatches in the same effective batch see a different teacher and center than earlier ones.
Changing only the accumulation setting therefore changes the teacher targets and the student gradients.
### Environment Details
**Environment:** LightlyTrain 0.17.0, commit [c1e2b0dc](https://github.com/lightly-ai/lightly-train/blob/c1e2b0dcabfacc3bd3574204c643bd43e6aa6157/src/lightly_train/_methods/dino/dino.py), Lightly commit `e7c81fa123baa39334d1590662813798bab9b826`, run through `lightly_train.pretrain`.
The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in [295f04f8](https://github.com/lightly-ai/lightly-train/blob/295f04f8e1208d39119d99589763a61d728819a2/src/lightly_train/_methods/dino/dino.py); the full experiment has not been rerun on that commit.
Separate current-branch regression environment: Ubuntu 26.04 LTS (x86_64), Python 3.11.15, PyTorch 2.13.0+cpu, PyTorch Lightning 2.6.5, and Lightly 1.5.26. LightlyTrain is imported from the source checkout at baseline `295f04f8` plus PR #969, with dependencies installed in a uv-managed virtual environment and `PYTHONPATH=src` for the tests. These are the CPU regression details, not a diagnostic dump from the pinned GPU experiment.
### Steps to Reproduce
The test uses the same 128 ordered synthetic images, deterministic two-view transforms, and the same initial student/teacher checkpoint, and compares:
- one batch of 128 with `accumulate_grad_batches=1`;
- four batches of 32 with `accumulate_grad_batches=4`.
Both perform one SGD update at the same resolved learning rate `0.01`, with teacher and center momentum fixed at `0.9` and gradient clipping at `3.0`.
| Difference between the two settings | Current code (relative L2) | Once-per-window repair (relative L2) |
| --- | ---: | ---: |
| Teacher state after the window | `0.1497115` | `0` |
| Center after the window | `0.6940674` | `0` |
| Student gradient after clipping | `7.45e-4` | `5.94e-8` |
Holding both momentum values at 1 removes the teacher and center differences. The repaired run also matches an independent calculation of the teacher EMA, the DINO loss, and the window-wide center. Reproduced on a second host.
The regression code is in [PR #969](https://github.com/lightly-ai/lightly-train/blob/aa10c27ef71369f19d72aa5620fb99b1c9baeb19/tests/_methods/dino/test_dino.py). From a checkout of that PR with its test dependencies installed:
```bash
python -m pytest -q tests/_methods/dino/test_dino.py
```
The CPU run passed 12 tests. Both native accumulation comparisons fail on the unmodified baseline.
### Expected behavior
With the same logical batch, deterministic views, initial state, and resolved learning rate, gradient accumulation should keep the teacher and center fixed throughout the optimizer window and update them once per window. Their states should agree with the non-accumulated reference.
### Root cause
The teacher and center updates are tied to each training-step call, while Lightning accumulates gradients across several such calls before stepping the optimizer.
Even with the student held fixed, applying the teacher EMA four times changes the weight on the old teacher from `0.9` to `0.9^4 = 0.6561`. Updating the center between microbatches also changes the centering used to build the next microbatch's targets.
### Proposed fix
Reproduce the non-accumulated update order at the window boundary:
1. Update both teacher modules once at the start of the window.
2. Keep that teacher and the previous center fixed for every microbatch in the window.
3. Accumulate teacher-output statistics across the window and update the center once at its end.
The tested repair keeps all four microbatches and follows this order. The teacher EMA and the center need different timing: moving both after `optimizer.step()` would not match the tested non-accumulated path.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/lightly_train/_methods/dino/dino.py and compare its training-step update timing with Lightning gradient accumulation. Run python -m pytest -q tests/_methods/dino/test_dino.py, using the regression comparisons described in PR #969. Done means teacher and center remain fixed across a logical accumulation window and match the non-accumulated reference.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100