PaperlayerView exact revision on arXiv

Source-linked AI summary

Bandits via Additive Quantized Representations

Ami Tavory, Noam Touitou, Tal Sarig, Frank Cheng, Ido Guy

arXiv:2610.02440v1cs.LG

TL;DR

Contextual bandits face a gap between nonlinear models that require growing replay memory and linear models that have bounded memory but restricted reward structure. This paper uses offline-trained Residual Quantization codebooks for additive bandit models with fixed memory. Across datasets, RQ variants improve on most within-class baselines while LinTS-RQ is competitive with XGBoost at 1000× less memory.

  • Problem

    Tree and neural bandit methods model nonlinear rewards but rely on replay buffers that can grow with the horizon, while fixed-memory linear methods are structurally restricted to linear rewards.

    Source: §1 · p. 1

  • Method

    The paper uses an offline-trained Residual Quantization codebook as a representation layer for additive bandit models with fixed O(1) memory and no replay buffer.

    Source: §1 · pp. 1–2

  • Results

    RQ methods improve over their within-class baselines on 11 of 13 datasets, and LinTS-RQ is competitive with XGBoost (7W/4L/2T) at 1000× less memory.

    Source: §1 · p. 2

  • Takeaways & Limitations

    RQ variants improve on most datasets at orders-of-magnitude less memory than replay-based methods, while TS-RQ remains robust when LinTS-RQ’s covariance starves.

    Source: §6 · p. 9

    LinTS-RQ can degrade when data per parameter is insufficient; on year_prediction, its N/(bKd^2) ratio is 0.04 and performance degrades even at depth 1.

    Source: §7 · p. 9

Abstract

from arXiv · show

Contextual bandits require balancing nonlinear reward modeling with online efficiency. Tree ensembles and neural methods capture nonlinearities but require periodic retraining and large replay buffers. Linear models update efficiently per observation with O(1) memory, but are fundamentally restricted to linear reward structures. We propose Residual Quantization (RQ) as a representation layer to bridge this gap. An offline-trained RQ codebook maps continuous contexts into discrete centroid assignments across multiple levels, set dynamically through a shadow mechanism. This enables a spectrum of additive bandit algorithms that achieve nonlinear expressivity with strictly bounded memory. Across 13 datasets, RQ variants beat their non-RQ counterparts on 11 of 13 datasets, often by wide margins, while matching doubling-retrain XGBoost and neural baselines using up to 1000 times less memory.

1 Introduction

Contextual bandits need nonlinear reward models without the replay memory and retraining of tree and neural methods, while linear models remain restricted to linear rewards. The paper uses Residual Quantization (RQ) for additive bandit models with fixed memory and reports broad improvements over within-class baselines.

1 Introduction

  • Tree and neural methods capture nonlinear rewards but use replay buffers that can grow with the horizon, whereas linear methods have fixed memory but are restricted to linear rewards.At five million rounds on a 28-feature dataset, a doubling-retrain bandit uses over 1 GB of replay memory.Source: §1 · p. 1
  • RQ codebooks trained by unsupervised k-means map contexts across levels, yielding bℓ effective cells with ℓb parameters and fixed O(1) online memory.The codebook is trained on unlabeled contexts and frozen before the bandit starts; online learning needs no replay buffer or retraining.Source: §1 · pp. 1–2
  • A shadow model trains at the next depth without affecting arm selection and is promoted at pre-committed checks only if a statistical test finds improved prediction.Otherwise, the method remains at its initial depth.Source: §1 · p. 1
  • Under three base-learner convergence properties, asymptotic regret is bounded by additive misspecification, and the shadow mechanism has false-promotion probability at most α.The analysis also decomposes depth-dependent approximation error and proves geometric convergence to an irreducible floor under geometric contraction of reconstruction and additive-tail errors.Source: §1 · p. 2
  • RQ variants improve over their within-class baselines on 11 of 13 datasets, and LinTS-RQ is competitive with XGBoost (7W/4L/2T) at 1000× less memory.Experiments cover 13 datasets, up to 11M rounds, and 30 seeds each.Source: §1 · p. 2, Figure 2 · p. 3

2 Method

The method freezes an RQ codebook trained on a separate pool of unlabeled contexts, then fits additive bandit learners over its levels. A shadow level can be promoted by a prediction-accuracy test, while the architecture’s memory–regret comparison contrasts RQ with linear and doubling methods.

2 Method

  • RQ levels partition contexts into cells, while the codebook is trained once by sequential k-means on a separate pool of unlabeled contexts and frozen before bandit learning.Greedy assignments produce a code sequence and residual contexts used as per-level feature inputs.Source: §2 · p. 3
  • Each level fits the reward prediction on its residual; deeper levels predict the gap between reward and the cumulative prediction of shallower levels.A damping factor controls each level’s contribution.Source: §2 · p. 3
  • The variants share the same outer loop and differ only in the per-cell learner.LinTS-RQ pseudo-targets shrink at deeper levels, with regularization and numerical scaling addressed separately.Source: §2 · p. 3
  • A shadow level trains alongside the active model without affecting arm selection; at pre-committed checks, it is promoted if a statistical test finds better prediction accuracy.Otherwise, the active depth is unchanged.Source: §2 · p. 3
  • RQ methods occupy the fixed-memory, low-regret region of the memory–quality comparison, matching tree and neural regret at 1000× less memory.Each plotted point represents a method–dataset pair; the figure contrasts RQ and linear methods at O(1) memory with doubling methods that use O(T) replay buffers.Source: Figure 3 · p. 4

3 Analysis

The analysis bounds asymptotic regret through additive misspecification, studies how depth reduces that error, and controls false promotions by the shadow mechanism.

3.1 Asymptotic regret

  • Asymptotic regret is bounded by the stagewise approximation error, with equality to twice the additive misspecification when stagewise and joint optima coincide.The guarantee assumes the base learner satisfies three convergence properties; with nonzero misspecification, regret may grow linearly at the residual error rate.Source: §3.1 · p. 5

3.2 Misspecification and depth

  • Depth-dependent error separates into saturated quantization error and additive-interaction error; geometric decay requires both reconstruction and additive-tail errors to contract geometrically.The additive readout uses O(ℓb) parameters even as the saturated benchmark has b^ℓ cells.Source: §3.2 · pp. 5–6
  • Across 39 dataset–method pairs, median R2 reaches 0.99 for exponential fits, consistent with—but not proof of—the conditional geometric depth-rate law.Within-cell reward variance contracts with median per-level ratio 0.83.Source: §3.2 · p. 6, Figure 4 · p. 6

3.3 Shadow false-promotion control

  • The shadow is falsely promoted with probability at most α over the full horizon when it has no lower expected squared prediction error under the active policy’s tested-window distribution.Freedman’s martingale inequality avoids an i.i.d. requirement, and Bonferroni correction distributes significance across the checks.Source: §3.3 · p. 6

4 Experiments

Across 13 datasets, RQ methods improve on their within-family baselines on most tasks, while LinTS-RQ is competitive with replay-based nonlinear baselines at substantially lower memory. Results also expose data-sufficiency and model-specific conditions that shape performance.

4.1 Setup

  • The evaluation uses 13 tabular datasets, a 50K-context unlabeled codebook-training holdout, 30 seeds, and fixed hyperparameters without per-dataset tuning.Regret rate is cumulative regret divided by T; win/loss calls use 95% confidence intervals.Source: §4.1 · p. 7

4.2 Results

  • Doubling-XGB and Doubling-NeuralCB match LinTS-RQ regret but require O(T) replay memory.Regret is normalized by the context-free oracle rate, with 1.0 corresponding to majority-class guessing.Source: §4.2 · p. 7
  • LinTS-RQ wins on 10 of 13 datasets, SGD-LinTS-RQ on 12 of 13, and TS-RQ on 11 of 13 versus their within-family baselines.Table 3 summarizes win/loss/tie outcomes across 13 datasets using 30 seeds and 95% confidence intervals.Source: §4.2 · p. 7, Table 3 · p. 8
  • At T=5M, LinTS-RQ uses 620 KB versus XGB’s 1.1 GB while remaining competitive with Doubling-XGB (7W/4L/2T).Against TabNet-LinTS, LinTS-RQ wins on 11 of 13 datasets (11W/2L).Source: §4.2 · p. 8
  • On 12 of 13 datasets, N/(bKd^2) exceeds 1.1 and LinTS-RQ wins; on year_prediction, the ratio is 0.04 and LinTS-RQ degrades even at depth 1.TS-RQ matches the TS baseline on year_prediction (93.0% vs 93.4%), and Figure 5 shows year_prediction as the sole real-dataset loss.Source: §4.2 · p. 8, Figure 5 · p. 8
  • On skin_segmentation and miniboone, TS-RQ achieves 2.6% and 14.7% normalized regret, respectively, where low effective rank hurts SGD-LinTS.The passage attributes the diagonal approximation’s weakness to allocating parameters across noisy coordinates, while TS-RQ captures dominant directions through centroid placement.Source: §4.2 · p. 8
  • On covertype, shadow promotion discovers depth 4 and gains 3.4 percentage points over fixed depth 2.On higgs, it settles at depth 3 because the marginal depth-4 gain falls within the Freedman test’s noise floor.Source: §4.2 · p. 9, Table 4 · p. 9

5 Related Work

Related work spans linear and nonlinear bandits, adaptive partitioning, model selection, and representation compression. The paper positions residual quantization as combining additive representations with sequential depth promotion.

Linear and GLM bandits

  • LinTS and LinUCB attain O(d√KT log T) regret under linear rewards, while misspecification can yield a linear regret bound.The cited misspecification bound from Lattimore et al. (2020) underlies the paper’s asymptotic analysis.Source: §5 · p. 9

Neural and tree-based bandits

  • Neural and tree-based bandits model nonlinear rewards but can require per-step gradient and inverse-Hessian computation or refitting on accumulated history.The passage contrasts these costs with the paper’s offline codebook, which avoids gradients and refitting.Source: §5 · p. 9

Partition-based bandits

Model selection and quantization

  • Unlike CORRAL and Regret Balancing, shadow promotion monitors one additional depth sequentially without influencing arm selection.The passage also describes RQ as compressing bℓ cells into ℓb parameters and constructing its codebook through unsupervised k-means.Source: §5 · p. 9

6 Conclusion

The conclusion presents RQ codebooks as a fixed-memory contextual-bandit representation and reports gains on most datasets. It also identifies robustness when LinTS-RQ lacks sufficient covariance data and outlines extensions.

Conclusion

  • RQ codebooks are repurposed as a contextual-bandit representation layer for gradient-boosted reward models with fixed O(1) memory.The conclusion describes RQ variants as improving on most datasets at orders-of-magnitude less memory than replay-based methods.Source: §6 · p. 9
  • TS-RQ remains robust when LinTS-RQ’s covariance starves.The conclusion names continuous rewards, online codebook adaptation, and alternative per-level features such as PCA and neural-linear as extensions.Source: §6 · p. 9
  • The conclusion identifies online codebook adaptation and alternative per-level features as possible extensions.Examples given include PCA and neural-linear features.Source: §6 · p. 9

7 Limitations

The limitations section identifies data sufficiency for full-covariance LinTS-RQ, offline codebook training, and assumptions behind the asymptotic analysis as boundaries on the method’s scope.

Data sufficiency

  • LinTS-RQ’s data-sufficiency condition is N/(bKd^2) ≫ 1; on year_prediction, with ratio 0.04, it degrades even at depth 1.A p-dimensional projection would change the condition to N/(bKp^2).Source: §7 · p. 9

Analysis and reward limitations

  • The fixed-depth asymptotic analysis is conditional on base-learner convergence properties, and post-promotion performance depends on the new policy’s data distribution.The section also notes that linear variants use squared loss for binary rewards and that reward-signal biases may be amplified, though the method neither exacerbates nor mitigates them.Source: §7 · pp. 9–10

Codebook limitations

  • The unsupervised k-means codebook minimizes reconstruction rather than reward-prediction error, and its offline holdout requirement may cause degradation in cold-start settings.A reward-aware partition could reduce misspecification directly but would require labeled data for codebook construction.Source: §7 · p. 10

8 Broader Impact

The broader-impact statement presents the work as a general-purpose bandit representation, not one aimed at a specific application domain. The accompanying appendix describes shared cell operations and verifies convergence properties for the three base learners.

8 Broader Impact

  • The work targets no specific application domain and identifies no risks specific to residual quantization, while noting standard fairness and privacy concerns for bandit algorithms.Source: §8 · p. 10

A Base Learner Details

  • Each cell exposes sample, mean, and update operations; linear variants append an intercept to context residuals to ensure feature norms are at least 1.The lower bound is needed for the exploration guarantee, including when deep-level residuals are small.Source: A Base Learner Details · p. 13

B Base Learner Properties

  • The base-learner properties require eventual bounded score means, positive variance under starvation, and prediction consistency under infinite visits and vanishing target perturbations.Source: B Base Learner Properties · p. 14

B.1 Verification for LinTS-RQ

  • For LinTS-RQ, bounded targets and finite updates bound estimates; with infinitely many visits, predictions converge, including in prediction space when a centroid covariance is rank-deficient.The rank-deficient case is handled because features for that centroid lie in the covariance’s column space.Source: §B.1 · pp. 14–15

B.2 Verification for SGD-LinTS-RQ

  • SGD-LinTS-RQ verifies the properties coordinate-wise, but its coordinate-wise optimum can differ from the full-linear optimum when features are correlated.Source: §B.2 · p. 16

B.3 Verification for TS-RQ

  • TS-RQ is a scalar model using only a count and sum; its adaptive-sampling argument requires every centroid to be visited infinitely often, while a fully formal treatment under misspecified adaptive TS remains open generally.Source: §B.3 · p. 16

C Fixed-Depth Regret Analysis

For fixed depth, the analysis relates asymptotic regret to the converged sequential residual-fitting predictor and establishes exploration and prediction convergence under base-learner properties. It also proves a depth-rate result, with geometric convergence under geometric error-contraction assumptions.

C Fixed-Depth Regret Analysis

  • The asymptotic regret result applies to base learners satisfying Properties 1–3; its bound is at least 2∆ℓ(M), with equality when stagewise and joint optima coincide.Source: C Fixed-Depth Regret Analysis · p. 16
  • Every arm visits each positive-probability centroid infinitely often, and its level predictions converge under the stated base-learner properties.The proof proceeds inductively because deeper-level targets depend on estimates from shallower levels.Source: C Fixed-Depth Regret Analysis · p. 17
  • After predictions converge, the per-round regret is governed by the approximation error of the sequential residual-fitting limit, which is at least the best additive approximation error.The two errors are equal when the stagewise and joint optima coincide.Source: C Fixed-Depth Regret Analysis · p. 19

C.1 Proof of the depth-rate result

  • The depth-rate proof combines approximation-error components, including reconstruction and interaction-tail terms, to establish the stated depth bound.Source: §C.1 · pp. 19–20
  • Under the two geometric assumptions, the resulting bound has rate ρ = max{√ρx, ρτ}.Source: §C.1 · p. 20

D Promotion Proofs

The shadow promotion analysis controls false promotions under the active policy’s data distribution and accounts for the number of tests. It distinguishes that pre-promotion guarantee from post-promotion behavior, when the policy and data distribution can change.

D Promotion Proofs

  • The shadow trains without influencing arm selection, and the formal promotion test uses Freedman’s inequality without an i.i.d. assumption on paired prediction errors.The appendix notes that the paired t-test used as a practical heuristic is not formally calibrated for the evolving policy’s data.Source: D Promotion Proofs · p. 20
  • The shadow’s test compares paired squared-error differences over a trailing window, with positive differences indicating closer predictions by the shadow.Source: D Promotion Proofs · p. 20
  • Testing only the current depth frontier makes the multiplicity penalty additive O(nc + ℓmax), with nc = O(log T) under the doubling schedule.Source: D Promotion Proofs · p. 22
  • With probability at least 1 − α, every promotion corresponds to a true improvement in expected squared prediction error under the active policy’s tested data distribution.Source: D Promotion Proofs · p. 22
  • The guarantee is pre-promotion: after promotion, the policy changes and the data distribution shifts, so the guarantee does not establish post-promotion prediction performance.Source: D Promotion Proofs · p. 22

E.1 Exponential Decay of Misspecification

The appendix examines geometric depth behavior, adaptive depth selection, and supporting benchmark and implementation details. It also presents a full-covariance analytical variant and per-dataset comparisons, while the supplied passages do not report the table cell outcomes.

E.1 Exponential Decay of Misspecification

  • The exponential regret-decay fit has median R2 0.99, and held-out-depth predictions have median error 0.3%, with 73% within 1% of observed values.The fit evaluates R(T)/T against depth; Figure 7 reports predicted-versus-actual values and error by held-out depth.Source: §E.1 · p. 23, Figure 7 · p. 23

E.2 Within-Cell Reward Variance

  • Across all 13 datasets, within-cell reward variance decays roughly linearly on a log scale, with median per-level ratio ρ = 0.83 (range 0.24–0.93).The variance is weighted across cells defined by the first ℓ centroid assignments.Source: §E.2 · p. 24, Figure 9 · p. 24

E.3 Fixed-Depth vs. Shadow Promotion

  • On covertype, shadow promotion reaches depth 4 and gains 3.4 percentage points over fixed depth 2; on higgs, it settles at depth 3 with regret matching fixed depth 2.Source: §E.3 · p. 24

F Analytical Variant: Finite-Time Bound

  • The analytical variant uses a full joint posterior covariance; its RQ additive structure has dimension D = bℓp while retaining the approximation power of bℓ cells.Its finite-time bound complements the practical algorithm’s asymptotic guarantee; the supplied analysis cautions that the misspecified Thompson Sampling decomposition is heuristic.Source: F Analytical Variant: Finite-Time Bound · pp. 25, 27

G Dataset Details

  • The evaluation uses publicly available OpenML and UCI datasets, with codebooks trained on 50K held-out contexts; preprocessing label-encodes categorical features without feature selection, imputation, or augmentation.Source: G Dataset Details · p. 27, Table 6 · p. 28

H Full Per-Dataset Results

I Per-Dataset Regret Curves

  • Figure 10 plots regret rate against horizon T for all nine methods on each dataset, with shaded bands showing ±1 standard deviation across seeds.Source: Figure 10 · p. 30

J Regret Scaling Analysis

Across datasets, TS-RQ regret grows nearly linearly with horizon, consistent with the additive misspecification bias term dominating exploration at tested horizons. The analysis therefore points to reducing approximation bias as more impactful than reducing exploration cost.

J Regret Scaling Analysis

  • At large T, the O(√T log T) exploration term becomes negligible relative to O(∆ℓT), as Theorem 3.1 predicts.The paper identifies reducing ∆ℓ through deeper codebooks or richer base learners as more impactful than reducing exploration cost.Source: J Regret Scaling Analysis · p. 31
  • Across all 13 datasets, fitted regret exponents were α ∈[0.87, 1.00], indicating regret was overwhelmingly dominated by the linear bias term at experimental horizons up to T=11M.Source: J Regret Scaling Analysis · p. 30
  • Regret rate converges to a constant on selected datasets, consistent with the O(∆ℓT) bias term in Theorem 3.1.The direct fit separates exploration and bias contributions.Source: Figure 11 · p. 31, J Regret Scaling Analysis · p. 31

K Data Sufficiency: Synthetic Validation

The appendix tests when RQ bandits have enough data per parameter, examines failures and comparisons across model classes, and reports timing and memory behavior. LinTS-RQ competes with replay-based nonlinear baselines, while data sufficiency and effective rank qualify which variants perform well.

K Data Sufficiency: Synthetic Validation

L year_prediction Depth Analysis

  • On year_prediction, LinTS-RQ worsens with depth (96.3% at ℓ=1 to 97.4% at ℓ=4), while TS-RQ remains at 93.0% across depths.The reported data-to-parameter ratio is 0.04; context-free TS reaches 92.7%, and the features carry negligible signal for the 84-way target.Source: L year_prediction Depth Analysis · p. 32

M Reconstruction Error

  • At depth 4, normalized reconstruction error ranges from 0.005 on skin_segmentation to 0.648 on year_prediction, whose feature geometry is poorly captured by 16 centroids in R90.The paper links year_prediction’s high error to its status as the sole LinTS-RQ failure.Source: M Reconstruction Error · p. 32, Table 10 · p. 33

N Effective Rank and SGD-LinTS Failure

O XGBoost and NeuralCB Comparison

P TabNet-LinTS Comparison

Q Timing

  • TS-RQ costs 50−750 µs per step, while LinTS-RQ on year_prediction (d=90, K=84) reaches 28 ms.TS-RQ is consistently the cheapest RQ variant.Source: Figure 16 · p. 36

R Memory Formulas

  • Memory scales with model class: TS-RQ stores 2ℓbK values, SGD-LinTS-RQ 2ℓbdK, and LinTS-RQ (d2+d)ℓbK; XGB replay grows linearly with T.Adapted depths vary by dataset: TS-RQ reaches ℓ=8, whereas LinTS-RQ generally stays at ℓ=2−3.Source: R Memory Formulas · p. 36

Related papers

Loading 2610.02440v1…