Off-Policy Learning in Large Action Spaces: Optimization Matters More Than Estimation

Technical Summary: Off-Policy Learning in Large Action Spaces: Optimization Matters More Than Estimation

Problem Statement
Off-policy learning (OPL) in offline contextual bandits aims to learn a policy $\hat{\pi}$ from logged data $D_n = \{(X_i, A_i, R_i)\}$ that maximizes the expected reward $V(\pi)$. The dominant paradigm involves optimizing an off-policy evaluation (OPE) estimator, typically based on Inverse Propensity Scoring (IPS), such as $\hat{\pi} = \arg\max_\pi \hat{V}_{IPS}(\pi)$. While recent research has focused on developing sophisticated estimators with improved statistical properties (e.g., reduced variance or bias) to approximate $V(\pi)$ more accurately, this paper argues that this estimator-centric approach overlooks a critical practical obstacle: the optimization landscape.

In large action spaces (where the number of actions $K$ can reach millions), IPS-based objectives paired with standard policy parametrizations (e.g., softmax) exhibit severe optimization pathologies. These include long plateaus where gradient ascent stalls and an exponential number of poor local maxima. The paper posits that even estimators with superior statistical guarantees fail to yield good policies because the underlying optimization problem becomes intractable as $K$ increases.

Methodology and Analysis
The authors provide both theoretical analysis and empirical evidence to support the claim that optimization stability is more critical than estimation accuracy in large action spaces.

1. Theoretical Analysis of IPS-Based Objectives:

   • The paper analyzes "oracle policies" (the theoretical maximizers of the expected objective) for various IPS-based estimators, including standard IPS, Clipped IPS (cIPS), Exponential Smoothing (ES), Doubly Robust (DR), Marginalized IPS (MIPS), OffCEM, POTEC, and Policy Convolution (PC).
   • Optimization Pathologies: The authors prove that for linear-in-$\pi$ estimators with softmax policies, gradient ascent can remain trapped in suboptimal regions for $O(K_{eff})$ iterations, where $K_{eff}$ is the size of the effective action space. Furthermore, the landscape contains a number of local maxima exponential in $K_{eff}$.
   • The Role of $K_{eff}$: Even methods designed for large action spaces (like MIPS or OffCEM) often fail to mitigate this if they optimize over the full action space $A$ ($K_{eff} = K$). The only exception identified is POTEC, which restricts optimization to a cluster level, effectively reducing $K_{eff}$.
   • Objective-Aware Parametrization: The authors propose aligning the policy class with the estimator's inductive bias. For instance, restricting the policy support to the logging policy's support $S_0$ for IPS-based methods, or using a two-stage decomposition for cluster-based methods, can reduce $K_{eff}$ and improve trainability without changing the statistical target. However, this only mitigates symptoms; it does not cure the fundamental non-concavity of IPS-based objectives.

2. Policy-Weighted Log-Likelihood (PWLL) Objectives:

   • As a fundamental alternative, the paper explores PWLL objectives, which take the form $\hat{U}_g(\pi) = \frac{1}{n} \sum g(R_i, \pi_0(A_i|X_i)) \log \pi(A_i|X_i)$.
   • Concavity: Unlike IPS objectives, PWLL objectives are inherently concave (and strongly concave with $\ell_2$ regularization) when using linear-softmax policies. This is because the map from scores to log-probabilities is concave.
   • Optimization Guarantees: The authors establish that stochastic gradient ascent on PWLL objectives achieves global convergence rates of $O(1/\sqrt{T})$ generally, accelerating to geometric rates with regularization.
   • Variants: The study evaluates specific PWLL variants including Local Policy Improvement (LPI), Clipped LPI (cLPI), and Regularized KL (RegKL). Notably, cLPI is shown to induce the same greedy decision rule as cIPS but via a convex optimization problem.

Key Contributions

1. Optimization Theory Scaling with Action Space: The paper formalizes how optimization pathologies (plateaus and local maxima) scale with the number of actions $K$, explaining why IPS-based objectives become increasingly difficult to optimize as $K$ grows.
2. Objective-Aware Parametrization: The authors identify minimal sufficient policy structures for various objectives (derived from oracle policies) and demonstrate that restricting or factorizing the policy to match these structures improves trainability and performance.
3. PWLL as an Optimization-Friendly Objective: The paper demonstrates that for linear-softmax policies, PWLL objectives yield well-behaved, concave optimization landscapes, offering a robust alternative to direct value maximization.
4. Large-Scale Empirical Stress Tests: Extensive experiments on three large-scale datasets (MovieLens with $K=60k$, Twitch with $K=200k$, and GoodReads with $K=1M$) show that IPS-based methods are highly sensitive to hyperparameters (batch size, learning rate schedules), whereas PWLL methods are robust and consistently achieve higher rewards.
5. Actionable Guidance: The paper distills design rules for large action spaces: prioritize optimization stability by using objective-aware parametrizations and prefer PWLL-style objectives over complex IPS estimators for policy learning.

Results

• Optimization Sensitivity: In controlled sweeps, IPS-based objectives (including state-of-the-art methods like POTEC) showed high variance and performance collapse under suboptimal hyperparameters. In contrast, PWLL methods (specifically cLPI) maintained high performance across all configurations.
• Estimation vs. Learning: The paper highlights a counter-intuitive finding: methods with lower Mean Squared Error (MSE) as value estimators (e.g., MIPS) did not necessarily produce better policies. Conversely, PWLL methods, which are poor value estimators (high MSE), produced the best policies. This suggests that in large action spaces, the ability to optimize the objective is a stronger predictor of policy quality than the statistical accuracy of the value estimate.
• Parametrization Impact: Restricting the effective action space to match the objective's inductive bias (e.g., using support-restricted policies for IPS) significantly improved the stability and performance of IPS methods, though they still lagged behind PWLL methods.
• Robustness: PWLL methods demonstrated superior robustness to reward noise and different logging policy supports compared to IPS-based approaches.

Significance and Claims
The paper claims that the field of off-policy learning has been overly focused on refining estimators while neglecting the optimization landscape, a critical bottleneck in large action spaces. The authors argue that:

• Optimization is the bottleneck: In large action spaces, the difficulty of optimizing non-concave IPS objectives undermines the practical utility of even the most statistically sophisticated estimators.
• Simplicity wins: Simpler, optimization-friendly objectives like PWLL (specifically cLPI) outperform complex, variance-reduced IPS estimators because they allow for stable and reliable convergence.
• Design Shift: The development of OPL algorithms for large action spaces must explicitly address optimization considerations. The paper advocates for a shift from "estimator-centric" design to "optimization-centric" design, suggesting that ensuring a tractable optimization landscape is more important than achieving high-fidelity value estimation for the purpose of learning a policy.

The authors note limitations: PWLL objectives are not suitable for off-policy evaluation (OPE) or policy selection where accurate value estimates are required. Additionally, the concavity guarantees for PWLL hold strictly for linear-softmax policies; while deep networks may lose formal concavity, the gradient structure remains more benign than that of IPS-based objectives.