surestop 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- surestop-0.2.0/LICENSE +21 -0
- surestop-0.2.0/PKG-INFO +173 -0
- surestop-0.2.0/README.md +148 -0
- surestop-0.2.0/pyproject.toml +36 -0
- surestop-0.2.0/setup.cfg +4 -0
- surestop-0.2.0/src/surestop/__init__.py +17 -0
- surestop-0.2.0/src/surestop/core.py +201 -0
- surestop-0.2.0/src/surestop/optuna_pruner.py +99 -0
- surestop-0.2.0/src/surestop.egg-info/PKG-INFO +173 -0
- surestop-0.2.0/src/surestop.egg-info/SOURCES.txt +13 -0
- surestop-0.2.0/src/surestop.egg-info/dependency_links.txt +1 -0
- surestop-0.2.0/src/surestop.egg-info/requires.txt +9 -0
- surestop-0.2.0/src/surestop.egg-info/top_level.txt +1 -0
- surestop-0.2.0/tests/test_core.py +88 -0
- surestop-0.2.0/tests/test_optuna_pruner.py +48 -0
surestop-0.2.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Adam Noonan
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
surestop-0.2.0/PKG-INFO
ADDED
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: surestop
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Stop bad runs early, with a conformal bound on how often you stop one that would have ended well
|
|
5
|
+
Author: Adam Noonan
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/ACNoonan/surestop
|
|
8
|
+
Project-URL: Issues, https://github.com/ACNoonan/surestop/issues
|
|
9
|
+
Keywords: hyperparameter-optimization,early-stopping,pruning,conformal,optuna
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Intended Audience :: Science/Research
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
14
|
+
Requires-Python: >=3.10
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
License-File: LICENSE
|
|
17
|
+
Requires-Dist: numpy>=1.24
|
|
18
|
+
Requires-Dist: scipy>=1.10
|
|
19
|
+
Provides-Extra: optuna
|
|
20
|
+
Requires-Dist: optuna>=4; extra == "optuna"
|
|
21
|
+
Provides-Extra: test
|
|
22
|
+
Requires-Dist: pytest; extra == "test"
|
|
23
|
+
Requires-Dist: optuna>=4; extra == "test"
|
|
24
|
+
Dynamic: license-file
|
|
25
|
+
|
|
26
|
+
# SureStop
|
|
27
|
+
|
|
28
|
+
**Stop bad training runs early, with a stated bound on how often you stop one that would have
|
|
29
|
+
ended well.** You pick α. The rule kills runs so that the expected fraction of runs that were
|
|
30
|
+
good *and* got killed stays at or below α.
|
|
31
|
+
|
|
32
|
+
Every pruner in common use (median, percentile, successive halving, Hyperband) is a heuristic,
|
|
33
|
+
and none says how often it kills a run that would have finished well. This one does, under a
|
|
34
|
+
condition you can check: the runs you calibrate on and the runs you judge must come from the same
|
|
35
|
+
random process.
|
|
36
|
+
|
|
37
|
+
It was measured on 400 real LoRA fine-tuning runs under a pre-registered protocol. The curves are
|
|
38
|
+
in this repository, and a script checks that the package reproduces the recorded numbers exactly.
|
|
39
|
+
|
|
40
|
+
## Install
|
|
41
|
+
|
|
42
|
+
```
|
|
43
|
+
pip install surestop # numpy + scipy
|
|
44
|
+
pip install "surestop[optuna]" # adds the Optuna pruner
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
## Use
|
|
48
|
+
|
|
49
|
+
Plain numpy, on any training loop:
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from surestop import KillRule
|
|
53
|
+
|
|
54
|
+
# completed_curves: array of shape (runs, evaluations), lower is better
|
|
55
|
+
rule = KillRule(alpha=0.05, good_threshold=2.16, min_peek=17).fit(completed_curves)
|
|
56
|
+
|
|
57
|
+
rule.should_kill(curve_so_far) # True -> stop this run now
|
|
58
|
+
rule.kill_step(full_curve) # where it would have fired, or None
|
|
59
|
+
rule.evaluate(held_out_curves) # false kills, joint rate, savings
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
With Optuna, the first `n_calibration` trials run unpruned to calibrate, then the rule freezes:
|
|
63
|
+
|
|
64
|
+
```python
|
|
65
|
+
import optuna
|
|
66
|
+
from surestop import ConformalPruner
|
|
67
|
+
|
|
68
|
+
study = optuna.create_study(sampler=optuna.samplers.RandomSampler(),
|
|
69
|
+
pruner=ConformalPruner(n_calibration=60, alpha=0.05))
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Pass `maximize=True` (or a `direction="maximize"` study) when higher is better.
|
|
73
|
+
|
|
74
|
+
## The rule
|
|
75
|
+
|
|
76
|
+
A run is *good* if its final value ends at or below a threshold τ. At every evaluation, a
|
|
77
|
+
predictor forecasts the final value from the curve so far. The default predictor is the latest
|
|
78
|
+
observed value.
|
|
79
|
+
|
|
80
|
+
To calibrate, take n runs that have already finished. For each good one, record its worst moment:
|
|
81
|
+
the largest gap between forecast and τ at any evaluation. Runs that are not good contribute
|
|
82
|
+
nothing. Set λ to the ⌈(1−α)(n+1)⌉-th smallest of these scores. A new run is killed at the first
|
|
83
|
+
evaluation where its forecast exceeds τ + λ.
|
|
84
|
+
|
|
85
|
+
This is conformal risk control on the loss "good and killed". Taking the worst moment across all
|
|
86
|
+
evaluations means one threshold covers every look at a run, so there is no correction for
|
|
87
|
+
multiple looks. The guarantee needs the calibration runs and the new runs to be exchangeable.
|
|
88
|
+
Random sampling gives that. Adaptive samplers such as TPE do not.
|
|
89
|
+
|
|
90
|
+
## What was measured
|
|
91
|
+
|
|
92
|
+
Experiment CC-03c, pre-registered before any run finished:
|
|
93
|
+
|
|
94
|
+
- **Runs:** 400 LoRA fine-tunes of Qwen2.5-0.5B on dolly-15k, 750 steps each, validation loss
|
|
95
|
+
measured 50 times per run.
|
|
96
|
+
- **Configs:** i.i.d. draws from a frozen hyperparameter prior (`docs/hyperparameter-prior.md`).
|
|
97
|
+
- **τ:** pinned from an earlier pool before the first run landed.
|
|
98
|
+
- **Splits:** two, one random and one past-to-future by night, each about half calibration and
|
|
99
|
+
half test.
|
|
100
|
+
|
|
101
|
+
The pre-registered rule used a fitted power-law predictor. On both splits it kept the bound and
|
|
102
|
+
saved about half the evaluation compute:
|
|
103
|
+
|
|
104
|
+
| split | false kills | joint rate (95% CI) | eval compute saved |
|
|
105
|
+
|---|---|---|---|
|
|
106
|
+
| random | 13 of 200 | 6.5% (3.5–10.9%) | 52.2% |
|
|
107
|
+
| past-to-future | 11 of 203 | 5.4% (2.7–9.5%) | 52.8% |
|
|
108
|
+
|
|
109
|
+
A rule that kills every run immediately was scored by the same gate and failed on both splits.
|
|
110
|
+
|
|
111
|
+
**The flaw that run found:** every kill fired at the 3rd of 50 evaluations. Early rankings only
|
|
112
|
+
match final rankings from about evaluation 18. The rule stayed within budget because the bound
|
|
113
|
+
held. The fitted power law ranked runs worse than the latest observed value at every early
|
|
114
|
+
evaluation.
|
|
115
|
+
|
|
116
|
+
**What this package ships instead,** and what the reproduction script checks:
|
|
117
|
+
|
|
118
|
+
| predictor | hold until | random split | past-to-future split | 2000 resplits |
|
|
119
|
+
|---|---|---|---|---|
|
|
120
|
+
| last value | none | 18 of 200 (9.0%), 78% saved | 7 of 203 (3.4%), 75% saved | joint 5.0%, breach 8.5% |
|
|
121
|
+
| last value | eval 18 | 18 of 200 (9.0%), 57% saved | 7 of 203 (3.4%), 56% saved | joint 5.0%, breach 8.6% |
|
|
122
|
+
|
|
123
|
+
The random split's 18 of 200 is a breach. It sits at the 95th percentile of that rule's own
|
|
124
|
+
resplit distribution, whose mean is 5.0%, so it reads as a split tail rather than a leak. A
|
|
125
|
+
correct rule at this sample size lands on a breach in about 7% of splits. The held version was
|
|
126
|
+
chosen *after* seeing these 400 runs, so its numbers are a hypothesis. It halves
|
|
127
|
+
how often the single best run is killed at the same budget. Its confirmation needs a fresh pool.
|
|
128
|
+
|
|
129
|
+
Full account: `docs/writeup.md`. Run the check yourself:
|
|
130
|
+
|
|
131
|
+
```
|
|
132
|
+
python reproduce/reproduce_cc03c.py
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
## What the bound does not cover
|
|
136
|
+
|
|
137
|
+
- **It needs exchangeable runs.** With a random or quasi-random sampler the bound holds. With
|
|
138
|
+
TPE, Gaussian-process or CMA-ES samplers it does not. The Optuna pruner warns once and carries
|
|
139
|
+
on, but do not rely on the bound there.
|
|
140
|
+
- **It bounds the joint rate rather than the share of good runs killed.** When good runs are rare, a large
|
|
141
|
+
share of them can still die: about a third in the measurement above, mostly borderline ones.
|
|
142
|
+
Lower α if that matters.
|
|
143
|
+
- **An estimated τ gives a measured bound rather than a proven one.** `good_quantile` estimates τ from the
|
|
144
|
+
calibration runs. Simulations put the realized rate at 0.88 to 0.97 times α. Pass a real
|
|
145
|
+
`good_threshold` when you have one.
|
|
146
|
+
- **α = 0.05 needs at least 19 calibration runs**, and the certified threshold gets loose with
|
|
147
|
+
few good runs. Sweeps of a few dozen trials are too small for this.
|
|
148
|
+
- **One setting only so far.** One model, one dataset, one horizon.
|
|
149
|
+
|
|
150
|
+
## Layout
|
|
151
|
+
|
|
152
|
+
| path | what |
|
|
153
|
+
|---|---|
|
|
154
|
+
| `src/surestop/` | `KillRule` (numpy/scipy) and `ConformalPruner` (Optuna) |
|
|
155
|
+
| `tests/` | the bound holds on synthetic exchangeable runs, and the same assertion fails on a rule that leaks 2× its budget |
|
|
156
|
+
| `reproduce/` | the 400 curves, the pre-registered splits and targets, the script, and its committed output |
|
|
157
|
+
| `docs/` | the write-up, the prior-art sweep, and the frozen hyperparameter prior the curves were drawn from |
|
|
158
|
+
| `optunahub/` | the package as an OptunaHub registry entry |
|
|
159
|
+
|
|
160
|
+
## Related work
|
|
161
|
+
|
|
162
|
+
Calibrating a stop threshold on a predictor with a conformal guarantee is an established pattern. Xie et al.
|
|
163
|
+
(arXiv:2602.13935) use the same device, a threshold on the maximum of a running score calibrated
|
|
164
|
+
on one class only, to stop LLM reasoning traces. Related work does the same for mixed-integer
|
|
165
|
+
solvers and LLM agent episodes. As far as a search found, nobody has applied it to pruning
|
|
166
|
+
training runs, and no major tuning library ships a pruner with a stated error guarantee.
|
|
167
|
+
`docs/prior-art.md` has the sweep.
|
|
168
|
+
|
|
169
|
+
## Citing
|
|
170
|
+
|
|
171
|
+
Concept DOI: [10.5281/zenodo.22726440](https://doi.org/10.5281/zenodo.22726440), which always resolves
|
|
172
|
+
to the newest version. `CITATION.cff` has the full entry. The curves are a measurement, released
|
|
173
|
+
under the same MIT licence as the code.
|
surestop-0.2.0/README.md
ADDED
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
# SureStop
|
|
2
|
+
|
|
3
|
+
**Stop bad training runs early, with a stated bound on how often you stop one that would have
|
|
4
|
+
ended well.** You pick α. The rule kills runs so that the expected fraction of runs that were
|
|
5
|
+
good *and* got killed stays at or below α.
|
|
6
|
+
|
|
7
|
+
Every pruner in common use (median, percentile, successive halving, Hyperband) is a heuristic,
|
|
8
|
+
and none says how often it kills a run that would have finished well. This one does, under a
|
|
9
|
+
condition you can check: the runs you calibrate on and the runs you judge must come from the same
|
|
10
|
+
random process.
|
|
11
|
+
|
|
12
|
+
It was measured on 400 real LoRA fine-tuning runs under a pre-registered protocol. The curves are
|
|
13
|
+
in this repository, and a script checks that the package reproduces the recorded numbers exactly.
|
|
14
|
+
|
|
15
|
+
## Install
|
|
16
|
+
|
|
17
|
+
```
|
|
18
|
+
pip install surestop # numpy + scipy
|
|
19
|
+
pip install "surestop[optuna]" # adds the Optuna pruner
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
## Use
|
|
23
|
+
|
|
24
|
+
Plain numpy, on any training loop:
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
from surestop import KillRule
|
|
28
|
+
|
|
29
|
+
# completed_curves: array of shape (runs, evaluations), lower is better
|
|
30
|
+
rule = KillRule(alpha=0.05, good_threshold=2.16, min_peek=17).fit(completed_curves)
|
|
31
|
+
|
|
32
|
+
rule.should_kill(curve_so_far) # True -> stop this run now
|
|
33
|
+
rule.kill_step(full_curve) # where it would have fired, or None
|
|
34
|
+
rule.evaluate(held_out_curves) # false kills, joint rate, savings
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
With Optuna, the first `n_calibration` trials run unpruned to calibrate, then the rule freezes:
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
import optuna
|
|
41
|
+
from surestop import ConformalPruner
|
|
42
|
+
|
|
43
|
+
study = optuna.create_study(sampler=optuna.samplers.RandomSampler(),
|
|
44
|
+
pruner=ConformalPruner(n_calibration=60, alpha=0.05))
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Pass `maximize=True` (or a `direction="maximize"` study) when higher is better.
|
|
48
|
+
|
|
49
|
+
## The rule
|
|
50
|
+
|
|
51
|
+
A run is *good* if its final value ends at or below a threshold τ. At every evaluation, a
|
|
52
|
+
predictor forecasts the final value from the curve so far. The default predictor is the latest
|
|
53
|
+
observed value.
|
|
54
|
+
|
|
55
|
+
To calibrate, take n runs that have already finished. For each good one, record its worst moment:
|
|
56
|
+
the largest gap between forecast and τ at any evaluation. Runs that are not good contribute
|
|
57
|
+
nothing. Set λ to the ⌈(1−α)(n+1)⌉-th smallest of these scores. A new run is killed at the first
|
|
58
|
+
evaluation where its forecast exceeds τ + λ.
|
|
59
|
+
|
|
60
|
+
This is conformal risk control on the loss "good and killed". Taking the worst moment across all
|
|
61
|
+
evaluations means one threshold covers every look at a run, so there is no correction for
|
|
62
|
+
multiple looks. The guarantee needs the calibration runs and the new runs to be exchangeable.
|
|
63
|
+
Random sampling gives that. Adaptive samplers such as TPE do not.
|
|
64
|
+
|
|
65
|
+
## What was measured
|
|
66
|
+
|
|
67
|
+
Experiment CC-03c, pre-registered before any run finished:
|
|
68
|
+
|
|
69
|
+
- **Runs:** 400 LoRA fine-tunes of Qwen2.5-0.5B on dolly-15k, 750 steps each, validation loss
|
|
70
|
+
measured 50 times per run.
|
|
71
|
+
- **Configs:** i.i.d. draws from a frozen hyperparameter prior (`docs/hyperparameter-prior.md`).
|
|
72
|
+
- **τ:** pinned from an earlier pool before the first run landed.
|
|
73
|
+
- **Splits:** two, one random and one past-to-future by night, each about half calibration and
|
|
74
|
+
half test.
|
|
75
|
+
|
|
76
|
+
The pre-registered rule used a fitted power-law predictor. On both splits it kept the bound and
|
|
77
|
+
saved about half the evaluation compute:
|
|
78
|
+
|
|
79
|
+
| split | false kills | joint rate (95% CI) | eval compute saved |
|
|
80
|
+
|---|---|---|---|
|
|
81
|
+
| random | 13 of 200 | 6.5% (3.5–10.9%) | 52.2% |
|
|
82
|
+
| past-to-future | 11 of 203 | 5.4% (2.7–9.5%) | 52.8% |
|
|
83
|
+
|
|
84
|
+
A rule that kills every run immediately was scored by the same gate and failed on both splits.
|
|
85
|
+
|
|
86
|
+
**The flaw that run found:** every kill fired at the 3rd of 50 evaluations. Early rankings only
|
|
87
|
+
match final rankings from about evaluation 18. The rule stayed within budget because the bound
|
|
88
|
+
held. The fitted power law ranked runs worse than the latest observed value at every early
|
|
89
|
+
evaluation.
|
|
90
|
+
|
|
91
|
+
**What this package ships instead,** and what the reproduction script checks:
|
|
92
|
+
|
|
93
|
+
| predictor | hold until | random split | past-to-future split | 2000 resplits |
|
|
94
|
+
|---|---|---|---|---|
|
|
95
|
+
| last value | none | 18 of 200 (9.0%), 78% saved | 7 of 203 (3.4%), 75% saved | joint 5.0%, breach 8.5% |
|
|
96
|
+
| last value | eval 18 | 18 of 200 (9.0%), 57% saved | 7 of 203 (3.4%), 56% saved | joint 5.0%, breach 8.6% |
|
|
97
|
+
|
|
98
|
+
The random split's 18 of 200 is a breach. It sits at the 95th percentile of that rule's own
|
|
99
|
+
resplit distribution, whose mean is 5.0%, so it reads as a split tail rather than a leak. A
|
|
100
|
+
correct rule at this sample size lands on a breach in about 7% of splits. The held version was
|
|
101
|
+
chosen *after* seeing these 400 runs, so its numbers are a hypothesis. It halves
|
|
102
|
+
how often the single best run is killed at the same budget. Its confirmation needs a fresh pool.
|
|
103
|
+
|
|
104
|
+
Full account: `docs/writeup.md`. Run the check yourself:
|
|
105
|
+
|
|
106
|
+
```
|
|
107
|
+
python reproduce/reproduce_cc03c.py
|
|
108
|
+
```
|
|
109
|
+
|
|
110
|
+
## What the bound does not cover
|
|
111
|
+
|
|
112
|
+
- **It needs exchangeable runs.** With a random or quasi-random sampler the bound holds. With
|
|
113
|
+
TPE, Gaussian-process or CMA-ES samplers it does not. The Optuna pruner warns once and carries
|
|
114
|
+
on, but do not rely on the bound there.
|
|
115
|
+
- **It bounds the joint rate rather than the share of good runs killed.** When good runs are rare, a large
|
|
116
|
+
share of them can still die: about a third in the measurement above, mostly borderline ones.
|
|
117
|
+
Lower α if that matters.
|
|
118
|
+
- **An estimated τ gives a measured bound rather than a proven one.** `good_quantile` estimates τ from the
|
|
119
|
+
calibration runs. Simulations put the realized rate at 0.88 to 0.97 times α. Pass a real
|
|
120
|
+
`good_threshold` when you have one.
|
|
121
|
+
- **α = 0.05 needs at least 19 calibration runs**, and the certified threshold gets loose with
|
|
122
|
+
few good runs. Sweeps of a few dozen trials are too small for this.
|
|
123
|
+
- **One setting only so far.** One model, one dataset, one horizon.
|
|
124
|
+
|
|
125
|
+
## Layout
|
|
126
|
+
|
|
127
|
+
| path | what |
|
|
128
|
+
|---|---|
|
|
129
|
+
| `src/surestop/` | `KillRule` (numpy/scipy) and `ConformalPruner` (Optuna) |
|
|
130
|
+
| `tests/` | the bound holds on synthetic exchangeable runs, and the same assertion fails on a rule that leaks 2× its budget |
|
|
131
|
+
| `reproduce/` | the 400 curves, the pre-registered splits and targets, the script, and its committed output |
|
|
132
|
+
| `docs/` | the write-up, the prior-art sweep, and the frozen hyperparameter prior the curves were drawn from |
|
|
133
|
+
| `optunahub/` | the package as an OptunaHub registry entry |
|
|
134
|
+
|
|
135
|
+
## Related work
|
|
136
|
+
|
|
137
|
+
Calibrating a stop threshold on a predictor with a conformal guarantee is an established pattern. Xie et al.
|
|
138
|
+
(arXiv:2602.13935) use the same device, a threshold on the maximum of a running score calibrated
|
|
139
|
+
on one class only, to stop LLM reasoning traces. Related work does the same for mixed-integer
|
|
140
|
+
solvers and LLM agent episodes. As far as a search found, nobody has applied it to pruning
|
|
141
|
+
training runs, and no major tuning library ships a pruner with a stated error guarantee.
|
|
142
|
+
`docs/prior-art.md` has the sweep.
|
|
143
|
+
|
|
144
|
+
## Citing
|
|
145
|
+
|
|
146
|
+
Concept DOI: [10.5281/zenodo.22726440](https://doi.org/10.5281/zenodo.22726440), which always resolves
|
|
147
|
+
to the newest version. `CITATION.cff` has the full entry. The curves are a measurement, released
|
|
148
|
+
under the same MIT licence as the code.
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "surestop"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "Stop bad runs early, with a conformal bound on how often you stop one that would have ended well"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
authors = [{ name = "Adam Noonan" }]
|
|
13
|
+
requires-python = ">=3.10"
|
|
14
|
+
dependencies = ["numpy>=1.24", "scipy>=1.10"]
|
|
15
|
+
keywords = ["hyperparameter-optimization", "early-stopping", "pruning", "conformal", "optuna"]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Development Status :: 3 - Alpha",
|
|
18
|
+
"Intended Audience :: Science/Research",
|
|
19
|
+
"Programming Language :: Python :: 3",
|
|
20
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
[project.optional-dependencies]
|
|
24
|
+
optuna = ["optuna>=4"]
|
|
25
|
+
test = ["pytest", "optuna>=4"]
|
|
26
|
+
|
|
27
|
+
[project.urls]
|
|
28
|
+
Homepage = "https://github.com/ACNoonan/surestop"
|
|
29
|
+
Issues = "https://github.com/ACNoonan/surestop/issues"
|
|
30
|
+
|
|
31
|
+
[tool.setuptools.packages.find]
|
|
32
|
+
where = ["src"]
|
|
33
|
+
include = ["surestop*"]
|
|
34
|
+
|
|
35
|
+
[tool.pytest.ini_options]
|
|
36
|
+
testpaths = ["tests"]
|
surestop-0.2.0/setup.cfg
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""surestop — stop bad runs early, with a bound on how often you stop one that would have ended well."""
|
|
2
|
+
|
|
3
|
+
from .core import KillRule, NotEnoughCalibration, last_value
|
|
4
|
+
|
|
5
|
+
__version__ = "0.2.0"
|
|
6
|
+
__all__ = ["KillRule", "NotEnoughCalibration", "last_value", "ConformalPruner"]
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
from typing import Any
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def __getattr__(name: str) -> Any:
|
|
13
|
+
# Optuna is optional: import the pruner only when someone asks for it.
|
|
14
|
+
if name == "ConformalPruner":
|
|
15
|
+
from .optuna_pruner import ConformalPruner
|
|
16
|
+
return ConformalPruner
|
|
17
|
+
raise AttributeError(name)
|
|
@@ -0,0 +1,201 @@
|
|
|
1
|
+
"""A kill rule for training runs that bounds how often it kills a run that would have ended well.
|
|
2
|
+
|
|
3
|
+
The construction is the one CC-03c pre-registered and measured on 400 real LoRA runs
|
|
4
|
+
(`experiments/2026-08-17-CC-03c-validity-gate-that-can-fire/`):
|
|
5
|
+
|
|
6
|
+
1. Calibrate on n **completed** runs, each a learning curve on a shared step grid.
|
|
7
|
+
2. A run is *good* if its final value is at or below a threshold τ (lower is better).
|
|
8
|
+
3. A predictor forecasts the final value from the curve so far, at every step.
|
|
9
|
+
4. Each good calibration run scores max over eligible steps of (prediction − τ).
|
|
10
|
+
Runs that are not good score −∞.
|
|
11
|
+
5. λ̂ is the ⌈(1−α)(n+1)⌉-th smallest score. This is conformal risk control on the loss
|
|
12
|
+
"this run is good AND the rule kills it".
|
|
13
|
+
6. A new run is killed at the first eligible step where prediction − λ̂ > τ.
|
|
14
|
+
|
|
15
|
+
**Guarantee.** If the calibration runs and the new runs are exchangeable — for example, drawn
|
|
16
|
+
i.i.d. by the same random sampler — then the expected fraction of new runs that are good and
|
|
17
|
+
killed is at most α. One kill decision covers every step at once, with no multiplicity
|
|
18
|
+
correction.
|
|
19
|
+
|
|
20
|
+
**What it does not promise.**
|
|
21
|
+
- **It bounds the joint rate, not the share of good runs killed.** With a small good fraction,
|
|
22
|
+
the conditional rate can be large: ~30% on CC-03c.
|
|
23
|
+
- **It needs exchangeability.** An adaptive sampler such as TPE shifts later runs away from the
|
|
24
|
+
calibration runs, and the guarantee no longer applies.
|
|
25
|
+
- **With τ estimated from the calibration finals** (`good_quantile`), the bound is measured, not
|
|
26
|
+
proven: 0.88–0.97× α in CC-03c's simulations. Pass `good_threshold` when you have a real target.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
from __future__ import annotations
|
|
30
|
+
|
|
31
|
+
import math
|
|
32
|
+
import warnings
|
|
33
|
+
from dataclasses import dataclass, field
|
|
34
|
+
from typing import Any, Callable
|
|
35
|
+
|
|
36
|
+
import numpy as np
|
|
37
|
+
from numpy.typing import ArrayLike
|
|
38
|
+
from scipy.stats import rankdata
|
|
39
|
+
|
|
40
|
+
Predictor = Callable[[np.ndarray], np.ndarray]
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def last_value(curves: ArrayLike) -> np.ndarray:
|
|
44
|
+
"""Forecast the final value as the latest observed value.
|
|
45
|
+
|
|
46
|
+
A predictor maps curves of shape (n, T) to forecasts of shape (n, T) and must be causal:
|
|
47
|
+
column t may use only columns 0..t. On CC-03c's pool this beat a fitted pow3 curve at every
|
|
48
|
+
early step.
|
|
49
|
+
"""
|
|
50
|
+
return np.asarray(curves, dtype=float)
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
class NotEnoughCalibration(ValueError):
|
|
54
|
+
"""Too few calibration runs for any threshold to certify at this α."""
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _spearman_cols(M: np.ndarray, y: np.ndarray) -> np.ndarray:
|
|
58
|
+
rm = rankdata(M, axis=0)
|
|
59
|
+
rm -= rm.mean(axis=0)
|
|
60
|
+
ry = rankdata(y)
|
|
61
|
+
ry -= ry.mean()
|
|
62
|
+
denom = np.sqrt((rm ** 2).sum(axis=0) * (ry ** 2).sum())
|
|
63
|
+
with np.errstate(invalid="ignore", divide="ignore"):
|
|
64
|
+
return np.where(denom > 0, (rm * ry[:, None]).sum(axis=0) / denom, 0.0)
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
@dataclass
|
|
68
|
+
class KillRule:
|
|
69
|
+
"""Certified early-kill rule. Fit on completed curves, then ask `should_kill` on partial ones.
|
|
70
|
+
|
|
71
|
+
Parameters
|
|
72
|
+
----------
|
|
73
|
+
alpha : bound on the expected fraction of runs that are good and killed.
|
|
74
|
+
good_threshold : a run is good if its final value is at or below this. Overrides `good_quantile`.
|
|
75
|
+
good_quantile : if no threshold is given, τ is this quantile of the calibration finals.
|
|
76
|
+
min_peek : index of the earliest step at which a kill may fire.
|
|
77
|
+
hold_until_resolved : if set (e.g. 0.90), also hold kills until the first step at or after
|
|
78
|
+
`min_peek` where Spearman(prediction, final) on the calibration runs reaches this value.
|
|
79
|
+
Estimated per calibration set, so it varies between sets; pin `min_peek` instead when a
|
|
80
|
+
prior pool tells you where the ranking settles.
|
|
81
|
+
maximize : set True when higher values are better (accuracy, reward).
|
|
82
|
+
predictor : causal forecaster of the final value. Default: last observed value.
|
|
83
|
+
"""
|
|
84
|
+
|
|
85
|
+
alpha: float = 0.05
|
|
86
|
+
good_threshold: float | None = None
|
|
87
|
+
good_quantile: float = 0.20
|
|
88
|
+
min_peek: int = 0
|
|
89
|
+
hold_until_resolved: float | None = None
|
|
90
|
+
maximize: bool = False
|
|
91
|
+
predictor: Predictor = last_value
|
|
92
|
+
|
|
93
|
+
tau_: float = field(init=False, default=math.nan)
|
|
94
|
+
lambda_: float = field(init=False, default=math.nan)
|
|
95
|
+
first_peek_: int = field(init=False, default=-1)
|
|
96
|
+
n_steps_: int = field(init=False, default=0)
|
|
97
|
+
n_calibration_: int = field(init=False, default=0)
|
|
98
|
+
n_good_: int = field(init=False, default=0)
|
|
99
|
+
resolution_: np.ndarray | None = field(init=False, default=None)
|
|
100
|
+
|
|
101
|
+
# ── fitting ────────────────────────────────────────────────────────────────
|
|
102
|
+
def _oriented(self, curves: ArrayLike) -> np.ndarray:
|
|
103
|
+
Y = np.asarray(curves, dtype=float)
|
|
104
|
+
return -Y if self.maximize else Y
|
|
105
|
+
|
|
106
|
+
def fit(self, curves: ArrayLike) -> "KillRule":
|
|
107
|
+
Y = self._oriented(curves)
|
|
108
|
+
if Y.ndim != 2 or Y.shape[1] < 2:
|
|
109
|
+
raise ValueError("curves must be a 2-D array (runs × steps) with at least 2 steps")
|
|
110
|
+
if not np.all(np.isfinite(Y)):
|
|
111
|
+
raise ValueError("calibration curves must be finite; drop incomplete runs first")
|
|
112
|
+
if not 0 < self.alpha < 1:
|
|
113
|
+
raise ValueError("alpha must lie in (0, 1)")
|
|
114
|
+
n, T = Y.shape
|
|
115
|
+
order = math.ceil((1 - self.alpha) * (n + 1))
|
|
116
|
+
if order > n:
|
|
117
|
+
raise NotEnoughCalibration(
|
|
118
|
+
f"alpha={self.alpha} needs at least {math.ceil(1 / self.alpha) - 1} calibration "
|
|
119
|
+
f"runs for any threshold to certify; got {n}")
|
|
120
|
+
|
|
121
|
+
finals = Y[:, -1]
|
|
122
|
+
if self.good_threshold is not None:
|
|
123
|
+
tau = -self.good_threshold if self.maximize else float(self.good_threshold)
|
|
124
|
+
else:
|
|
125
|
+
tau = float(np.sort(finals)[math.ceil(self.good_quantile * n) - 1])
|
|
126
|
+
P = np.asarray(self.predictor(Y), dtype=float)
|
|
127
|
+
if P.shape != Y.shape:
|
|
128
|
+
raise ValueError(f"predictor returned shape {P.shape}, expected {Y.shape}")
|
|
129
|
+
|
|
130
|
+
last_eligible = T - 2 # a kill at the final step saves nothing
|
|
131
|
+
j0 = self.min_peek
|
|
132
|
+
if self.hold_until_resolved is not None:
|
|
133
|
+
rho = _spearman_cols(P[:, : T - 1], finals)
|
|
134
|
+
self.resolution_ = rho
|
|
135
|
+
hit = np.flatnonzero(rho[self.min_peek:] >= self.hold_until_resolved)
|
|
136
|
+
j0 = self.min_peek + int(hit[0]) if hit.size else T - 1
|
|
137
|
+
|
|
138
|
+
good = finals <= tau
|
|
139
|
+
if j0 <= last_eligible:
|
|
140
|
+
scores = np.where(good, P[:, j0: last_eligible + 1].max(axis=1) - tau, -np.inf)
|
|
141
|
+
else:
|
|
142
|
+
scores = np.full(n, -np.inf)
|
|
143
|
+
lam = float(np.sort(scores)[order - 1])
|
|
144
|
+
if lam == -np.inf and j0 <= last_eligible:
|
|
145
|
+
warnings.warn(
|
|
146
|
+
f"only {int(good.sum())} good calibration runs: at alpha={self.alpha} no finite "
|
|
147
|
+
"threshold certifies, so the rule kills every run at its first eligible step. "
|
|
148
|
+
"Raise good_quantile or add calibration runs.", RuntimeWarning, stacklevel=2)
|
|
149
|
+
|
|
150
|
+
self.tau_, self.lambda_, self.first_peek_ = tau, lam, j0
|
|
151
|
+
self.n_steps_, self.n_calibration_, self.n_good_ = T, n, int(good.sum())
|
|
152
|
+
return self
|
|
153
|
+
|
|
154
|
+
# ── deciding ───────────────────────────────────────────────────────────────
|
|
155
|
+
def _check_fitted(self) -> None:
|
|
156
|
+
if self.n_steps_ == 0:
|
|
157
|
+
raise RuntimeError("call fit() first")
|
|
158
|
+
|
|
159
|
+
def should_kill(self, partial_curve: ArrayLike) -> bool:
|
|
160
|
+
"""True if a run whose curve so far is `partial_curve` should be killed at its latest step."""
|
|
161
|
+
self._check_fitted()
|
|
162
|
+
y = self._oriented(partial_curve)
|
|
163
|
+
t = y.shape[0] - 1
|
|
164
|
+
if t < self.first_peek_ or t > self.n_steps_ - 2:
|
|
165
|
+
return False
|
|
166
|
+
pred = float(np.asarray(self.predictor(y[None, :]))[0, t])
|
|
167
|
+
return bool(pred - self.lambda_ > self.tau_)
|
|
168
|
+
|
|
169
|
+
def kill_step(self, curve: ArrayLike) -> int | None:
|
|
170
|
+
"""Index of the first step at which the rule would kill this curve, or None."""
|
|
171
|
+
self._check_fitted()
|
|
172
|
+
Y = self._oriented(curve)[None, :]
|
|
173
|
+
P = np.asarray(self.predictor(Y), dtype=float)[0]
|
|
174
|
+
lo, hi = self.first_peek_, min(self.n_steps_, Y.shape[1]) - 2
|
|
175
|
+
if lo > hi:
|
|
176
|
+
return None
|
|
177
|
+
fire = np.flatnonzero(P[lo: hi + 1] - self.lambda_ > self.tau_)
|
|
178
|
+
return lo + int(fire[0]) if fire.size else None
|
|
179
|
+
|
|
180
|
+
def evaluate(self, curves: ArrayLike) -> dict[str, Any]:
|
|
181
|
+
"""Score the rule on held-out COMPLETED curves: false kills, joint rate, savings."""
|
|
182
|
+
self._check_fitted()
|
|
183
|
+
Y = self._oriented(curves)
|
|
184
|
+
n, T = Y.shape
|
|
185
|
+
P = np.asarray(self.predictor(Y), dtype=float)
|
|
186
|
+
lo, hi = self.first_peek_, T - 2
|
|
187
|
+
if lo <= hi:
|
|
188
|
+
fire = P[:, lo: hi + 1] - self.lambda_ > self.tau_
|
|
189
|
+
killed = fire.any(axis=1)
|
|
190
|
+
step = np.where(killed, lo + np.argmax(fire, axis=1), T - 1)
|
|
191
|
+
else:
|
|
192
|
+
killed, step = np.zeros(n, bool), np.full(n, T - 1)
|
|
193
|
+
good = Y[:, -1] <= self.tau_
|
|
194
|
+
fk = int((killed & good).sum())
|
|
195
|
+
return {
|
|
196
|
+
"n": n, "n_good": int(good.sum()), "n_killed": int(killed.sum()),
|
|
197
|
+
"false_kills": fk, "joint_false_kill_rate": fk / n,
|
|
198
|
+
"conditional_false_kill_rate": fk / int(good.sum()) if good.any() else None,
|
|
199
|
+
"savings": float(1 - (step + 1).sum() / (T * n)),
|
|
200
|
+
"median_kill_step": float(np.median(step[killed])) if killed.any() else None,
|
|
201
|
+
}
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
"""Optuna integration: a pruner that runs in shadow mode first, then kills with a certified bound.
|
|
2
|
+
|
|
3
|
+
pruner = ConformalPruner(n_calibration=60, alpha=0.05)
|
|
4
|
+
study = optuna.create_study(sampler=optuna.samplers.RandomSampler(), pruner=pruner)
|
|
5
|
+
|
|
6
|
+
**How it behaves.**
|
|
7
|
+
- **Shadow mode first.** The first `n_calibration` completed trials, by trial number, run
|
|
8
|
+
unpruned. Their reported curves calibrate the rule once.
|
|
9
|
+
- **Then frozen.** The rule never recalibrates on later trials: those are survivors of the
|
|
10
|
+
rule's own kills, and fitting on them would bias it.
|
|
11
|
+
- **The final outcome is the last reported intermediate value,** not `trial.value`.
|
|
12
|
+
- **Every calibration trial must report on the same step grid.**
|
|
13
|
+
|
|
14
|
+
**The guarantee needs the later trials to be exchangeable with the calibration trials.** With
|
|
15
|
+
`RandomSampler`, `QMCSampler` or a shuffled `GridSampler` they are. With `TPESampler`, `GPSampler`
|
|
16
|
+
or `CmaEsSampler` they are not, so the pruner warns once and the bound should not be relied on.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
import warnings
|
|
22
|
+
from typing import Any
|
|
23
|
+
|
|
24
|
+
import numpy as np
|
|
25
|
+
import optuna
|
|
26
|
+
from optuna.pruners import BasePruner
|
|
27
|
+
from optuna.study import StudyDirection
|
|
28
|
+
from optuna.trial import TrialState
|
|
29
|
+
|
|
30
|
+
from .core import KillRule, Predictor, last_value
|
|
31
|
+
|
|
32
|
+
ArrayLike = Any
|
|
33
|
+
|
|
34
|
+
_EXCHANGEABLE = (optuna.samplers.RandomSampler, optuna.samplers.QMCSampler,
|
|
35
|
+
optuna.samplers.GridSampler, optuna.samplers.BruteForceSampler)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
class ConformalPruner(BasePruner):
|
|
39
|
+
def __init__(self, n_calibration: int = 60, alpha: float = 0.05, *,
|
|
40
|
+
good_threshold: float | None = None, good_quantile: float = 0.20,
|
|
41
|
+
min_peek: int = 0, hold_until_resolved: float | None = None,
|
|
42
|
+
predictor: Predictor = last_value) -> None:
|
|
43
|
+
self.n_calibration = n_calibration
|
|
44
|
+
self._kw: dict[str, Any] = dict(alpha=alpha, good_threshold=good_threshold, good_quantile=good_quantile,
|
|
45
|
+
min_peek=min_peek, hold_until_resolved=hold_until_resolved,
|
|
46
|
+
predictor=predictor)
|
|
47
|
+
self.rule_: KillRule | None = None
|
|
48
|
+
self.steps_: list[int] | None = None
|
|
49
|
+
self.calibration_trials_: list[int] = []
|
|
50
|
+
self._warned = False
|
|
51
|
+
|
|
52
|
+
@classmethod
|
|
53
|
+
def from_curves(cls, curves: ArrayLike, steps: Any, *, maximize: bool = False,
|
|
54
|
+
**kw: Any) -> "ConformalPruner":
|
|
55
|
+
"""A pruner already calibrated offline, e.g. on a previous sweep's completed curves."""
|
|
56
|
+
p = cls(n_calibration=len(curves), **kw)
|
|
57
|
+
p.rule_ = KillRule(maximize=maximize, **p._kw).fit(curves)
|
|
58
|
+
p.steps_ = list(steps)
|
|
59
|
+
return p
|
|
60
|
+
|
|
61
|
+
def _calibrate(self, study: optuna.study.Study) -> bool:
|
|
62
|
+
done = sorted((t for t in study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,))
|
|
63
|
+
if t.intermediate_values), key=lambda t: t.number)[: self.n_calibration]
|
|
64
|
+
if len(done) < self.n_calibration:
|
|
65
|
+
return False
|
|
66
|
+
grid = sorted(done[0].intermediate_values)
|
|
67
|
+
rows = [[t.intermediate_values[s] for s in grid] for t in done
|
|
68
|
+
if sorted(t.intermediate_values) == grid]
|
|
69
|
+
if len(rows) < len(done):
|
|
70
|
+
warnings.warn(f"{len(done) - len(rows)} calibration trials reported a different step "
|
|
71
|
+
"grid and were skipped", RuntimeWarning, stacklevel=3)
|
|
72
|
+
return False
|
|
73
|
+
maximize = study.direction == StudyDirection.MAXIMIZE
|
|
74
|
+
self.rule_ = KillRule(maximize=maximize, **self._kw).fit(np.asarray(rows))
|
|
75
|
+
self.steps_ = grid
|
|
76
|
+
self.calibration_trials_ = [t.number for t in done]
|
|
77
|
+
return True
|
|
78
|
+
|
|
79
|
+
def prune(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> bool:
|
|
80
|
+
if not self._warned and not isinstance(study.sampler, _EXCHANGEABLE):
|
|
81
|
+
warnings.warn(f"{type(study.sampler).__name__} adapts to earlier trials, so later trials "
|
|
82
|
+
"are not exchangeable with the calibration trials and the false-kill "
|
|
83
|
+
"bound does not hold. Use RandomSampler or QMCSampler for a certified "
|
|
84
|
+
"bound.", RuntimeWarning, stacklevel=2)
|
|
85
|
+
self._warned = True
|
|
86
|
+
step = trial.last_step
|
|
87
|
+
if step is None:
|
|
88
|
+
return False
|
|
89
|
+
if self.rule_ is None and not self._calibrate(study):
|
|
90
|
+
return False
|
|
91
|
+
assert self.rule_ is not None and self.steps_ is not None
|
|
92
|
+
if trial.number in self.calibration_trials_ or step not in self.steps_:
|
|
93
|
+
return False
|
|
94
|
+
idx = self.steps_.index(step)
|
|
95
|
+
values = trial.intermediate_values
|
|
96
|
+
prefix = [values.get(s) for s in self.steps_[: idx + 1]]
|
|
97
|
+
if any(v is None for v in prefix):
|
|
98
|
+
return False
|
|
99
|
+
return self.rule_.should_kill(np.asarray(prefix, dtype=float))
|
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: surestop
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Stop bad runs early, with a conformal bound on how often you stop one that would have ended well
|
|
5
|
+
Author: Adam Noonan
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/ACNoonan/surestop
|
|
8
|
+
Project-URL: Issues, https://github.com/ACNoonan/surestop/issues
|
|
9
|
+
Keywords: hyperparameter-optimization,early-stopping,pruning,conformal,optuna
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Intended Audience :: Science/Research
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
14
|
+
Requires-Python: >=3.10
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
License-File: LICENSE
|
|
17
|
+
Requires-Dist: numpy>=1.24
|
|
18
|
+
Requires-Dist: scipy>=1.10
|
|
19
|
+
Provides-Extra: optuna
|
|
20
|
+
Requires-Dist: optuna>=4; extra == "optuna"
|
|
21
|
+
Provides-Extra: test
|
|
22
|
+
Requires-Dist: pytest; extra == "test"
|
|
23
|
+
Requires-Dist: optuna>=4; extra == "test"
|
|
24
|
+
Dynamic: license-file
|
|
25
|
+
|
|
26
|
+
# SureStop
|
|
27
|
+
|
|
28
|
+
**Stop bad training runs early, with a stated bound on how often you stop one that would have
|
|
29
|
+
ended well.** You pick α. The rule kills runs so that the expected fraction of runs that were
|
|
30
|
+
good *and* got killed stays at or below α.
|
|
31
|
+
|
|
32
|
+
Every pruner in common use (median, percentile, successive halving, Hyperband) is a heuristic,
|
|
33
|
+
and none says how often it kills a run that would have finished well. This one does, under a
|
|
34
|
+
condition you can check: the runs you calibrate on and the runs you judge must come from the same
|
|
35
|
+
random process.
|
|
36
|
+
|
|
37
|
+
It was measured on 400 real LoRA fine-tuning runs under a pre-registered protocol. The curves are
|
|
38
|
+
in this repository, and a script checks that the package reproduces the recorded numbers exactly.
|
|
39
|
+
|
|
40
|
+
## Install
|
|
41
|
+
|
|
42
|
+
```
|
|
43
|
+
pip install surestop # numpy + scipy
|
|
44
|
+
pip install "surestop[optuna]" # adds the Optuna pruner
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
## Use
|
|
48
|
+
|
|
49
|
+
Plain numpy, on any training loop:
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from surestop import KillRule
|
|
53
|
+
|
|
54
|
+
# completed_curves: array of shape (runs, evaluations), lower is better
|
|
55
|
+
rule = KillRule(alpha=0.05, good_threshold=2.16, min_peek=17).fit(completed_curves)
|
|
56
|
+
|
|
57
|
+
rule.should_kill(curve_so_far) # True -> stop this run now
|
|
58
|
+
rule.kill_step(full_curve) # where it would have fired, or None
|
|
59
|
+
rule.evaluate(held_out_curves) # false kills, joint rate, savings
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
With Optuna, the first `n_calibration` trials run unpruned to calibrate, then the rule freezes:
|
|
63
|
+
|
|
64
|
+
```python
|
|
65
|
+
import optuna
|
|
66
|
+
from surestop import ConformalPruner
|
|
67
|
+
|
|
68
|
+
study = optuna.create_study(sampler=optuna.samplers.RandomSampler(),
|
|
69
|
+
pruner=ConformalPruner(n_calibration=60, alpha=0.05))
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Pass `maximize=True` (or a `direction="maximize"` study) when higher is better.
|
|
73
|
+
|
|
74
|
+
## The rule
|
|
75
|
+
|
|
76
|
+
A run is *good* if its final value ends at or below a threshold τ. At every evaluation, a
|
|
77
|
+
predictor forecasts the final value from the curve so far. The default predictor is the latest
|
|
78
|
+
observed value.
|
|
79
|
+
|
|
80
|
+
To calibrate, take n runs that have already finished. For each good one, record its worst moment:
|
|
81
|
+
the largest gap between forecast and τ at any evaluation. Runs that are not good contribute
|
|
82
|
+
nothing. Set λ to the ⌈(1−α)(n+1)⌉-th smallest of these scores. A new run is killed at the first
|
|
83
|
+
evaluation where its forecast exceeds τ + λ.
|
|
84
|
+
|
|
85
|
+
This is conformal risk control on the loss "good and killed". Taking the worst moment across all
|
|
86
|
+
evaluations means one threshold covers every look at a run, so there is no correction for
|
|
87
|
+
multiple looks. The guarantee needs the calibration runs and the new runs to be exchangeable.
|
|
88
|
+
Random sampling gives that. Adaptive samplers such as TPE do not.
|
|
89
|
+
|
|
90
|
+
## What was measured
|
|
91
|
+
|
|
92
|
+
Experiment CC-03c, pre-registered before any run finished:
|
|
93
|
+
|
|
94
|
+
- **Runs:** 400 LoRA fine-tunes of Qwen2.5-0.5B on dolly-15k, 750 steps each, validation loss
|
|
95
|
+
measured 50 times per run.
|
|
96
|
+
- **Configs:** i.i.d. draws from a frozen hyperparameter prior (`docs/hyperparameter-prior.md`).
|
|
97
|
+
- **τ:** pinned from an earlier pool before the first run landed.
|
|
98
|
+
- **Splits:** two, one random and one past-to-future by night, each about half calibration and
|
|
99
|
+
half test.
|
|
100
|
+
|
|
101
|
+
The pre-registered rule used a fitted power-law predictor. On both splits it kept the bound and
|
|
102
|
+
saved about half the evaluation compute:
|
|
103
|
+
|
|
104
|
+
| split | false kills | joint rate (95% CI) | eval compute saved |
|
|
105
|
+
|---|---|---|---|
|
|
106
|
+
| random | 13 of 200 | 6.5% (3.5–10.9%) | 52.2% |
|
|
107
|
+
| past-to-future | 11 of 203 | 5.4% (2.7–9.5%) | 52.8% |
|
|
108
|
+
|
|
109
|
+
A rule that kills every run immediately was scored by the same gate and failed on both splits.
|
|
110
|
+
|
|
111
|
+
**The flaw that run found:** every kill fired at the 3rd of 50 evaluations. Early rankings only
|
|
112
|
+
match final rankings from about evaluation 18. The rule stayed within budget because the bound
|
|
113
|
+
held. The fitted power law ranked runs worse than the latest observed value at every early
|
|
114
|
+
evaluation.
|
|
115
|
+
|
|
116
|
+
**What this package ships instead,** and what the reproduction script checks:
|
|
117
|
+
|
|
118
|
+
| predictor | hold until | random split | past-to-future split | 2000 resplits |
|
|
119
|
+
|---|---|---|---|---|
|
|
120
|
+
| last value | none | 18 of 200 (9.0%), 78% saved | 7 of 203 (3.4%), 75% saved | joint 5.0%, breach 8.5% |
|
|
121
|
+
| last value | eval 18 | 18 of 200 (9.0%), 57% saved | 7 of 203 (3.4%), 56% saved | joint 5.0%, breach 8.6% |
|
|
122
|
+
|
|
123
|
+
The random split's 18 of 200 is a breach. It sits at the 95th percentile of that rule's own
|
|
124
|
+
resplit distribution, whose mean is 5.0%, so it reads as a split tail rather than a leak. A
|
|
125
|
+
correct rule at this sample size lands on a breach in about 7% of splits. The held version was
|
|
126
|
+
chosen *after* seeing these 400 runs, so its numbers are a hypothesis. It halves
|
|
127
|
+
how often the single best run is killed at the same budget. Its confirmation needs a fresh pool.
|
|
128
|
+
|
|
129
|
+
Full account: `docs/writeup.md`. Run the check yourself:
|
|
130
|
+
|
|
131
|
+
```
|
|
132
|
+
python reproduce/reproduce_cc03c.py
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
## What the bound does not cover
|
|
136
|
+
|
|
137
|
+
- **It needs exchangeable runs.** With a random or quasi-random sampler the bound holds. With
|
|
138
|
+
TPE, Gaussian-process or CMA-ES samplers it does not. The Optuna pruner warns once and carries
|
|
139
|
+
on, but do not rely on the bound there.
|
|
140
|
+
- **It bounds the joint rate rather than the share of good runs killed.** When good runs are rare, a large
|
|
141
|
+
share of them can still die: about a third in the measurement above, mostly borderline ones.
|
|
142
|
+
Lower α if that matters.
|
|
143
|
+
- **An estimated τ gives a measured bound rather than a proven one.** `good_quantile` estimates τ from the
|
|
144
|
+
calibration runs. Simulations put the realized rate at 0.88 to 0.97 times α. Pass a real
|
|
145
|
+
`good_threshold` when you have one.
|
|
146
|
+
- **α = 0.05 needs at least 19 calibration runs**, and the certified threshold gets loose with
|
|
147
|
+
few good runs. Sweeps of a few dozen trials are too small for this.
|
|
148
|
+
- **One setting only so far.** One model, one dataset, one horizon.
|
|
149
|
+
|
|
150
|
+
## Layout
|
|
151
|
+
|
|
152
|
+
| path | what |
|
|
153
|
+
|---|---|
|
|
154
|
+
| `src/surestop/` | `KillRule` (numpy/scipy) and `ConformalPruner` (Optuna) |
|
|
155
|
+
| `tests/` | the bound holds on synthetic exchangeable runs, and the same assertion fails on a rule that leaks 2× its budget |
|
|
156
|
+
| `reproduce/` | the 400 curves, the pre-registered splits and targets, the script, and its committed output |
|
|
157
|
+
| `docs/` | the write-up, the prior-art sweep, and the frozen hyperparameter prior the curves were drawn from |
|
|
158
|
+
| `optunahub/` | the package as an OptunaHub registry entry |
|
|
159
|
+
|
|
160
|
+
## Related work
|
|
161
|
+
|
|
162
|
+
Calibrating a stop threshold on a predictor with a conformal guarantee is an established pattern. Xie et al.
|
|
163
|
+
(arXiv:2602.13935) use the same device, a threshold on the maximum of a running score calibrated
|
|
164
|
+
on one class only, to stop LLM reasoning traces. Related work does the same for mixed-integer
|
|
165
|
+
solvers and LLM agent episodes. As far as a search found, nobody has applied it to pruning
|
|
166
|
+
training runs, and no major tuning library ships a pruner with a stated error guarantee.
|
|
167
|
+
`docs/prior-art.md` has the sweep.
|
|
168
|
+
|
|
169
|
+
## Citing
|
|
170
|
+
|
|
171
|
+
Concept DOI: [10.5281/zenodo.22726440](https://doi.org/10.5281/zenodo.22726440), which always resolves
|
|
172
|
+
to the newest version. `CITATION.cff` has the full entry. The curves are a measurement, released
|
|
173
|
+
under the same MIT licence as the code.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
src/surestop/__init__.py
|
|
5
|
+
src/surestop/core.py
|
|
6
|
+
src/surestop/optuna_pruner.py
|
|
7
|
+
src/surestop.egg-info/PKG-INFO
|
|
8
|
+
src/surestop.egg-info/SOURCES.txt
|
|
9
|
+
src/surestop.egg-info/dependency_links.txt
|
|
10
|
+
src/surestop.egg-info/requires.txt
|
|
11
|
+
src/surestop.egg-info/top_level.txt
|
|
12
|
+
tests/test_core.py
|
|
13
|
+
tests/test_optuna_pruner.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
surestop
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
"""The bound, a negative control that proves the bound check can fail, and the rule's edges."""
|
|
2
|
+
|
|
3
|
+
import math
|
|
4
|
+
|
|
5
|
+
import numpy as np
|
|
6
|
+
import pytest
|
|
7
|
+
|
|
8
|
+
from surestop import KillRule, NotEnoughCalibration
|
|
9
|
+
|
|
10
|
+
T = 30
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def synth(n, rng):
|
|
14
|
+
"""Exchangeable learning curves: a final value, plus a decaying offset and noise that shrink with step."""
|
|
15
|
+
finals = rng.normal(2.17, 0.013, n)
|
|
16
|
+
t = np.arange(T)
|
|
17
|
+
offset = rng.uniform(0.05, 0.3, n)[:, None] * np.exp(-t / 6.0)
|
|
18
|
+
noise = rng.normal(0, 1, (n, T)) * 0.012 * np.exp(-t / 10.0)
|
|
19
|
+
Y = finals[:, None] + offset + noise
|
|
20
|
+
Y[:, -1] = finals
|
|
21
|
+
return Y
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def mean_joint_rate(alpha_fit, reps=400, n_cal=200, n_test=200, seed=0, **kw):
|
|
25
|
+
rng = np.random.default_rng(seed)
|
|
26
|
+
rates = []
|
|
27
|
+
for _ in range(reps):
|
|
28
|
+
cal, test = synth(n_cal, rng), synth(n_test, rng)
|
|
29
|
+
rule = KillRule(alpha=alpha_fit, good_threshold=2.159, **kw).fit(cal)
|
|
30
|
+
rates.append(rule.evaluate(test)["joint_false_kill_rate"])
|
|
31
|
+
rates = np.array(rates)
|
|
32
|
+
return rates.mean(), rates.std(ddof=1) / math.sqrt(reps)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def test_joint_false_kill_rate_is_bounded():
|
|
36
|
+
mean, se = mean_joint_rate(0.05)
|
|
37
|
+
assert mean <= 0.05 + 3 * se, (mean, se)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def test_negative_control_the_bound_check_can_fail():
|
|
41
|
+
# A rule that leaks 2x its budget: certify at 0.10, hold the result to 0.05. The assertion
|
|
42
|
+
# the real test makes must fail here. (At 0.25 the synthetic pool has too few good runs for a
|
|
43
|
+
# finite threshold, so the rule kills everything — too easy a control to prove anything.)
|
|
44
|
+
mean, se = mean_joint_rate(0.10, reps=400)
|
|
45
|
+
assert not (mean <= 0.05 + 3 * se), (mean, se)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_the_rule_actually_saves_compute():
|
|
49
|
+
rng = np.random.default_rng(1)
|
|
50
|
+
rule = KillRule(alpha=0.05, good_threshold=2.159).fit(synth(300, rng))
|
|
51
|
+
assert rule.evaluate(synth(300, rng))["savings"] > 0.2
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def test_min_peek_blocks_earlier_kills():
|
|
55
|
+
rng = np.random.default_rng(2)
|
|
56
|
+
rule = KillRule(alpha=0.05, good_threshold=2.159, min_peek=12).fit(synth(300, rng))
|
|
57
|
+
steps = [rule.kill_step(c) for c in synth(300, rng)]
|
|
58
|
+
assert all(s is None or s >= 12 for s in steps)
|
|
59
|
+
assert any(s is not None for s in steps)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def test_should_kill_agrees_with_kill_step():
|
|
63
|
+
rng = np.random.default_rng(3)
|
|
64
|
+
rule = KillRule(alpha=0.05, good_threshold=2.159, min_peek=4).fit(synth(300, rng))
|
|
65
|
+
for c in synth(50, rng):
|
|
66
|
+
first = next((t for t in range(T) if rule.should_kill(c[: t + 1])), None)
|
|
67
|
+
assert first == rule.kill_step(c)
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def test_hold_until_resolved_moves_the_first_peek_later():
|
|
71
|
+
rng = np.random.default_rng(4)
|
|
72
|
+
cal = synth(300, rng)
|
|
73
|
+
held = KillRule(alpha=0.05, good_threshold=2.159, hold_until_resolved=0.9).fit(cal)
|
|
74
|
+
assert held.first_peek_ > 0
|
|
75
|
+
assert held.resolution_[held.first_peek_] >= 0.9
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def test_too_few_calibration_runs_raises():
|
|
79
|
+
with pytest.raises(NotEnoughCalibration):
|
|
80
|
+
KillRule(alpha=0.05).fit(synth(18, np.random.default_rng(5)))
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def test_maximize_mirrors_minimize():
|
|
84
|
+
rng = np.random.default_rng(6)
|
|
85
|
+
cal, test = synth(200, rng), synth(100, rng)
|
|
86
|
+
lo = KillRule(alpha=0.05, good_threshold=2.159).fit(cal)
|
|
87
|
+
hi = KillRule(alpha=0.05, good_threshold=-2.159, maximize=True).fit(-cal)
|
|
88
|
+
assert [lo.kill_step(c) for c in test] == [hi.kill_step(-c) for c in test]
|
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
"""The Optuna pruner: shadow mode first, pruning after, a warning for adaptive samplers."""
|
|
2
|
+
|
|
3
|
+
import numpy as np
|
|
4
|
+
import pytest
|
|
5
|
+
|
|
6
|
+
optuna = pytest.importorskip("optuna")
|
|
7
|
+
from surestop import ConformalPruner # noqa: E402
|
|
8
|
+
|
|
9
|
+
optuna.logging.set_verbosity(optuna.logging.WARNING)
|
|
10
|
+
T = 20
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def objective(trial):
|
|
14
|
+
x = trial.suggest_float("x", 0.0, 1.0)
|
|
15
|
+
rng = np.random.default_rng(trial.number)
|
|
16
|
+
final = 2.15 + 0.05 * x + rng.normal(0, 0.005)
|
|
17
|
+
for step in range(T):
|
|
18
|
+
value = final + 0.2 * np.exp(-step / 4) + rng.normal(0, 0.004) * np.exp(-step / 8)
|
|
19
|
+
if step == T - 1:
|
|
20
|
+
value = final
|
|
21
|
+
trial.report(float(value), step)
|
|
22
|
+
if trial.should_prune():
|
|
23
|
+
raise optuna.TrialPruned()
|
|
24
|
+
return final
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def test_calibrates_in_shadow_mode_then_prunes():
|
|
28
|
+
pruner = ConformalPruner(n_calibration=40, alpha=0.05)
|
|
29
|
+
study = optuna.create_study(sampler=optuna.samplers.RandomSampler(seed=0), pruner=pruner)
|
|
30
|
+
study.optimize(objective, n_trials=160)
|
|
31
|
+
states = [t.state for t in study.trials]
|
|
32
|
+
assert all(s == optuna.trial.TrialState.COMPLETE for s in states[:40])
|
|
33
|
+
assert pruner.rule_ is not None and np.isfinite(pruner.rule_.lambda_)
|
|
34
|
+
assert sum(s == optuna.trial.TrialState.PRUNED for s in states) > 0
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_warns_for_an_adaptive_sampler():
|
|
38
|
+
pruner = ConformalPruner(n_calibration=20)
|
|
39
|
+
study = optuna.create_study(sampler=optuna.samplers.TPESampler(seed=0), pruner=pruner)
|
|
40
|
+
with pytest.warns(RuntimeWarning, match="not exchangeable"):
|
|
41
|
+
study.optimize(objective, n_trials=2)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def test_from_curves_is_calibrated_before_the_first_trial():
|
|
45
|
+
rng = np.random.default_rng(0)
|
|
46
|
+
curves = 2.17 + rng.normal(0, 0.01, (60, 1)) + 0.2 * np.exp(-np.arange(T) / 4)
|
|
47
|
+
pruner = ConformalPruner.from_curves(curves, steps=range(T))
|
|
48
|
+
assert pruner.rule_.n_calibration_ == 60 and pruner.steps_ == list(range(T))
|