Source

Shen et al. (2024), LIT-146 — the Power scheduler.

The law

Under a warmup-stable-decay schedule (SOTA-140), the optimal learning rate follows a power law in the number of training tokens and the batch size. The Power scheduler applies the law directly: the learning rate is a power of the tokens seen, capped at a maximum, with the WSD decay at the end. A value found on a short run at one batch size then holds when the run is lengthened or the batch changed, and with µP (SOTA-143) when the model is widened, so one sweep serves the whole family.

The law and the scheduler built on it are LIT-146’s. Its authors studied how the optimal learning rate correlates with batch size and token count under WSD, found the power law, and trained their 3B dense and MoE models (PowerLM-3B and PowerMoE-3B) with the Power scheduler combined with µP, which is the case for one setting carrying across model sizes.

Conditions, and what the one adoption counts for

Fitted under WSD; the exponents are the paper’s and were measured on its own model family, so a new family should re-fit them from a few short runs. Falcon-H1-Tiny (LIT-119) is the record’s independent adoption — a square-root decay of the learning rate from 100 GT onward, inside its WSD schedule. It used the law rather than testing it, so under ADR-017 that adoption is evidence about the field and not about the claim: it moves the consensus reading, not the evidence base.

Known implementations

  • PowerLM-3B and PowerMoE-3B; Falcon-H1-Tiny