From 79d6d403d8999de4aea8b388ca10905c8bf3de33 Mon Sep 17 00:00:00 2001 From: Nikhil1869 Date: Sat, 11 Apr 2026 12:01:04 +0530 Subject: [PATCH] Fix data leakage in normalization window (#227) --- finetune/dataset.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/finetune/dataset.py b/finetune/dataset.py index f955ec1..ae4f724 100644 --- a/finetune/dataset.py +++ b/finetune/dataset.py @@ -90,21 +90,7 @@ class QlibDataset(Dataset): return self.n_samples def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: - """ - Retrieves a random sample from the dataset. - - Note: The `idx` argument is ignored. Instead, a random index is drawn - from the pre-computed `self.indices` list using `self.py_rng`. This - ensures random sampling over the entire dataset for each call. - - Args: - idx (int): Ignored. - - Returns: - tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - x_tensor (torch.Tensor): The normalized feature tensor. - - x_stamp_tensor (torch.Tensor): The time feature tensor. - """ + # Select a random sample from the entire pool of indices. random_idx = self.py_rng.randint(0, len(self.indices) - 1) symbol, start_idx = self.indices[random_idx] @@ -118,8 +104,15 @@ class QlibDataset(Dataset): x = win_df[self.feature_list].values.astype(np.float32) x_stamp = win_df[self.time_feature_list].values.astype(np.float32) - # Perform instance-level normalization. - x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0) + # Normalize the window. Mean and std are calculated strictly on the + # lookback window (past data) to prevent future data leakage. + past_len = self.config.lookback_window + past_x = x[:past_len] + + x_mean = np.mean(past_x, axis=0) + x_std = np.std(past_x, axis=0) + + # Apply normalization and robust clipping to the entire sequence x = (x - x_mean) / (x_std + 1e-5) x = np.clip(x, -self.config.clip, self.config.clip)