Aditionnal parameters for the edsnlp.tune function

[LoickChardon]

Feature type

New parameters

Description

The eds.tune function tends to shorten trials too quickly and too often. As observed when tuning a Span Classifier, only very few trials are completed. The results shown in the plots are therefore less reliable than expected, even when setting a very high number of trials (e.g. only 10 out of 500 trials are actually completed). According to the Optuna documentation, the pruner currently used in eds.tune works as follows:

Prune if the trial’s best intermediate result is worse than the median of intermediate results of previous trials at the same step. It stops unpromising trials early based on the intermediate results compared against the median of previous completed trials.

I suggest adding three parameters to the eds.tune function:

disable_pruning: bool — Whether to disable Optuna’s pruning functionality.
Warning: we recommend setting a lower number of trials if True, as it means every trial will be completed.

n_min_trials: int — The minimum number of trials that must be completed before enabling pruning. A higher value increases computation time but leads to a more robust estimation of the median of intermediate results. (This parameter seems to have the same role as n_startup_trials ?)

n_warmup_steps: int — The minimum number of steps before a trial can be pruned. A higher value increases computation time but ensures that Optuna explores each trial long enough before deciding to prune.