autogluon.timeseries 1.1.2b20241112__py3-none-any.whl → 1.1.2b20241114__py3-none-any.whl

autogluon/timeseries/models/chronos/pipeline/chronos.py

@@ -65,9 +65,12 @@ class ChronosTokenizer:
     which concrete classes must implement.
     """
 
-    def input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Any]:
+    def context_input_transform(
+        self,
+        context: torch.Tensor,
+    ) -> Tuple:
         """
-        Turn a batch of time series into token IDs, attention map, and scale.
+        Turn a batch of time series into token IDs, attention mask, and tokenizer_state.
 
         Parameters
         ----------
@@ -87,9 +90,40 @@ class ChronosTokenizer:
             which input observations are not ``torch.nan`` (i.e. not
             missing nor padding).
         tokenizer_state
-            An object that will be passed to ``output_transform``.
-            Contains the relevant context to decode output samples into
-            real values, such as location and scale parameters.
+            An object that can be passed to ``label_input_transform``
+            and ``output_transform``. Contains the relevant information
+            to decode output samples into real values,
+            such as location and scale parameters.
+        """
+        raise NotImplementedError()
+
+    def label_input_transform(self, label: torch.Tensor, tokenizer_state: Any) -> Tuple:
+        """
+        Turn a batch of label slices of time series into token IDs and attention mask
+        using the ``tokenizer_state`` provided by ``context_input_transform``.
+
+        Parameters
+        ----------
+        label
+            A tensor shaped (batch_size, time_length), containing the
+            timeseries label, i.e., the ground-truth future values.
+        tokenizer_state
+            An object returned by ``context_input_transform`` containing
+            relevant information to preprocess data, such as location and
+            scale. The nature of this depends on the specific tokenizer.
+            This is used for tokenizing the label, in order to use the same
+            scaling used to tokenize the context.
+
+        Returns
+        -------
+        token_ids
+            A tensor of integers, shaped (batch_size, time_length + 1)
+            if ``config.use_eos_token`` and (batch_size, time_length)
+            otherwise, containing token IDs for the input series.
+        attention_mask
+            A boolean tensor, same shape as ``token_ids``, indicating
+            which input observations are not ``torch.nan`` (i.e. not
+            missing nor padding).
         """
         raise NotImplementedError()
 
@@ -117,6 +151,11 @@ class ChronosTokenizer:
 
 
 class MeanScaleUniformBins(ChronosTokenizer):
+    """
+    A tokenizer that performs mean scaling and then quantizes the scaled time series into
+    uniformly-spaced bins between some bounds on the real line.
+    """
+
     def __init__(self, low_limit: float, high_limit: float, config: ChronosConfig) -> None:
         self.config = config
         self.centers = torch.linspace(
@@ -132,15 +171,15 @@ class MeanScaleUniformBins(ChronosTokenizer):
             )
         )
 
-    def input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        batch_size, length = context.shape
+    def _input_transform(
+        self, context: torch.Tensor, scale: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        attention_mask = ~torch.isnan(context)
 
-        if length > self.config.context_length:
-            context = context[..., -self.config.context_length :]
+        if scale is None:
+            scale = torch.nansum(torch.abs(context) * attention_mask, dim=-1) / torch.nansum(attention_mask, dim=-1)
+            scale[~(scale > 0)] = 1.0
 
-        attention_mask = ~torch.isnan(context)
-        scale = torch.nansum(torch.abs(context) * attention_mask, dim=-1) / torch.nansum(attention_mask, dim=-1)
-        scale[~(scale > 0)] = 1.0
         scaled_context = context / scale.unsqueeze(dim=-1)
         token_ids = (
             torch.bucketize(
@@ -153,15 +192,42 @@ class MeanScaleUniformBins(ChronosTokenizer):
             + self.config.n_special_tokens
         )
         token_ids[~attention_mask] = self.config.pad_token_id
+        token_ids.clamp_(0, self.config.n_tokens - 1)
 
-        if self.config.use_eos_token:
-            eos_tokens = torch.full((batch_size, 1), fill_value=self.config.eos_token_id)
-            token_ids = torch.concat((token_ids, eos_tokens), dim=1)
-            eos_mask = torch.full((batch_size, 1), fill_value=True)
-            attention_mask = torch.concat((attention_mask, eos_mask), dim=1)
+        return token_ids, attention_mask, scale
+
+    def _append_eos_token(
+        self, token_ids: torch.Tensor, attention_mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size = token_ids.shape[0]
+        eos_tokens = torch.full((batch_size, 1), fill_value=self.config.eos_token_id)
+        token_ids = torch.concat((token_ids, eos_tokens), dim=1)
+        eos_mask = torch.full((batch_size, 1), fill_value=True)
+        attention_mask = torch.concat((attention_mask, eos_mask), dim=1)
+
+        return token_ids, attention_mask
+
+    def context_input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        length = context.shape[-1]
+
+        if length > self.config.context_length:
+            context = context[..., -self.config.context_length :]
+
+        token_ids, attention_mask, scale = self._input_transform(context=context)
+
+        if self.config.use_eos_token and self.config.model_type == "seq2seq":
+            token_ids, attention_mask = self._append_eos_token(token_ids=token_ids, attention_mask=attention_mask)
 
         return token_ids, attention_mask, scale
 
+    def label_input_transform(self, label: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        token_ids, attention_mask, _ = self._input_transform(context=label, scale=scale)
+
+        if self.config.use_eos_token:
+            token_ids, attention_mask = self._append_eos_token(token_ids=token_ids, attention_mask=attention_mask)
+
+        return token_ids, attention_mask
+
     def output_transform(self, samples: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
         scale_unsqueezed = scale.unsqueeze(-1).unsqueeze(-1)
         indices = torch.clamp(
@@ -302,6 +368,7 @@ class ChronosPipeline(BaseChronosPipeline):
     forecast_type: ForecastType = ForecastType.SAMPLES
 
     def __init__(self, tokenizer, model):
+        super().__init__(inner_model=model.model)
         self.tokenizer = tokenizer
         self.model = model
 
@@ -330,7 +397,7 @@ class ChronosPipeline(BaseChronosPipeline):
             provided, and the extra 1 is for EOS.
         """
         context = self._prepare_and_validate_context(context=context)
-        token_ids, attention_mask, tokenizer_state = self.tokenizer.input_transform(context)
+        token_ids, attention_mask, tokenizer_state = self.tokenizer.context_input_transform(context)
         embeddings = self.model.encode(
             input_ids=token_ids.to(self.model.device),
             attention_mask=attention_mask.to(self.model.device),
@@ -402,7 +469,7 @@ class ChronosPipeline(BaseChronosPipeline):
         remaining = prediction_length
 
         while remaining > 0:
-            token_ids, attention_mask, scale = self.tokenizer.input_transform(context)
+            token_ids, attention_mask, scale = self.tokenizer.context_input_transform(context)
             samples = self.model(
                 token_ids.to(self.model.device),
                 attention_mask.to(self.model.device),

autogluon/timeseries/models/chronos/pipeline/chronos_bolt.py

@@ -289,7 +289,7 @@ class ChronosBoltModelForForecasting(T5PreTrainedModel):
             # normalize target
             target, _ = self.instance_norm(target, loc_scale)
             target = target.unsqueeze(1)  # type: ignore
-            assert self.chronos_config.prediction_length == target.shape[-1]
+            assert self.chronos_config.prediction_length >= target.shape[-1]
 
             target = target.to(quantile_preds.device)
             target_mask = (
@@ -297,6 +297,12 @@ class ChronosBoltModelForForecasting(T5PreTrainedModel):
             )
             target[~target_mask] = 0.0
 
+            # pad target and target_mask if they are shorter than model's prediction_length
+            if self.chronos_config.prediction_length > target.shape[-1]:
+                padding_shape = (*target.shape[:-1], self.chronos_config.prediction_length - target.shape[-1])
+                target = torch.cat([target, torch.zeros(padding_shape).to(target)], dim=-1)
+                target_mask = torch.cat([target_mask, torch.zeros(padding_shape).to(target_mask)], dim=-1)
+
             loss = (
                 2
                 * torch.abs(
@@ -373,6 +379,7 @@ class ChronosBoltPipeline(BaseChronosPipeline):
     _aliases = ["PatchedT5Pipeline"]
 
     def __init__(self, model: ChronosBoltModelForForecasting):
+        super().__init__(inner_model=model)
         self.model = model
 
     @property

autogluon/timeseries/models/chronos/pipeline/utils.py

@@ -1,15 +1,212 @@
+import logging
 import os
 import re
 import time
+from itertools import chain, cycle
 from pathlib import Path
-from typing import Callable, List, Optional
+from typing import TYPE_CHECKING, Callable, Iterable, Iterator, List, Literal, Optional
 
 import numpy as np
 import torch
+from gluonts.dataset.field_names import FieldName
+from gluonts.transform import ExpectedNumInstanceSampler, InstanceSplitter, ValidationSplitSampler
+from torch.utils.data import IterableDataset
+from transformers import TrainerCallback
 
 from autogluon.common.loaders.load_s3 import download, list_bucket_prefix_suffix_contains_s3
 from autogluon.core.utils.exceptions import TimeLimitExceeded
 from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
+from autogluon.timeseries.models.gluonts.abstract_gluonts import SimpleGluonTSDataset
+
+if TYPE_CHECKING:
+    # TODO: fix the underlying reason for this circular import, the pipeline should handle tokenization
+    from autogluon.timeseries.models.chronos.pipeline.chronos import ChronosTokenizer
+
+
+logger = logging.getLogger("autogluon.timeseries.models.chronos")
+
+
+class PseudoShuffledIterableDataset(IterableDataset):
+    """
+    Shuffle entries from an iterable by temporarily accumulating them
+    in an intermediate buffer.
+
+    Parameters
+    ----------
+    base_dataset
+        The original iterable object, representing the dataset.
+    shuffle_buffer_size
+        Size of the buffer use to shuffle entries from the base dataset.
+    """
+
+    def __init__(self, base_dataset, shuffle_buffer_size: int = 100) -> None:
+        super().__init__()
+        assert shuffle_buffer_size > 0
+        self.base_dataset = base_dataset
+        self.shuffle_buffer_size = shuffle_buffer_size
+        self.generator = torch.Generator()
+
+    def __iter__(self):
+        shuffle_buffer = []
+
+        for element in self.base_dataset:
+            shuffle_buffer.append(element)
+            if len(shuffle_buffer) >= self.shuffle_buffer_size:
+                idx = torch.randint(len(shuffle_buffer), size=(), generator=self.generator)
+                yield shuffle_buffer.pop(idx)
+
+        while shuffle_buffer:
+            idx = torch.randint(len(shuffle_buffer), size=(), generator=self.generator)
+            yield shuffle_buffer.pop(idx)
+
+
+class ChronosFineTuningDataset(IterableDataset):
+    """
+    Dataset wrapper to convert a ``TimeSeriesDataFrame`` into an iterable dataset
+    compatible with Chronos models.
+
+    When a ``tokenizer`` is provided, data is converted into HuggingFace-compatible set of
+    ``input_ids``, ``attention_mask`` and ``labels``, used by the original Chronos models.
+
+    When the ``tokenizer`` is omitted, data is converted into the format compatible with
+    ChronosBolt models, i.e., ``context`` and ``target``.
+
+    Parameters
+    ----------
+    target_df : TimeSeriesDataFrame
+        The ``TimeSeriesDataFrame`` to be converted
+    target_column : str, default = "target"
+        The name of the column which contains the target time series, by default "target"
+    context_length : int, default = 512
+        The length of the historical context
+    prediction_length : int, default = 64
+        The prediction_length, i.e., length of label or target
+    tokenizer : ``ChronosTokenizer``, default = None
+        When a ``ChronosTokenizer`` object is provided, data will be converted into the
+        HuggingFace format accepted by the original Chronos models using this ``ChronosTokenizer``.
+        If None, data will be converted into the format accepted by ChronosBolt models.
+    mode : Literal["training", "validation"], default = "training"
+        When ``training``, random slices from the time series will be returned for training purposes.
+        If ``validation``, the last slice of each time series returned in the original order.
+    """
+
+    def __init__(
+        self,
+        target_df: TimeSeriesDataFrame,
+        target_column: str = "target",
+        context_length: int = 512,
+        prediction_length: int = 64,
+        tokenizer: Optional["ChronosTokenizer"] = None,
+        mode: Literal["training", "validation"] = "training",
+    ) -> None:
+        super().__init__()
+
+        assert mode in ("training", "validation")
+
+        # A dummy hourly freq is used because the model doesn't actually need the freq
+        self.gluonts_dataset = SimpleGluonTSDataset(target_df=target_df, freq="h", target_column=target_column)
+        self.tokenizer = tokenizer
+        self.context_length = context_length
+        self.prediction_length = prediction_length
+        self.mode = mode
+
+    def _create_instance_splitter(self, mode: str):
+        instance_sampler = {
+            "training": ExpectedNumInstanceSampler(
+                num_instances=1.0, min_future=self.prediction_length, min_instances=1
+            ),
+            "validation": ValidationSplitSampler(min_future=self.prediction_length),
+        }[mode]
+
+        return InstanceSplitter(
+            target_field=FieldName.TARGET,
+            is_pad_field=FieldName.IS_PAD,
+            start_field=FieldName.START,
+            forecast_start_field=FieldName.FORECAST_START,
+            instance_sampler=instance_sampler,
+            past_length=self.context_length,
+            future_length=self.prediction_length,
+            dummy_value=np.nan,
+        )
+
+    def _create_training_data(self, data: Iterable[dict]):
+        data = chain.from_iterable(cycle([data]))
+        split_transform = self._create_instance_splitter("training")
+        data = split_transform.apply(data, is_train=True)
+        return data
+
+    def _create_validation_data(self, data: Iterable[dict]):
+        data = self._create_instance_splitter("validation").apply(data, is_train=False)
+        return data
+
+    def to_chronos_format(self, entry: dict) -> dict:
+        """Converts an entry from GluonTS data format with past and future targets
+        to the HuggingFace format accepted by the original Chronos models using the ChronosTokenizer.
+
+        Parameters
+        ----------
+        entry : dict
+            time series data entry in GluonTS format with ``past_target`` and ``future_target`` keys
+
+        Returns
+        -------
+        dict
+            time series data entry in HuggingFace format with ``input_ids``, ``attention_mask``, and ``labels``
+        """
+        assert self.tokenizer is not None, "A ChronosTokenizer is required to convert data into the Chronos format"
+        past_target = torch.tensor(entry[f"past_{FieldName.TARGET}"]).unsqueeze(0)
+        input_ids, attention_mask, scale = self.tokenizer.context_input_transform(past_target)
+        future_target = torch.tensor(entry[f"future_{FieldName.TARGET}"]).unsqueeze(0)
+        labels, labels_mask = self.tokenizer.label_input_transform(future_target, scale)
+        labels[labels_mask == 0] = -100
+
+        return {
+            "input_ids": input_ids.squeeze(0),
+            "attention_mask": attention_mask.squeeze(0),
+            "labels": labels.squeeze(0),
+        }
+
+    def to_chronos_bolt_format(self, entry: dict) -> dict:
+        """Converts an entry from GluonTS data format with past and future targets
+        to the format accepted by the ChronosBolt models.
+
+        Parameters
+        ----------
+        entry : dict
+            time series data entry in GluonTS format with ``past_target`` and ``future_target`` keys
+
+        Returns
+        -------
+        dict
+            time series data entry in ChronosBolt format with ``context`` and ``target``
+        """
+        past_target = torch.tensor(entry[f"past_{FieldName.TARGET}"])
+        future_target = torch.tensor(entry[f"future_{FieldName.TARGET}"])
+
+        return {"context": past_target, "target": future_target}
+
+    def __iter__(self) -> Iterator:
+        if self.mode == "training":
+            iterable = self._create_training_data(self.gluonts_dataset)
+        elif self.mode == "validation":
+            iterable = self._create_validation_data(self.gluonts_dataset)
+
+        format_transform_fn = self.to_chronos_format if self.tokenizer is not None else self.to_chronos_bolt_format
+        for entry in iterable:
+            yield format_transform_fn(entry)
+
+    def shuffle(self, shuffle_buffer_size: Optional[int] = None):
+        """Returns a (pseudo) shuffled version of this iterable dataset.
+
+        Parameters
+        ----------
+        shuffle_buffer_size : int, optional, default = None
+            The shuffle buffer size used for pseudo shuffling
+        """
+        assert shuffle_buffer_size is None or shuffle_buffer_size >= 0
+        if not shuffle_buffer_size:
+            return self
+        return PseudoShuffledIterableDataset(self, shuffle_buffer_size)
 
 
 def left_pad_and_stack_1D(tensors: List[torch.Tensor]) -> torch.Tensor:
@@ -91,12 +288,51 @@ class ChronosInferenceDataLoader(torch.utils.data.DataLoader):
             self.callback()
 
 
+class EvaluateAndSaveFinalStepCallback(TrainerCallback):
+    """Callback to evaluate and save the model at last training step."""
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if state.global_step >= state.max_steps:
+            control.should_log = True
+            control.should_evaluate = True
+            control.should_save = True
+
+
+class TimeLimitCallback(TrainerCallback):
+    def __init__(self, time_limit: int):
+        """
+        Callback to stop training once a specified time has elapsed.
+
+        Parameters
+        ----------
+        time_limit: int
+            maximum time allowed for training in seconds.
+        """
+        self.time_limit = time_limit
+        self.start_time = None
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self.start_time = time.monotonic()
+
+    def on_step_end(self, args, state, control, **kwargs):
+        elapsed_time = time.monotonic() - self.start_time
+        if elapsed_time > self.time_limit:
+            raise TimeLimitExceeded
+
+
+class LoggerCallback(TrainerCallback):
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        logs.pop("total_flos", None)
+        if state.is_local_process_zero:
+            logger.info(logs)
+
+
 def timeout_callback(seconds: Optional[float]) -> Callable:
     """Return a callback object that raises an exception if time limit is exceeded."""
-    start_time = time.time()
+    start_time = time.monotonic()
 
     def callback() -> None:
-        if seconds is not None and time.time() - start_time > seconds:
+        if seconds is not None and time.monotonic() - start_time > seconds:
             raise TimeLimitExceeded
 
     return callback

autogluon/timeseries/models/gluonts/abstract_gluonts.py

@@ -286,10 +286,6 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         self.negative_data = (dataset[self.target] < 0).any()
 
-    @property
-    def default_context_length(self) -> int:
-        return min(512, max(10, 2 * self.prediction_length))
-
     def preprocess(self, data: TimeSeriesDataFrame, is_train: bool = False, **kwargs) -> TimeSeriesDataFrame:
         # Copy data to avoid SettingWithCopyWarning from pandas
         data = data.copy()
@@ -357,25 +353,40 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             known_covariates[columns] = self._real_column_transformers["known"].transform(known_covariates[columns])
         return known_covariates
 
+    def _get_default_params(self):
+        """Gets default parameters for GluonTS estimator initialization that are available after
+        AbstractTimeSeriesModel initialization (i.e., before deferred initialization). Models may
+        override this method to update default parameters.
+        """
+        return {
+            "batch_size": 64,
+            "context_length": min(512, max(10, 2 * self.prediction_length)),
+            "predict_batch_size": 500,
+            "early_stopping_patience": 20,
+            "max_epochs": 100,
+            "lr": 1e-3,
+            "freq": self._dummy_gluonts_freq,
+            "prediction_length": self.prediction_length,
+            "quantiles": self.quantile_levels,
+        }
+
     def _get_model_params(self) -> dict:
         """Gets params that are passed to the inner model."""
-        init_args = super()._get_model_params().copy()
-        init_args.setdefault("batch_size", 64)
-        init_args.setdefault("context_length", self.default_context_length)
-        init_args.setdefault("predict_batch_size", 500)
-        init_args.setdefault("early_stopping_patience", 20)
-        init_args.update(
-            dict(
-                freq=self._dummy_gluonts_freq,
-                prediction_length=self.prediction_length,
-                quantiles=self.quantile_levels,
-                callbacks=self.callbacks,
-            )
-        )
-        # Support MXNet kwarg names for backwards compatibility
-        init_args.setdefault("lr", init_args.get("learning_rate", 1e-3))
-        init_args.setdefault("max_epochs", init_args.get("epochs", 100))
-        return init_args
+        # for backward compatibility with the old GluonTS MXNet API
+        parameter_name_aliases = {
+            "epochs": "max_epochs",
+            "learning_rate": "lr",
+        }
+
+        init_args = super()._get_model_params()
+        for alias, actual in parameter_name_aliases.items():
+            if alias in init_args:
+                if actual in init_args:
+                    raise ValueError(f"Parameter '{alias}' cannot be specified when '{actual}' is also specified.")
+                else:
+                    init_args[actual] = init_args.pop(alias)
+
+        return self._get_default_params() | init_args
 
     def _get_estimator_init_args(self) -> Dict[str, Any]:
         """Get GluonTS specific constructor arguments for estimator objects, an alias to `self._get_model_params`
@@ -395,7 +406,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         default_trainer_kwargs = {
             "limit_val_batches": 3,
             "max_epochs": init_args["max_epochs"],
-            "callbacks": init_args["callbacks"],
+            "callbacks": self.callbacks,
             "enable_progress_bar": False,
             "default_root_dir": self.path,
         }

autogluon/timeseries/models/gluonts/torch/models.py

@@ -209,15 +209,16 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
     supports_cat_covariates = True
     supports_static_features = True
 
-    @property
-    def default_context_length(self) -> int:
-        return min(512, max(64, 2 * self.prediction_length))
-
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.tft import TemporalFusionTransformerEstimator
 
         return TemporalFusionTransformerEstimator
 
+    def _get_default_params(self):
+        return super()._get_default_params() | {
+            "context_length": min(512, max(64, 2 * self.prediction_length)),
+        }
+
     def _get_estimator_init_args(self) -> Dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         if self.num_feat_dynamic_real > 0:
@@ -282,9 +283,10 @@ class DLinearModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
-    @property
-    def default_context_length(self) -> int:
-        return 96
+    def _get_default_params(self):
+        return super()._get_default_params() | {
+            "context_length": 96,
+        }
 
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.d_linear import DLinearEstimator
@@ -341,18 +343,16 @@ class PatchTSTModel(AbstractGluonTSModel):
 
     supports_known_covariates = True
 
-    @property
-    def default_context_length(self) -> int:
-        return 96
-
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.patch_tst import PatchTSTEstimator
 
         return PatchTSTEstimator
 
+    def _get_default_params(self):
+        return super()._get_default_params() | {"context_length": 96, "patch_len": 16}
+
     def _get_estimator_init_args(self) -> Dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
-        init_kwargs.setdefault("patch_len", 16)
         init_kwargs["num_feat_dynamic_real"] = self.num_feat_dynamic_real
         return init_kwargs
 
@@ -467,27 +467,27 @@ class TiDEModel(AbstractGluonTSModel):
         If False, past covariates will be used by the model if they are present in the dataset.
     feat_proj_hidden_dim : int, default = 4
         Size of the feature projection layer.
-    encoder_hidden_dim : int, default = 4
+    encoder_hidden_dim : int, default = 64
         Size of the dense encoder layer.
-    decoder_hidden_dim : int, default = 4
+    decoder_hidden_dim : int, default = 64
         Size of the dense decoder layer.
-    temporal_hidden_dim : int, default = 4
+    temporal_hidden_dim : int, default = 64
         Size of the temporal decoder layer.
-    distr_hidden_dim : int, default = 4
+    distr_hidden_dim : int, default = 64
         Size of the distribution projection layer.
-    num_layers_encoder : int, default = 1
+    num_layers_encoder : int, default = 2
         Number of layers in dense encoder.
-    num_layers_decoder : int, default = 1
+    num_layers_decoder : int, default = 2
         Number of layers in dense decoder.
-    decoder_output_dim : int, default = 4
+    decoder_output_dim : int, default = 16
         Output size of the dense decoder.
-    dropout_rate : float, default = 0.3
+    dropout_rate : float, default = 0.2
         Dropout regularization parameter.
     num_feat_dynamic_proj : int, default = 2
         Output size of feature projection layer.
     embedding_dimension : int, default = [16] * num_feat_static_cat
         Dimension of the embeddings for categorical features
-    layer_norm : bool, default = False
+    layer_norm : bool, default = True
         Should layer normalization be enabled?
     scaling : {"mean", "std", None}, default = "mean"
         Scaling applied to each *context window* during training & prediction.
@@ -496,13 +496,13 @@ class TiDEModel(AbstractGluonTSModel):
         Note that this is different from the `target_scaler` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
-    batch_size : int, default = 64
+    batch_size : int, default = 256
         Size of batches used during training
     predict_batch_size : int, default = 500
         Size of batches used during prediction.
     num_batches_per_epoch : int, default = 50
         Number of batches processed every epoch
-    lr : float, default = 1e-3,
+    lr : float, default = 1e-4,
         Learning rate used during training
     trainer_kwargs : dict, optional
         Optional keyword arguments passed to ``lightning.Trainer``.
@@ -515,15 +515,27 @@ class TiDEModel(AbstractGluonTSModel):
     supports_known_covariates = True
     supports_static_features = True
 
-    @property
-    def default_context_length(self) -> int:
-        return min(512, max(64, 2 * self.prediction_length))
-
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.tide import TiDEEstimator
 
         return TiDEEstimator
 
+    def _get_default_params(self):
+        return super()._get_default_params() | {
+            "context_length": min(512, max(64, 2 * self.prediction_length)),
+            "encoder_hidden_dim": 64,
+            "decoder_hidden_dim": 64,
+            "temporal_hidden_dim": 64,
+            "distr_hidden_dim": 64,
+            "num_layers_encoder": 2,
+            "num_layers_decoder": 2,
+            "decoder_output_dim": 16,
+            "dropout_rate": 0.2,
+            "layer_norm": True,
+            "lr": 1e-4,
+            "batch_size": 256,
+        }
+
     def _get_estimator_init_args(self) -> Dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat