Repository: sktime/sktime
Pull request: [ENH] add fit_strategy for TinyTimeMixerForecaster
<!--
Welcome to sktime, and thanks for contributing!
Please have a look at our contribution guide:
https://www.sktime.net/en/latest/get_involved/contributing.html
-->

#### Reference Issues/PRs
<!--
Example: Fixes #1234. See also #3456.

Please use keywords (e.g., Fixes) to create link to the issues or pull requests
you resolved, so that they will automatically be closed when your pull request
is merged. See https://github.com/blog/1506-closing-issues-via-pull-requests.
If no issue exists, you can open one here: https://github.com/sktime/sktime/issues
-->
Fixes: #7426 


#### What does this implement/fix? Explain your changes.
<!--
A clear and concise description of what you have implemented.
-->
Added fit strategy parameter for ttm forecaster to allow proper handling of model fine-tuning. Also refactored the comments accordingly.


<!-- This section is useful if you have added a test in addition to the existing ones. This will ensure that further changes to these files won't introduce the same kind of bug. It is considered good practice to add tests with newly added code to enforce the fact that the code actually works. This will reduce the chance of introducing logical bugs.
-->

#### PR checklist
<!--
Please go through the checklist below. Please feel free to remove points if they are not applicable.
-->

##### For all contributions
- [x] The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [BUG] - bugfix, [MNT] - CI, test framework, [ENH] - adding or improving code, [DOC] - writing or improving documentation or docstrings.

<!--
Thanks for contributing!
-->


Problem statement:
[ENH] Generalize training options for pre-trained models
since we have a bunch of pre-trained models in sktime, we now want to make them more coherent and comparable functionality especially with training strategy.

I have collected the implemented training strategies for the following models
`HFPatchTSTForecaster`, `HFPatchTSTForecaster`,`TimesFMForecaster`, `TinyTimeMixerForecaster`

Based on how the models weights were initialized, this depends on the configuration, we can have following options, for each option we have the corresponding training strategies

| Model Initialization / Fit Strategy        | Zero-Shot (zero-shot)                                                                                                                          | PEFT  (peft)                                      | Train All Weights (full)                          | Train Only Random Weights (minimal)                                                                                   |
| ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------- | ------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------- |
| **Random Weights**                         | Should Raise Error                                                                                                                             | Should Raise Error                                | same as "minimal"                                 | `HFPatchTSTForecaster.mode == "untrained"`                                                                            |
| **Pre-trained Weights**                    | `HFPatchTSTForecaster.mode == "zero-shot"`, `TimesFMForecaster` (only option), `TinyTimeMixerForecaster` (selected automatically for this row) | `HFTransformersForecaster.fit_strategy == "peft"` | `HFPatchTSTForecaster.mode == "finetune"`         | nothing happens here                                                                                                  |
| **Some Pre-trained & Some Random Weights** | Should Raise Error                                                                                                                             | Should Raise Error                                | `HFTransformersForecaster.fit_strategy == "full"` | `HFTransformersForecaster.fit_strategy == "minimal"`, `TinyTimeMixerForecaster` (selected automatically for this row) |

FYI: @julian-fong @benHeid @pranavvp16 @XinyuWuu

Hints from issue discussion:
Would an extension template be helpful as a reference for new global forecasters that contributors want to implement?
> Would an extension template be helpful as a reference for new global forecasters that contributors want to implement?

yes that sounds good, that template should also cover the caching behavior as is implemented with `TimesFM` and `Chronos`
@fkiraly @geetu040
I want to work on this issue, can I get some pointers on how to approach it?
This issue basically collects how the training options for different pre-trained models work and how we can keep them more coherent
Is there a particular item or feature you have in mind to work on related to this?
> This issue basically collects how the training options for different pre-trained models work and how we can keep them more coherent Is there a particular item or feature you have in mind to work on related to this?

`TinyTimeMixer ` currently automatically adjusts which **fit_strategy** to use, we can take fit_strategy from user. I particularly want to work on it.
> `TinyTimeMixer ` currently automatically adjusts which **fit_strategy** to use, we can take fit_strategy from user. I particularly want to work on it.

yes, I think it would be great to have `fit_strategy` option for `TinyTimeMixer`. We have discussed this on discord before.

FYI: @fkiraly @benHeid 

`TinyTimeMixer` currently doesnot take a `fit_strategy` from user. It trains only the weights that have been newly initialized because of a change in configuration, rest of the weights are always frozen, this mimics the `"minimal"` mode of `fit_strategy` compared to other pre-trained models. Do we want to give `fit_strategy` option to the user as another model parameter? This does not break the API and ensures backward compatibility.

Commit message: Added fit strategy parameter for ttm forecaster to allow proper handling of model fine-tuning. Also refactored the comments accordingly.

Patch to review:
diff --git a/sktime/forecasting/ttm.py b/sktime/forecasting/ttm.py
index cf77e9071b9..504c47968a3 100644
--- a/sktime/forecasting/ttm.py
+++ b/sktime/forecasting/ttm.py
@@ -31,12 +31,12 @@ class TinyTimeMixerForecaster(_BaseGlobalForecaster):
     open-sourced by IBM Research. With less than 1 Million parameters, TTM introduces
     the notion of the first-ever "tiny" pre-trained models for Time-Series Forecasting.
 
-    **Zero-shot and Fine-tuning**:
+    **Fit Strategies: Full, Minimal, and Zero-shot**
 
-    This model is designed as a versatile foundation model, capable of both
-    zero-shot forecasting and fine-tuning on custom datasets.
-    The choice between zero-shot and fine-tuning is managed internally and
-    is determined by the final configuration used to initialize the model.
+    This model supports three fit strategies: *zero-shot* for direct predictions without
+    training, *minimal* fine-tuning for lightweight adaptation to new data, and *full*
+    fine-tuning for comprehensive model training. The selected strategy is determined
+    by the model's fit_strategy parameter
 
     **Initialization Process**:
 
@@ -75,10 +75,14 @@ class TinyTimeMixerForecaster(_BaseGlobalForecaster):
       aligned with the *model architechture*, fine-tuing part is bypassed and
       the model preforms zero-short forecasting.
 
-    - **Fine-tuning**: When not all the *pre-trained weights* are correctly aligned
-      with the *model architechture*, rather some weights are re-initialized,
+    - **Minimal Fine-tuning**: When not all the *pre-trained weights* are correctly
+      aligned with the *model architechture*, rather some weights are re-initialized,
       these re-initialized weights are fine-tuned on the provided data.
 
+    - **Full Fine-tuning**:  The model is *fully fine-tuned* on new data, updating *all
+      parameters*. This approach offers maximum adaptation to the dataset but requires
+      more computational resources.
+
     Parameters
     ----------
     model_path : str, default="ibm/TTM"
@@ -125,6 +129,14 @@ class TinyTimeMixerForecaster(_BaseGlobalForecaster):
         fit and predict on it. The broadcasting is happening inside automatically,
         from the outerside api perspective, the input and output are the same,
         only one multiindex output from ``predict``.
+    fit_strategy : str, default="minimal"
+        Strategy to use for fitting (fine-tuning) the model. This can be one of
+        the following:
+        - "zero-shot": Uses pre-trained model as it is.
+        - "minimal": Fine-tunes only a small subset of the model parameters,
+          allowing for quick adaptation with limited computational resources.
+        - "full": Fine-tunes all model parameters, which may result in better
+          performance but requires more computational power and time.
 
     use_source_package : bool, default=False
         If True, the model and configuration will be loaded directly from the source
@@ -151,7 +163,7 @@ class TinyTimeMixerForecaster(_BaseGlobalForecaster):
     >>> from sktime.forecasting.ttm import TinyTimeMixerForecaster
     >>> from sktime.datasets import load_airline
     >>> y = load_airline()
-    >>> forecaster = TinyTimeMixerForecaster() # doctest: +SKIP
+    >>> forecaster = TinyTimeMixerForecaster(fit_strategy="zero-shot") # doctest: +SKIP
     >>> # performs zero-shot forecasting, as default config (unchanged) is used
     >>> forecaster.fit(y, fh=[1, 2, 3]) # doctest: +SKIP
     >>> y_pred = forecaster.predict() # doctest: +SKIP
@@ -179,7 +191,7 @@ class TinyTimeMixerForecaster(_BaseGlobalForecaster):
     ...     },
     ... ) # doctest: +SKIP
     >>>
-    >>> # model is fine-tuned, as a config different from default is provided.
+    >>> # model is fine-tuned with the minimal strategy due to a non-default config.
     >>> forecaster.fit(y, fh=[1, 2, 3]) # doctest: +SKIP
     >>> y_pred = forecaster.predict() # doctest: +SKIP
     """
@@ -221,6 +233,7 @@ def __init__(
         compute_metrics=None,
         callbacks=None,
         broadcasting=False,
+        fit_strategy="minimal",
         use_source_package=False,
     ):
         super().__init__()
@@ -235,6 +248,7 @@ def __init__(
         self.callbacks = callbacks
         self.broadcasting = broadcasting
         self.use_source_package = use_source_package
+        self.fit_strategy = fit_strategy
 
         if self.broadcasting:
             self.set_tags(
@@ -351,19 +365,32 @@ def _fit(self, y, X, fh):
             ignore_mismatched_sizes=True,
         )
 
-        if len(info["mismatched_keys"]) == 0:
-            return  # No need to fit
-
-        # Freeze all loaded parameters
-        for param in self.model.parameters():
-            param.requires_grad = False
-
-        # Reininit the weights of all layers that have mismatched sizes
-        for key, _, _ in info["mismatched_keys"]:
-            _model = self.model
-            for attr_name in key.split(".")[:-1]:
-                _model = getattr(_model, attr_name)
-            _model.weight.requires_grad = True
+        if self.fit_strategy == "zero-shot":
+            if len(info["mismatched_keys"]) > 0:
+                raise ValueError(
+                    "Fit strategy is 'zero-shot', but the model weights in the"
+                    "configuration are mismatched compared to the pretrained model."
+                    "Please ensure they match."
+                )
+            return
+        elif self.fit_strategy == "minimal":
+            if len(info["mismatched_keys"]) == 0:
+                return  # No need to fit
+            # Freeze all loaded parameters
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+            # Reininit the weights of all layers that have mismatched sizes
+            for key, _, _ in info["mismatched_keys"]:
+                _model = self.model
+                for attr_name in key.split(".")[:-1]:
+                    _model = getattr(_model, attr_name)
+                _model.weight.requires_grad = True
+        elif self.fit_strategy == "full":
+            for param in self.model.parameters():
+                param.requires_grad = True
+        else:
+            raise ValueError("Unknown fit strategy")
 
         y_train, y_test = temporal_train_test_split(y, test_size=self.validation_split)
 


Code review comments:
diff_hunk: @@ -125,6 +129,14 @@ class TinyTimeMixerForecaster(_BaseGlobalForecaster):
         fit and predict on it. The broadcasting is happening inside automatically,
         from the outerside api perspective, the input and output are the same,
         only one multiindex output from ``predict``.
+    fit_strategy : str, default="minimal"
+        Strategy to use for fitting (fine-tuning) the model. This can be one of
+        the following:
+        - "zero-shot": Uses pre-trained model as it is.
+        - "minimal": Fine-tunes only a small subset of the model parameters,
+          allowing for quick adaptation with limited computational resources.
+        - "full": Fine-tunes all model parameters, which may result in better
+          performance but requires more computational power and time.
path: sktime/forecasting/ttm.py
text: @user1:
these need more description, as what happens when `model_path is None`