Technical Summary: Flow-Corrected Thompson Sampling for Non-Stationary Contextual Bandits

Problem Formulation

The paper addresses the challenge of non-stationary linear contextual bandits, a setting where the relationship between contexts, actions, and rewards evolves over time. In classical contextual bandits, the parameter vector $w_t$ is assumed stationary. However, in real-world applications (e.g., recommendation systems, dynamic pricing, finance), user preferences and market conditions drift, rendering historical data systematically biased if used naively.

The core difficulty lies in the trade-off between retaining historical data to reduce posterior variance and discarding old data to avoid bias from outdated models. Existing approaches—such as sliding windows, exponential discounting, and periodic restarts—treat historical data primarily as a liability to be managed by uniformly discarding or down-weighting it. The authors argue that when non-stationarity exhibits predictable structure (e.g., gradual drift, periodicity, or recurring regimes), historical observations contain valuable information that can be leveraged if corrected for the intervening environmental changes.

Methodology: Flow-Corrected Thompson Sampling (FC-TS)

The authors propose Flow-Corrected Thompson Sampling (FC-TS), a Bayesian framework that transforms historical observations into "present-time pseudo-observations" before incorporating them into the posterior. The method decouples the handling of non-stationarity into two distinct components: a transport operator and a confidence weighting rule.

1. Core Mechanism

At each round $t$, for every past observation $(x_s, a_s, r_s)$, FC-TS constructs a transported reward $\hat{r}_{s \to t}$:
$$\hat{r}_{s \to t} = r_s + \hat{\Delta}_{s \to t}(x_s, a_s)$$
where $\hat{\Delta}_{s \to t}$ estimates the shift in the mean reward from time $s$ to $t$. This effectively "flows" the past reward forward to align with the current reward model.

To prevent noise from inaccurate drift estimates, each transported observation is assigned a confidence weight $\omega_{s,t} \in [0, 1]$. This weight controls the influence of the observation on the posterior, inflating its effective noise variance if the correction is uncertain.

The posterior for action $a$ is updated using these weighted, transported observations:
$$\Lambda_t^{(a)} = \lambda I + \sigma^{-2} \sum_{s<t, a_s=a} \omega_{s,t} x_s x_s^\top$$
$$\eta_t^{(a)} = \lambda m_t^{(a)} + \sigma^{-2} \sum_{s<t, a_s=a} \omega_{s,t} \hat{r}_{s \to t} x_s$$
$$\mu_t^{(a)} = (\Lambda_t^{(a)})^{-1} \eta_t^{(a)}$$
Action selection proceeds via standard Thompson sampling using the Gaussian posterior $\mathcal{N}(\mu_t^{(a)}, (\Lambda_t^{(a)})^{-1})$.

2. Structured Instantiations

The framework provides a unified template with specific instantiations for three common non-stationary patterns:

• Linear Drift: Assumes parameters evolve as $w_t^{(a)} \approx \theta^{(a)} + t \delta^{(a)}$. The transport operator uses an online ridge regression estimate $\hat{\delta}_t^{(a)}$ to shift rewards: $\hat{r}_{s \to t} = r_s + (t-s)\langle \hat{\delta}_t^{(a)}, x_s \rangle$. Weights decay exponentially ($\omega_{s,t} = \gamma^{t-s}$) to account for extrapolation uncertainty.
• Periodic Variation: Assumes parameters repeat with period $P$. The transport operator is identity ($\hat{r}_{s \to t} = r_s$) for observations in the same phase. Weights are determined by the circular distance between phases, allowing the algorithm to reuse data from previous cycles ($\omega_{s,t}$ depends on phase proximity).
• Recurring Regimes: Handles environments switching between a finite set of regimes. The algorithm maintains a bank of regime-specific posteriors. Upon detecting a changepoint (via CUSUM on prediction residuals), it reloads the posterior of the re-occurring regime rather than restarting from scratch.

The framework also supports composition, allowing simultaneous handling of linear drift and periodicity by combining the respective transport and weighting rules.

Key Contributions

1. Flow-Corrected Posterior Updates: The paper introduces a transport-based view of non-stationarity where past rewards are corrected into present-time pseudo-observations and incorporated via confidence-weighted Bayesian updates, rather than being uniformly discarded.
2. Modular Algorithm Design: FC-TS is defined by a generic transport operator and weighting rule, instantiated for linear drift, periodic variation, and recurring regimes. This includes an adaptive variant combining trends and multiple periodicities.
3. Empirical Validation: Extensive evaluation across five case studies (linear drift, periodic variation, recurring regimes, compound non-stationarity, and semi-synthetic portfolio selection) demonstrates that FC-TS outperforms state-of-the-art forgetting-based baselines.

Experimental Results

The authors evaluated FC-TS against four baselines: Standard TS (STD-TS), Sliding Window TS (SW-TS), Discounted TS (D-TS), and Restarting TS (R-TS).

• Performance Gains: FC-TS achieved the lowest cumulative regret in all five settings, with improvements ranging from 14.1% to 59.0% over the strongest baseline in each specific environment.

• Case-Specific Insights:

  • Linear Drift: FC-TS matched SW-TS during the warm-up phase but progressively separated as drift estimates stabilized, winning 8 of 10 seeds.
  • Periodic Variation: FC-TS showed the largest synthetic gain (50.0%) due to "phase memory," accumulating data from previous cycles. Regret per cycle dropped from ~70 to ~4, while baselines plateaued.
  • Recurring Regimes: FC-TS achieved a 59.0% improvement by reloading regime-specific posteriors upon detection of a switch, avoiding the "cold start" penalty of restarts.
  • Compound Non-Stationarity: In environments with both drift and abrupt changepoints, FC-TS successfully combined drift correction (for smooth evolution) and regime resets (for abrupt changes), yielding a 29.5% improvement.
  • Semi-Synthetic Portfolio: In a financial task with overlapping monthly/quarterly cycles and growth-to-value drift, the adaptive FC-TS variant improved regret by 28.1% over the best baseline (D-TS). Ablation studies confirmed that drift correction was the dominant component of this gain.

• Ablation Findings:

  • Effective Sample Size: Retaining more data without transport correction (the "no-transport" ablation) resulted in higher regret than methods that discard data, confirming that stale data introduces systematic bias.
  • Transport Value: The transport operator itself accounted for a 47–57% regret reduction compared to the no-transport variant across various drift scales, proving that the correction mechanism, not just the posterior rebuild, drives performance.

Significance and Claims

The paper claims that structured non-stationarity is better addressed by correcting and reusing experience than by uniformly discarding it.

The authors emphasize that when the environment's drift is predictable (gradual, periodic, or recurring), historical data is an asset rather than a liability. By explicitly modeling the "flow" of the reward distribution over time, FC-TS maintains a much larger effective sample size than forgetting-based methods. The results suggest that the key to sample efficiency in non-stationary bandits lies in reweighting and correcting historical observations based on the specific temporal structure of the environment, rather than applying a single, uniform forgetting rule.

The paper concludes modestly, noting that future work includes learning drift structures online with calibrated uncertainty, extending the framework to generalized linear or neural bandits, and developing safeguards for high-stakes settings with unmodeled distribution shifts.