Add early stopping utils

[queensferryme]

This pull request introduces a utility function early_stopping(f; min_delta=0, patience=3) for conveniently implementing early stopping. Its motive was discussed in #227 back in 2018. Also AFAIU, early stopping is widely adopted in training large-scale neural networks nowadays as a mechanism to prevent overfitting. Therefore, I believe a built-in utility function for early stopping would be beneficial for Flux.jl users.

The implementation is heavily based on tf.keras.callbacks.EarlyStopping and ignite.handlers.EarlyStopping.

This is a draft implementation. Documentation and tests are not added yet. As I am new to Flux.jl, I think advice and opinions from maintainers would be helpful before I proceed.

Example Usage:

# test_ea.jl
using Flux

function loss()
    v = 1
    return () -> v += 1
end

ea = Flux.early_stopping(loss(), patience=5)
Flux.@epochs 10 begin
    ea() || break
end

Output:

➜ julia test_ea.jl
 Activating environment at `~/Developer/Flux.jl/Project.toml`
[ Info: Epoch 1
[ Info: Epoch 2
[ Info: Epoch 3
[ Info: Epoch 4
[ Info: Epoch 5
[ Info: Epoch 6
➜

PR Checklist

• Tests are added
• Entry in NEWS.md
• Documentation, if applicable
• API changes require approval from a committer (different from the author, if applicable)

[DhairyaLGandhi]

Thanks for the contribution! Seems generally handy and useful. Do you think it would make more sense as a callback?

[DhairyaLGandhi]

We can always do it that way later though. For now some docs and tests would be helpful.

[DhairyaLGandhi]

Something like this might make sense too once we merge #1471

struct EarlyStopping
  l::AbstractVector
  patience
  tolerance
end

EarlyStopping(memory, patience, tolerance) = EarlyStopping(CircularBuffer(memory), patience, tolerance)

function (es::EarlyStopping)(l, gs, d)
  roi = es.l[end-patience:end]
  if l within tolerance
    push!(es.l, l)
  else
    Flux.stop()
  end
end

[darsnack]

Thanks for your contribution. Can you add a doc string, entry in the documentation, and some tests?

[darsnack]

No need to limit this utility by pigeon-holing it into a callback-biased implementation. I think the current approach in the PR is best. It can easily be used in a CB like:

ea = Flux.early_stopping(loss(), patience=5)
cb = () -> ea() && Flux.stop()

Perhaps we can have a version of stop() that accepts a function. i.e. Flux.stop(f) returns a function that calls Flux.stop() when f() returns true.

[DhairyaLGandhi]

Implementing it as a callback doesn't take away anything from it being used outside of a callback context.

We should also generally want this to be decoupled from the loss function. It would hold unnecessary references alive.

Anyway this would be expected to be used within a train loop 99% of the times, so it's worth it to make that api first class.

[darsnack]

The callback system uses Flux.stop() whereas this just returns true when a condition is met. This implementation is more generic because it can be used in more than Flux.train!. There are other ways to make the Flux.train! API "first-class" without limiting these utilities to only be used in the stopping implementation used by Flux.train!.

[darsnack]

Anyways, let's not open an argument. Like you said, we can do the callback (or not callback) later.

[darsnack]

As I am new to Flux.jl, I think advice and opinions from maintainers would be helpful before I proceed.

The tests should be added under test/utils.jl. You can probably do a test with a pseudo-loss function that you know will hit the early stop criterion in n iterations, run the loop with the early stopping criterion in place. Then test if n matches the number you expect.

Docs will include a docstring and an entry in docs/src/utilities.md.

[queensferryme]

Implementing it as a callback doesn't take away anything from it being used outside of a callback context.

We should also generally want this to be decoupled from the loss function. It would hold unnecessary references alive.

Anyway this would be expected to be used within a train loop 99% of the times, so it's worth it to make that api first class.

I have considered implementing a callback function for early stopping, which seems more natural and integrated into the current Flux workflow. However, here are some considerations that drive me to this closure-based implementation:

1. In my experience, most early stopping happens between loops instead of between batches (within one train loop) because loss generally tends to oscillate greatly between batches. That means early stopping bwtween batches is likely to break the entire training process even when the model is still underfitting.

2. As far as I understand, Flux.stop only terminates the current train loop, whereas in early stopping we'd like that the current train loop along with any follow-up loops are all stopped. Since there aren't any inter-loop hooks for Flux right now, I think manually breaking from loops like ea() || break becomes the only choice.

[darsnack]

Looking good. I think once you figure out the increasing/decreasing metric part, then this will likely be ready to go.

Your feedback on early stopping is appreciated. I think it is just more proof that we should do the most generic thing in this context, which is to write a utility that returns true or false. It will make the utility more flexible to use a variety of contexts. There are several options for how we integrate a boolean utility into the rest of the callback system.

[DhairyaLGandhi]

re 2. since Flux.stop throws, it will actually terminate all loops outside of train and will break the loop when used inside train, so it can be used either way.

[darsnack]

It can be used either way if the user is willing to a) write boiler plate try/catch code or b) accept an ugly exception message. A boolean utility truly can be used however the user wants, because it makes no impositions on the user.

[queensferryme]

re 2. since Flux.stop throws, it will actually terminate all loops outside of train and will break the loop when used inside train, so it can be used either way.

My personal experience tells otherwise:

I believe this is because StopException and SkipException is captured within Flux.train!

Flux.jl/src/optimise/train.jl

Lines 106 to 114 in 8100707

	catch ex
	if ex isa StopException
	break
	elseif ex isa SkipException
	continue
	else
	rethrow(ex)
	end
	end

Maybe this is an unexpected behavior?

[darsnack]

No that's expected. I think what Dhairya was saying is that if you call the callback in your outer loop (over epochs) then it will throw an exception and stop the program. That's not great IMO because it means you either accept an ugly exception, or you write some boilerplate try/catch code to make it pretty. I don't see why we would want the user to write more boilerplate when ea() || break works just fine.

[queensferryme]

Correct, it could get a condition function maybe which can operate on a fixed size of vector of some quantity of interest (typically loss).

I am also considering @DhairyaLGandhi's suggestion. Maybe we could provide a variant of early_stopping that holds a fixed size buffer of the last n collected metrics, as well as a condition function that operates on the buffer to decide whether we should early stop now. This could add more flexibility to the behavior of early stopping.

I will try implement if this sounds like a good idea to you both @DhairyaLGandhi @darsnack.

[darsnack]

Maybe we could provide a variant of early_stopping that holds a fixed size buffer of the last n collected metrics, as well as a condition function that operates on the buffer to decide whether we should early stop now.

My only concern here is calling this "early stopping," since it doesn't match the accepted understanding of the term. To @DhairyaLGandhi's point, the current approach doesn't stop anything, it just returns a boolean that can be used to stop something. If we were to move to this more generic and powerful approach, I would suggest we consider renaming the function.

The other alternative is to build this utility for applying a threshold on a window of values, then make early_stopping be syntactic sugar for passing the correct value for condition so that the it matches the conventional understanding.

[queensferryme]

The other alternative is to build this utility for applying a threshold on a window of values, then make early_stopping be syntactic sugar for passing the correct value for condition so that the it matches the conventional understanding.

Sounds somewhat troublesome and complicated to use 😂 I think I will stick with this implemention for now, which seems more intuitive.

[darsnack]

Since the current design allows for more than just stopping, do we still want to consider a name change? PyTorch's ReduceLROnPlateau is an example that uses the same "early stopping" logic to decide when to advance the schedule. We have the opportunity to do something more generic here. ReduceLROnPlateau is hardcoded to a specific criterion for advancing the schedule, and a specific type of schedule. We can write a more generic version (see here) that works with any underlying schedule and predicate for advancing.

But if we use the utility here as our predicate, then the name early_stopping is weird. And in general, this utility doesn't do any stopping itself.

[queensferryme]

I also thought we were going with patience instead of plateau? The former can be better associated with being more general purpose than plateau would indicate.

But this utility specifically will always be a "plateau." You could turn the patience off completely by setting it to 1. But it's still a plateau because the utility measures when the distance between two points is small enough (by definition a "plateau"). BTW I am not tied to plateau. I just think patience is too uninformative.

For the current implementation, yes it will always be a plateau, and patience may seem too general in this case. In my opinion, a general patience utility should have only three components: a maximum patience, an internal count as well as a predicate that operates over a set of variables and returns eithr true or false. When the predicate says true, the internal count is increased by one; otherwise it is reset to zero.

So, in our case, plateau is more like a specialization of patience, where the predicate is reduced to delta < min_delta. The current plateau will not be able to do something like "stop the training when the loss has been below some threshold for n times". We can do this more general patience thing in another pull request and provide some pre-baked routines maybe, if you are interested.

[darsnack]

What about something like:

function patience(predicate, patience)
  on_trigger = let count = 0
   (args...; kwargs...) -> begin
      count = predicate(args...; kwargs...) ? count + 1 : 0

      return count >= patience
    end
  end

  return on_trigger
end

function plateau(f, p; distance = -, min_dist = 0, init_score = 0)
  on_plateau = let best_score = init_score
    (args...; kwargs...) -> begin
        score = f(args...; kwargs...)
        Δ = distance(best_score, score)
        best_score = Δ < 0 ? best_score : score

        return Δ < min_dist
     end
  end

  return patience(on_plateau, p)
end

const early_stopping = plateau

Personally, I think the patience logic is fairly simple. The hard part is the score closure. If it wasn't for the best score tracking, then I would suggest completely foregoing plateau in the example above. But I think the logic in plateau is cumbersome enough that it warrants its own function.

[queensferryme]

What about something like:

function patience(predicate, patience)
  on_trigger = let count = 0
   (args...; kwargs...) -> begin
      count = predicate(args...; kwargs...) ? count + 1 : 0

      return count >= patience
    end
  end

  return on_trigger
end

function plateau(f, p; distance = -, min_dist = 0, init_score = 0)
  on_plateau = let best_score = init_score
    (args...; kwargs...) -> begin
        score = f(args...; kwargs...)
        Δ = distance(best_score, score)
        best_score = Δ < 0 ? best_score : score

        return Δ < min_dist
     end
  end

  return patience(on_plateau, p)
end

const early_stopping = plateau

Personally, I think the patience logic is fairly simple. The hard part is the score closure. If it wasn't for the best score tracking, then I would suggest completely foregoing plateau in the example above. But I think the logic in plateau is cumbersome enough that it warrants its own function.

This looks great!

[darsnack]

Looking good, thanks. I noticed a couple questions that I missed before.

[darsnack]

Looking really good. I like the doc additions. I think we should be GTM after this last round of revisions.

cc @DhairyaLGandhi so he can look over things

[DhairyaLGandhi]

Cool, I'll give this a once over this weekend

[DhairyaLGandhi]

on_ trigger should be the name of the function imo. Good to have it named for stacktraces

[darsnack]

How is that done with a let?

Like

let ...
  function name(...)
    ...
  end
end

?

[DhairyaLGandhi]

Yeah, I believe so

[darsnack]

@queensferryme If you can address the latest comments when you get a chance, we can finish this PR up! Thanks for all your work and patience.

[darsnack]

Great, thanks @queensferryme!

bors r+

[DhairyaLGandhi]

Thanks all

[queensferryme]

@DhairyaLGandhi @darsnack Thanks for your patience and time!