Skip to content

Fix foundation model fine-tuning metrics using all pretrained quantiles - #3318

Draft
aashoday wants to merge 2 commits into
unit8co:masterfrom
aashoday:fix/3218-foundation-finetuning-metrics
Draft

aashoday wants to merge 2 commits into
unit8co:masterfrom
aashoday:fix/3218-foundation-finetuning-metrics

Conversation

@aashoday

Copy link
Copy Markdown

Checklist before merging this PR:

  • Read and understood the contribution guidelines including AI/agent usage.
  • Linked at least one approved issue below (an issue with label ok).
  • Summarized the updates of this PR under Summary.
  • Added an entry under Unreleased in the Changelog.

Fixes #3218

Summary

Foundation models return all pretrained quantiles during fine-tuning which is needed for the quantile loss, but PLForecastingModule._update_metrics() used that output directly for the training metrics. With likelihood=None this crashes with a shape mismatch; with a user-selected quantile subset it silently pairs values with the wrong quantile levels.

This PR adds a hook to the base module and uses it so that metrics only see the user-selected predictions, just like validation does, while the training loss still uses all pretrained quantiles:

  • PLForecastingModule._metrics_prediction(output): new hook, returns output.prediction by default. _update_metrics() now calls it. Behaviour is unchanged for all other models.
  • _TiRexModule._metrics_prediction(): during fine-tuning, selects the user quantiles (the median if likelihood=None) from the all-quantile training output.

Work in progress:

  • Same override for TiRex2Model, Chronos2Model, TimesFM2p5Model, TimesFM3Model, PatchTSTFMModel.
  • Regression tests.
  • Changelog entry.
  • Docstrings: likelihood currently says it only affects prediction output.

Other Information

Question: T0Model supports arbitrary quantile levels at prediction time through interpolation, whereas its training output only contains its fixed pretrained quantiles. This means the requested quantiles cannot always be selected directly from the training output. How should this case be handled when computing training metrics?

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Foundation model fine-tuning metrics use all pretrained quantiles instead of user-selected predictions

1 participant