[BUG] Learning Rate finder does not find optimal lr anymore since version 2.5.3

[dennisbader]

Bug description

Hi all and thanks for the awesome library 🚀

Since #21055 was merged (part of the version 2.5.3 release), the Learning Rate finders do not find the optimal learning rate anymore with noisy data.

Steps to reproduce (you can run below once with python-lightning<2.5.3 and once with >=2.5.3) are below. The example shows the lr finder behavior on synthetic noise-free and noisy losses.

import matplotlib.pyplot as plt
import numpy as np

from pytorch_lightning.callbacks.lr_finder import _LRFinder

use_lightning_le_2_5_3 = True

np.random.seed(42)
n_lrs = 100

# generate losses (noise-free and noisy)
x = np.linspace(-50, 50, n_lrs)
losses = (-x**2 - 30*x + 4000) / 1000.
noise = np.random.uniform(-0.1, 0.1, n_lrs)
losses_noisy = losses + noise

# generate exponential lrs (as done in _LRFinder)
lr_min = 1e-8
lr_max = 1.
lrs = []
for i in range(n_lrs):
    r = (i + 1) / n_lrs
    lrs.append(lr_min * (lr_max / lr_min) ** r)

fig, axes = plt.subplots(ncols=2, figsize=(10, 4))
for ax, losses_, name in zip(
    axes,
    [losses, losses_noisy],
    ["noise-free", "noisy"],
):
    # fill in the results manually to simulate a lr_finder run without actual training
    lr_finder = _LRFinder(mode="exponential", lr_min=lr_min, lr_max=lr_max, num_training=n_lrs)
    lr_finder.results = {
        "lr": lrs,
        "loss": losses_
    }
    # plot with the suggested lr
    lr_finder.plot(suggest=True, ax=ax)
    ax.set_title(f"LR Finder ({name})")

plt.suptitle(f"pytorch-lightning{'<2.5.3' if use_lightning_le_2_5_3 else '>=2.5.3'}")
fig.tight_layout()
plt.show(block=True)

This will give you the following suggested lrs:

As you can see, before 2.5.3, the results weren't affected by noise in the data. After 2.5.3 it does result in suggesting a non-optimal lr in the lower learning rates, since the gradients using the newly introduced spacing have much larger magnitude in these areas.

Having noise in the losses is something that comes naturally in our actual use cases. I'm from the Darts forecasting library where we built our neural network framework around lightning.

Expected behavior

The LRFinder suggestions should ideally not be affected by minor noise in the data (of course with large noise that's a different story).

Proposed solutions

That's a good question.. Maybe we could de-noise / filtering / moving average the loss data in the LRFinder before the gradients are computed?

Thanks again, and if you need more info, let me know!

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

* CUDA:
	- GPU:               None
	- available:         False
	- version:           None
* Lightning:
	- lightning:         2.5.3
	- lightning-utilities: 0.15.2
	- pytorch-lightning: 2.5.2
	- torch:             2.8.0
	- torchmetrics:      1.8.1
* System:
	- OS:                Darwin
	- architecture:
		- 64bit
		- 
	- processor:         arm
	- python:            3.12.11
	- release:           24.3.0
	- version:           Darwin Kernel Version 24.3.0: Thu Jan  2 20:24:16 PST 2025; root:xnu-
* installed with `pip`

More info

No response

cc @lantiga

[daidahao]

Hi everyone, I learnt of this bug from Darts library as well thanks to Dennis's first discovery. I agree with Dennis's assessment that the non-optimal lr came from the new gradient calculation in #21055. Learning rate spacing is being used in 2.5.3 instead of the training step for <2.5.3. Because learning rate spacing is very small in early steps compared to later steps, any small noises in the early training steps would result in much larger gradients than in later steps. That lead to a bias of choosing small learning rates over large ones.

Proposed Solutions
To fix the bug, I suggest partially revert #21055 and calculate the gradients on step spacing (1), instead of on learning rate spacing, effectively reverting gradient calculation back to <2.5.3. While de-nosing/filtering/moving average might help, I do not think they could avoid the bias of choosing small learning rates. Not to mention smoothing (moving average) is already implemented in _LRCallback.

Training loss progresses by the training step, not by the learning rate, even though learning rate is exponentially increased at each step. Moving average in _LRCallback is also performed on step spacing. Therefore, it would make more sense to calculate the gradient on step spacing.

Proposed Changes
A fix should be trivial since we only need to drop spacing from call to torch.gradient(). But thanks to efforts in #21055, we should now have better handling of NaN losses which <2.5.3 did not.

Possible Concerns
I understand the original concerns in #14248 and #21055 where non-homogenous stepping size is not addressed. However, finding a learning rate in low (1e-8, 1e-7) carries a lot more risks than not not finding the "best one" in a range like (1e-3, 1e-2).

[SkafteNicki]

@dennisbader thanks for reporting on this issue and @daidahao thanks for sharing your concerns. I was the one that created PR #21055 so I am to blame for the changes to the suggestion algorithm.

What to do now is the bigger problem... we remove the change and it fix the suggestion for noisy curves which I definitely see the value in or we keep it and it solve the concerns from #14248. Overall I think the solution is to add a parameter to the suggestion algorithm called homogenous_step=True/False (or probably some better name) that enables and disables this behavior. The question is then what the default value should be? To answer this, we probably need to answer:

What scenario is the most common that users will end up in?

If we can clearly answer this, then I am more than willing to implement the change asap.

[dennisbader]

Thanks for quick reply @SkafteNicki and @daidahao. To me it looks like the beta / smoothing parameter of _LRCallback is not exposed to the user.

Giving the user control over this could help, since they get a visual feedback on the smoothing through _LRFinder.plot().

Before we fully put away the idea of using the moving average: I think it has definitely a value since the noise issue can also pop up for the previous <2.5.3 implementation.

We could give the user the option over how many past steps to average on (ma_window or something). See the results below for ma_window=5 and ma_window=10. Of course it would be an approximation to the real smooth curve and the suggested lr will have an error.

Updated code with dummy moving average smoothing

import matplotlib.pyplot as plt
import numpy as np

from pytorch_lightning.callbacks.lr_finder import _LRFinder

use_lightning_le_2_5_3 = False

np.random.seed(42)
n_lrs = 100
ma_window = 10

# generate losses (noise-free and noisy)
x = np.linspace(-50, 50, n_lrs)
losses = (-x**2 - 30*x + 4000) / 1000.
noise = np.random.uniform(-0.1, 0.1, n_lrs)
losses_noisy = losses + noise

def moving_average(a: np.ndarray, n: int):
    test = np.cumsum(a)
    test[n:] = test[n:] - test[:-n]
    return test[n - 1:] / n

losses = losses
losses_noisy = moving_average(losses_noisy, n=ma_window)

# generate exponential lrs (as done in _LRFinder)
lr_min = 1e-8
lr_max = 1.
lrs = []
for i in range(n_lrs):
    r = (i + 1) / n_lrs
    lrs.append(lr_min * (lr_max / lr_min) ** r)

lrs = lrs

fig, axes = plt.subplots(ncols=2, figsize=(10, 4))
for ax, losses_, lrs_, name in zip(
    axes,
    [losses, losses_noisy],
    [lrs, lrs[ma_window - 1:]],
    ["noise-free", "noisy"],
):
    # fill in the results manually to simulate a lr_finder run without actual training
    lr_finder = _LRFinder(mode="exponential", lr_min=lr_min, lr_max=lr_max, num_training=n_lrs)
    lr_finder.results = {
        "lr": lrs_,
        "loss": losses_
    }
    # plot with the suggested lr
    lr_finder.plot(suggest=True, ax=ax)
    ax.set_title(f"LR Finder ({name})")
    print(f"Suggested LR ({name}): {lr_finder.suggestion()}")

plt.suptitle(f"ma_window={ma_window}")
fig.tight_layout()
plt.show(block=True)

>>> ma_window = 5
Suggested LR (noise-free): 1.9054607179632464e-05
Suggested LR (noisy): 1.584893192461114e-05

>>> ma_window = 10
Suggested LR (noise-free): 1.9054607179632464e-05
Suggested LR (noisy): 3.9810717055349735e-05

[SkafteNicki]

@dennisbader I do like the introduction of the moving average parameter which gives better control to users on how smooth the learning rate curve should be and I think it is a fix that can solve most situations. What would you propose the default value of such parameter to be?

[dennisbader]

@SkafteNicki, since this would skip ma_window - 1 steps at the beginning, I think it should either be deactivated (None) by default or could have the same default as skip_begin + 1. If MA is set, then we should probably also enforce skip_begin > ma_window or throw a warning / info.

Regarding your question about the optimal learning rate:

What scenario is the most common that users will end up in?

I would say that cases are common where users use learning rate scheduling with LR reduction: In this case, finding a higher learning rate as the start point is preferred (as the lr would be reduced during training).

So your option to control the spacing with something like homogenous_step=True/False could indeed be helpful.

[SkafteNicki]

@dennisbader could I convince you to open a PR with this new feature / fix?

[dennisbader]

@SkafteNicki yes I can find some time within the next week :) Could we agree before on what should be implemented?

• add new param homogenous_step: bool = True (users are probably more used to the old (<2.5.3) behavior which is why I would set it to True by default. The param would be added to Tuner.lr_find and propagated down to _LRFinder.
• add new param ma_window: Optional[int] = None for the moving average. None would not performing averaging. The param would be added to Tuner.lr_find and propagated down to _LRFinder.
• And if set, then either
  • Enforce 1 < ma_window <= skip_begin
  • Or simply throw a warning or info for the case above?