🤖 machine learning

ReAugment: Model Zoo-Guided RL for Few-Shot Time Series Augmentation and Forecasting

ReAugment is a reinforcement learning-based framework for few-shot time series forecasting that utilizes a model zoo to identify overfitting-prone samples and adaptively generates targeted augmentations to enhance training diversity and improve prediction performance.

Original authors: Haochen Yuan, Yutong Wang, Yihong Chen, Yunbo Wang, Xiaokang Yang

Published 2026-08-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haochen Yuan, Yutong Wang, Yihong Chen, Yunbo Wang, Xiaokang Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Time series forecasting is the art of predicting the future based on patterns in the past. Whether it is guessing the electricity demand for tomorrow, the price of a stock next week, or the weather for the coming month, these tasks rely on finding rhythms in data that change over time. In recent years, powerful computer programs have learned to spot these complex rhythms with remarkable accuracy, but they have a significant weakness: they need vast amounts of historical data to work well. When a new event occurs or when data is scarce, these programs often struggle, memorizing the few examples they have instead of learning the underlying rules. This leads to a situation where the model performs perfectly on the data it has seen but fails miserably when asked to predict something new. To solve this, scientists have long tried to create more data artificially, a process known as augmentation. However, simply adding random noise or shuffling numbers often creates fake data that confuses the model rather than helping it. The challenge has been to generate new, realistic examples that specifically target the weak spots where the model is most likely to fail.

A team of researchers has introduced a new approach called ReAugment, which treats data creation as a guided learning process rather than a random one. Instead of guessing which data points need help, the system first identifies the specific moments in the training data where different prediction models disagree the most. Imagine a group of experts looking at the same set of historical records; if they all agree on a pattern, that pattern is likely solid. But if they argue or produce very different predictions for a specific time period, that disagreement signals a "bottleneck"—a place where the model is fragile and prone to overfitting, or memorizing the noise instead of the signal. The researchers built a "zoo" of different forecasting models to find these points of disagreement. By measuring how much the models' predictions vary for each piece of data, they can pinpoint exactly which samples are the most difficult and most critical for the system to learn.

Once these weak spots are identified, the system uses a reinforcement learning framework to generate new data specifically around them. Reinforcement learning is a method where a computer agent learns by trial and error, receiving feedback based on how well it performs a task. In this case, the agent acts as a data generator. It takes the difficult, overfit-prone samples and creates new, slightly altered versions of them. The goal is not just to make more data, but to make data that forces the forecasting model to become more robust. The system is guided by a reward mechanism that encourages the generation of samples that are both realistic and challenging. It wants to create variations that the current models find hard to predict, effectively filling in the gaps in the model's understanding. Crucially, this process is designed to be efficient; it avoids the complex mathematical calculations that usually slow down such systems, allowing the agent to learn quickly without getting stuck in calculation loops.

The researchers tested this method on a wide variety of real-world datasets, including electricity usage, traffic flow, and weather patterns. They compared their system against traditional methods that add random noise and other advanced techniques that try to learn data patterns automatically. The results showed that ReAugment consistently improved the accuracy of the forecasting models, particularly in situations where data was scarce. In tests where the models were only allowed to see a small fraction of the available training data, the new method helped them recover much of the performance they would have had if they had seen the full dataset. The system proved effective across different types of forecasting architectures, suggesting that the problem of overfitting is a general issue that can be addressed by targeting the right data points. Even when the models had access to all the data, the method still provided a measurable improvement, indicating that the targeted augmentation helps refine the model's understanding even when data is plentiful.

One of the key findings was that not all data is equally important for training. The experiments revealed that focusing the augmentation efforts on the top half of the most difficult samples yielded better results than trying to augment the entire dataset. This suggests that the models do not need more data everywhere; they need more data specifically where they are struggling to generalize. The researchers also found that their method worked well without needing to be tied to a specific type of forecasting model, making it a flexible tool that can be applied to various systems. By using a group of models to identify the weak spots and then using a learning agent to fill those spots with targeted examples, ReAugment offers a new way to teach computers how to predict the future with less information. The approach demonstrates that by understanding where a model fails, we can teach it to succeed with far fewer examples than previously thought possible.

Technical Summary: ReAugment

Problem Statement

Few-shot time series forecasting faces a critical challenge: deep learning models, despite their architectural advancements (e.g., Transformers), suffer from severe overfitting when trained on limited, high-quality labeled data. This scarcity leads to poor generalization on unseen future sequences. Traditional data augmentation techniques (e.g., Gaussian noise, scaling) are often "static" and heuristic, failing to account for the specific distribution of a dataset. Furthermore, existing automated augmentation methods are typically decoupled from downstream forecasting performance, meaning they do not explicitly target the specific regions where a model struggles to generalize.

Methodology: ReAugment

The authors propose ReAugment, a closed-loop, reinforcement learning (RL) framework designed to adaptively identify and augment critical samples to improve generalization. The system operates on the core insight that samples where different forecasting models disagree (high prediction diversity) represent "bottlenecks" for generalization and are most prone to overfitting.

The framework follows an "identify-augment-forecast" pipeline consisting of three main stages:

1. Anchor Points Identification (Model Zoo)

To locate the "overfit-prone" samples without prior knowledge of the test distribution, ReAugment constructs a Model Zoo:

Construction: The training set is partitioned into $K$ disjoint folds. A forecasting model is trained on $K-1$ folds for each fold, creating an ensemble of models where each has a unique "blind spot" regarding one specific subset of data.
Quantification: For any training sample, the Model-Zoo Variance ( $V(x)$ ) is calculated based on the discrepancy in prediction errors across the ensemble.
Selection: Samples with high variance are identified as "anchor points." These represent regions where the model's empirical risk minimization fluctuates significantly, signaling a high risk of overfitting.

2. Generative Initialization (VMAE)

A Variational Masked Autoencoder (VMAE) is trained on the identified overfit-prone anchor points to learn a robust base distribution.

Architecture: The VMAE includes a Prior Network (estimating latent distribution from masked data and timestamps), a Posterior Module, an Encoder, and a Decoder.
Objective: It minimizes the Evidence Lower Bound (ELBO), balancing reconstruction loss (fidelity) with KL divergence (regularization). This provides a stable starting point for the RL agent.

3. Policy Finetuning via Reinforcement Learning

The VMAE prior network is treated as a policy ( $\pi_\theta$ ) and finetuned using Group Relative Policy Optimization (GRPO).

Formulation: The process is modeled as a one-step contextual bandit.
- State: A masked time series and its timestamp.
- Action: Sampling a latent variable $z$ from the prior.
- Reward: A scalar feedback derived from the Model Zoo's predictive discrepancy. The reward function encourages generating samples that are "difficult" (high model-zoo variance) yet "realistic" (low reconstruction error).
Optimization: GRPO is used to compute a relative advantage by normalizing rewards within a group of sampled actions. This avoids the gradient vanishing problem associated with backpropagating through the entire model zoo and handles varying reward scales across different models.

Key Contributions

Closed-Loop Framework: ReAugment is a novel method that directly aligns data augmentation objectives with forecasting performance, creating a feedback loop where augmentation targets specific generalization bottlenecks.
Model-Zoo-Guided Mechanism: The paper introduces a method to locate anchor points within the training distribution that are most susceptible to overfitting by leveraging prediction diversity across an ensemble of models, rather than relying on static heuristics.
Gradient-Friendly Policy Optimization: By formulating augmentation as an RL task using GRPO, the framework optimizes the generative policy using forecasting performance as a reward signal without requiring backpropagation through the forecasting models, thereby mitigating gradient vanishing.
Empirical Validation: Extensive experiments demonstrate that ReAugment consistently outperforms existing heuristic and learning-based augmentation methods across various architectures (iTransformer, PatchTST, DLinear) and datasets.

Experimental Results

The authors evaluated ReAugment on five real-world benchmarks (ETT, Traffic, Electricity, Weather, Exchange) under both few-shot (10–20% training data) and standard settings.

Few-Shot Performance: ReAugment consistently achieved the lowest MAE and MSE across all benchmarks. For instance, on the ETTh2 dataset, it reduced errors significantly compared to the second-best method.
Model Agnosticism: The method improved performance when applied to diverse backbones, including Transformer-based models (iTransformer, PatchTST) and linear models (DLinear), suggesting the identified bottlenecks are data-level deficiencies shared across model families.
Recovery Factor: Using a new metric, the Augmentation Recovery Factor ( $F$ ), ReAugment showed positive recovery values (often approaching or exceeding 100%), indicating it effectively recovered the performance gap between few-shot and full-data training. In contrast, baselines like Gaussian noise or TimeGAN sometimes yielded negative recovery, indicating harmful noise.
Comparison with Foundation Models: Despite having significantly fewer parameters, the ReAugment-enhanced iTransformer outperformed large foundation models (Chronos-2, TimesFM) on several few-shot tasks, suggesting that targeted augmentation can be more effective than massive pre-training when domain gaps are large.
Efficiency: The training overhead for the augmentation stages (VMAE and GRPO) was significantly lower than training the forecasting backbone itself, and inference requires no additional computation.

Significance and Claims

The paper claims that ReAugment provides a principled approach to few-shot time series forecasting by shifting the focus from "augmenting everything" to "augmenting the bottlenecks." By explicitly identifying regions of high prediction diversity (where models disagree) and using RL to synthesize data specifically for those regions, the method forces the model to learn more robust temporal patterns.

The authors emphasize that their approach offers a complementary perspective to the scaling of foundation models: rather than relying solely on broad pre-training coverage, ReAugment explicitly adapts scarce target-domain training sets by synthesizing samples around regions where the forecasting model is most sensitive to data exposure. The work concludes that targeted, task-aware augmentation is a viable and efficient strategy for mitigating overfitting in data-scarce scenarios.