feat: warn on pathologically slow custom loss functions - #1301
MilesCranmerBot wants to merge 2 commits into
Conversation
When a user passes a custom loss (`elementwise_loss`, `loss_function`, or `loss_function_expression`), benchmark it inside their own Julia process on their real data against the default squared-error loss, timed identically in the same process so the ratio is hardware-independent (one warm-up call for JIT, GC, then the median of 5 timed passes). Warn — at most once per fit — only if the custom loss is both >=50x slower than the default and takes >=1 ms per pass over the data, so a fast loss that is merely '50x slower than nothing' never fires. The warning prints both measured times, likely causes (excessive allocations, Python callbacks via PythonCall, type instability — see `@code_warntype`), and the opt-out flag. The check is opted out with `check_loss_speed=False`, degrades silently if the benchmark itself fails, and adds zero overhead when no custom loss is used. Custom operators are out of scope for this check. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Custom losses are Julia strings evaluated in-process; Python callbacks are not a reachable failure mode. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2bedf7bd10
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| f(args...) | ||
| GC.gc() | ||
| # Median of a few timed passes over the dataset: | ||
| times = map(1:5) do _ | ||
| @elapsed f(args...) |
There was a problem hiding this comment.
For the pathologically slow losses this feature is intended to detect, the warm-up plus five timed evaluations can itself stall fit for an arbitrarily long time. This runs before equation_search, so timeout_in_seconds cannot limit it; for example, a 30-second full objective adds roughly three minutes before the search even starts. Benchmark a bounded sample or stop timing after a deadline rather than requiring six complete passes.
Useful? React with 👍 / 👎.
| else | ||
| nothing | ||
| end | ||
| dataset = Dataset(X, y_vec; weights=w_vec) |
There was a problem hiding this comment.
When a valid loss_function or loss_function_expression reads dataset metadata supplied to fit, such as X_units, y_units, variable names, or the output index, this newly constructed Dataset contains defaults instead of the values passed to equation_search. The benchmark can therefore take a different branch or throw; because the exception is swallowed, the speed check silently produces no warning even though the same objective works during the real search. Construct the benchmark dataset with the same metadata as the search dataset.
Useful? React with 👍 / 👎.
| if w === nothing | ||
| total = zero(eltype(y)) | ||
| @inbounds for i in eachindex(y) | ||
| total += loss(y[i], y[i]) |
There was a problem hiding this comment.
Passing y[i] as both prediction and target benchmarks only the zero-residual path. A custom loss that cheaply returns zero for an exact match but performs expensive robust, probabilistic, or allocation-heavy work for nonzero residuals will appear fast here even though nearly every candidate evaluated during the search takes the slow path, so the intended warning is missed. Use representative predictions that differ from the targets; the weighted branch has the same issue.
Useful? React with 👍 / 👎.
| # this is safe even though the operator sets differ): | ||
| default_options = SymbolicRegression.Options() | ||
| if kind === :loss_function | ||
| tree = options.node_type(T; val=one(T)) |
There was a problem hiding this comment.
The full-objective benchmark always supplies a single constant node, so it cannot detect costs that scale with expression size. This is especially relevant because loss_function explicitly supports symbolic preprocessing and regularization: an objective that traverses, copies, simplifies, or otherwise analyzes the candidate tree can be nearly free for this node but pathologically slow on the larger trees reached during the search. Benchmark at least one representative nontrivial tree rather than only a constant.
Useful? React with 👍 / 👎.
| total = zero(eltype(y)) | ||
| weight_total = zero(eltype(y)) | ||
| @inbounds for i in eachindex(y) | ||
| total += loss(y[i], y[i], w[i]) |
There was a problem hiding this comment.
When weights is supplied with a supported LossFunctions.jl object such as the documented LPDistLoss{3}(), the real evaluator computes weight * loss(prediction, target) because these SupervisedLoss objects do not accept a third scalar argument. This benchmark instead calls loss(y[i], y[i], w[i]), raising a MethodError that the surrounding Python handler silently swallows, so the speed check never runs for this valid configuration. Detect SupervisedLoss and apply the weight outside its two-argument call, matching the backend.
Useful? React with 👍 / 👎.
| end | ||
| function _pysr_loss_speed_benchmark(kind::Symbol, loss::F, options, X, y, w) where {F} | ||
| if kind === :elementwise_loss | ||
| y_flat = vec(y) |
There was a problem hiding this comment.
For multi-output y, vec(y) combines every output into one benchmark pass, while equation_search searches and evaluates each output on its own dataset. Consequently the measured absolute time is multiplied by the number of outputs: a loss taking less than the 1 ms threshold per real evaluation can cross the threshold here and emit a false warning solely because several outputs were concatenated. Benchmark each output separately, or use one representative output as the full-objective branch already does.
Useful? React with 👍 / 👎.
This branch has not been deployed
2 participants
Summary
When a user passes a custom loss (
elementwise_loss,loss_function, orloss_function_expression), PySR now benchmarks it at the start offitand warns if it is pathologically slow, using the built-in default squared-error loss as a yardstick:GC.gc(), then the median of 5@elapsedpasses. The default squared-error loss is timed identically in the same process, so the ratio (custom/default) is hardware/OS independent.@code_warntype), and the opt-out flag.check_loss_speed=False(on by default only when a custom loss is provided; there is zero overhead and no side effects when the default loss is used).Full-objective losses (
loss_function/loss_function_expression) are benchmarked with a constant tree on a realDatasetbuilt from the user's data, and the baseline iseval_lossunder default options (the user'soptionswould just re-time their own loss). Elementwise losses are benchmarked with the same element-by-element call pattern the search uses.Example warning:
Testing
TestLossSpeedWarning(new): warning fires on an allocation-heavyelementwise_loss(plain and weighted) and on a slowloss_function; no warning with the default loss, with a fast custom loss, or withcheck_loss_speed=False; a failing benchmark degrades silently. All fits use 1 trivial iteration so the tests stay fast.test_main.py(114 passed, 1 skipped),test_cli.py(6 passed), loss-relatedtest_startup.pytests (notebook test needs thenbvalplugin, skipped locally).black,isort, andmypy(CI invocation) are clean.