Bandits via Additive Quantized Representations

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

cs.LGOct 1, 2026

Bandits via Additive Quantized Representations

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

Organizations: Meta Platforms

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Partitioning the context space creates nonlinear reward models from linear per-cell learners, but finer partitions fragment the data: each cell sees fewer observations. All three schemes realize $b^{\ell}{=}64$ cells ( $b{=}4$ , $\ell{=}3$ ). Top/Bottom left: Flat and factorial parameterizations learn each cell from its own data alone. Bottom right: Additive RQ achieves the same $64$ -cell granularity with $\ell b{=}12$ shared parameters—each observation updates $\ell{=}3$ centroids, pooling information across cells. Fine partition without data starvation.

Figure 2: Regret gap (%) of each RQ method relative to its within-class baseline on 13 datasets (30 seeds, 95% CIs). Negative = RQ method has lower regret. Black: significant win (CI excludes zero). Gray: not significant. Red: significant loss. Left: LinTS-RQ vs LinTS wins on 11/13 datasets; the sole loss (year_prediction,

d{=}90

K{=}84

) is explained by the data sufficiency condition

N/(bKd^{2})\ll 1

(Section 4.2 ). Center: SGD-LinTS-RQ vs SGD-LinTS wins on 12/13. Right: TS-RQ vs TS wins on 11/13.

Figure 3: Memory vs. quality across 13 datasets. Each point is one (method, dataset) pair. Hollow gray: linear baselines (

O(1)

memory, high regret). Solid gray: RQ methods (

O(1)

memory, low regret). Red: doubling methods (

O(T)

replay buffers, low regret). RQ methods match tree/neural regret at

1000\times

less memory. The vertical drop from hollow to solid at the same

x

-position shows the regret gain from the RQ codebook.

Method	$\mathcal{M}$	Features used	Cost/round
TS-RQ	Count/sum	Index $k_{i}$ only	$O(\ell K)$
SGD-LinTS-RQ	Diagonal linear	Index $k_{i}$ + $[x^{(i-1)};1]\in\mathbb{R}^{d+1}$	$O(\ell dK)$
LinTS-RQ	Full-covariance LinTS	Index $k_{i}$ + $[x^{(i-1)};1]\in\mathbb{R}^{d+1}$	$O(\ell d^{2}K)$

Table 1: Model variants sharing Algorithm 1 , differing only in per-cell learner

\mathcal{M}

(full pseudocode in Appendix A ).

Figure 4: Empirical behavior consistent with the conditional geometric depth-rate law. Left: Median

R^{2}

of the exponential fit

a+b\rho^{\ell}

improves with horizon

T

. Right: Normalized regret decay

(\Delta_{\ell}-\Delta_{5})/(\Delta_{1}-\Delta_{5})

at terminal

T

for each dataset; curves cluster along exponential trajectories.

	TS-RQ	LinTS-RQ	XGB replay
higgs ( $T{=}5\text{M}$ )	12 KB	620 KB	1.1 GB
susy ( $T{=}5\text{M}$ )	22 KB	176 KB	800 MB
covertype ( $T{=}530\text{K}$ )	60 KB	10 MB	51 MB

Table 2: Model memory at shadow-promoted depths (

b{=}16

). XGB replay buffer grows

O(T)

; all RQ methods are fixed

O(1)

Comparison	W/L/T	Best win	Worst loss
LinTS-RQ vs LinTS	10W/1L/2T	skin_seg $-96\%$	year_pred $+4.5\%$
SGD-RQ vs SGD-LinTS	12W/1L	skin_seg $-97\%$	year_pred $+0.7\%$
TS-RQ vs TS	11W/2L	kddcup99 $-97\%$	year_pred $+0.4\%$

Table 3: Win/loss/tie (W/L/T) summary across 13 datasets (30 seeds, 95% CIs). Negative gap = RQ method has lower regret.

Figure 5: Data sufficiency diagnostic: regret gap of LinTS-RQ vs. LinTS as a function of

N/(bKd^{2})

. All real datasets with ratio

>1

(outside the shaded region) are wins or ties; year_prediction (ratio

=0.04

) is the sole loss. Synthetic Gaussian data (gray squares) confirm the threshold (Appendix K ).

	$\ell{=}1$	$\ell{=}2$	Adaptive
covertype	48.3%	42.2%	38.8% (depth 4)
higgs	45.4%	43.6%	43.7% (depth 3)

Table 4: TS-RQ regret rate (%) at fixed vs. adaptive depth.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Empirical behavior consistent with the conditional geometric depth-rate law, one column per method. Top: Median

R^{2}

of the exponential fit

a+b\rho^{\ell}

as a function of horizon

T

; fit quality improves with more data. Bottom: Normalized regret decay

(\Delta_{\ell}-\Delta_{5})/(\Delta_{1}-\Delta_{5})

at terminal

T

for each dataset.

Figure 7: Leave-one-out cross-validation of the exponential decay model

\Delta_{\ell}\approx a+b\rho^{\ell}

. Left: Predicted vs. actual regret rate

R(T)/T

for all (dataset, method, held-out depth) triples; the gray band shows

\pm 5\%

. Right: Prediction error by held-out depth; depths

2

–

5

are predicted within

\sim\!1\%

, while depth

1

(furthest from

\Delta_{\infty}

) shows larger extrapolation error.

Figure 8: Regret rate

R(T)/T

vs. depth

\ell

for each RQ method, with fitted exponential decay curves

a+b\rho^{\ell}

overlaid. Each line is a dataset; error bars show

\pm 1

std over 30 seeds. Median

R^{2}

is reported per method.

Figure 9: Normalized within-cell reward variance

\mathrm{WCV}(\ell)/\mathrm{WCV}(0)

vs. RQ depth

\ell

for all 13 datasets (log scale). Every dataset decays monotonically; the near-linear trajectories are consistent with geometric decay with median

\rho=0.83

$\mathcal{M}$	$p(\mathcal{M})$	Exploration	Joint memory	Actual memory
TS	$1$	$O(b\ell\sqrt{KT\log(T/\delta)})$	$O(b^{2}\ell^{2}K)$	$O(b\ell K)$
LinTS	$d$	$O(b\ell d\sqrt{KT\log(T/\delta)})$	$O(b^{2}\ell^{2}d^{2}K)$	$O(b\ell d^{2}K)$
SGD-LinTS	$d$	$O(b\ell d\sqrt{KT\log(T/\delta)})$	$O(b^{2}\ell^{2}d^{2}K)$	$O(b\ell dK)$

Appendix

Table 5: Exploration and memory scaling by base learner. “Joint” columns refer to the analytical variant; “Actual” is Algorithm 1 .

Dataset	$N$	$d$	$K$	$p_{\max}$	eff. rank
airlines_delay	500K	7	2	0.555	4.8
bng_elevators	1M	18	10	0.183	3.7
bng_letter	1M	16	26	0.041	4.6
covertype	581K	54	7	0.488	3.9
beer_reviews	500K	12	104	0.074	3.2
hepmass	5M	27	2	0.500	5.7

Appendix

Table 6: Datasets used in the evaluation.

N

= total samples,

d

= feature dimension,

K

= number of arms (classes),

p_{\max}

= majority class frequency, eff. rank = effective rank of the feature covariance. All datasets are publicly available from OpenML and UCI Dua and Graff (2019) ; codebook trained on

n_{\mathrm{rq}}{=}50\text{K}

held-out samples. Individual sources: HIGGS and SUSY Baldi et al. (2014) , HEPMASS Baldi et al. (2016) , covertype Blackard and Dean (1999) , MiniBooNE Roe et al. (2005) , year_prediction Bertin-Mahieux et al. (2011) .

Dataset	$d$	$K$	Base (%)	RQ (%)	Gap (%)	$T$
airlines	7	2	41.3	39.8	$-3.6\pm 0.1$	450K
beer_rev.	12	104	85.6	83.0	$-2.9\pm 0.5$	450K
bng_elev.	18	10	75.1	75.0	$-0.2\pm 0.0$	950K
bng_letter	16	26	52.4	46.6	$-11.0\pm 0.2$	950K
covertype	54	7	28.1	23.9	$-15.0\pm 0.2$	530K
hepmass	27	2	18.6	18.4	$-1.1\pm 0.1$	830K

Appendix

Table 7: LinTS-RQ vs LinTS: per-dataset regret rates (%) and gap with 95% CIs.

Dataset	$d$	$K$	Base (%)	RQ (%)	Gap (%)	$T$
airlines	7	2	45.6	42.8	$-6.2\pm 2.5$	450K
beer_rev.	12	104	92.7	89.8	$-3.2\pm 0.3$	450K
bng_elev.	18	10	81.7	77.9	$-4.6\pm 0.2$	950K
bng_letter	16	26	96.0	71.4	$-25.6\pm 0.5$	950K
covertype	54	7	51.3	38.8	$-24.4\pm 0.7$	530K
hepmass	27	2	50.0	21.7	$-56.6\pm 0.2$	830K

Appendix

Table 8: TS-RQ vs TS: per-dataset regret rates (%) and gap with 95% CIs.

Dataset	$d$	$K$	Base (%)	RQ (%)	Gap (%)	$T$
airlines	7	2	46.2	42.4	$-8.2\pm 1.1$	450K
beer_rev.	12	104	91.7	89.2	$-2.8\pm 0.2$	450K
bng_elev.	18	10	80.0	77.6	$-3.0\pm 0.6$	950K
bng_letter	16	26	81.6	64.3	$-21.2\pm 0.7$	950K
covertype	54	7	48.9	32.0	$-34.6\pm 1.4$	530K
hepmass	27	2	20.5	19.1	$-7.1\pm 0.2$	830K

Appendix

Table 9: SGD-RQ vs SGD-LinTS: per-dataset regret rates (%) and gap with 95% CIs.

Figure 10: Regret rate vs.

T

for all 9 methods on each dataset. Shaded bands show

\pm 1

std across seeds.

Figure 11: Regret scaling analysis for TS-RQ. Left: Power-law exponent

\alpha

per dataset (

R(T)\sim T^{\alpha}

); all values are near

1.0

, indicating bias-dominated regret. Right: Regret rate

R(T)/T

vs.

T

for selected datasets, converging to a constant consistent with the

O(\Delta_{\ell}T)

bias term in Theorem 3.1 .

Figure 12: Data sufficiency:

N/(bKd^{2})

vs. LinTS-RQ regret gap. Real datasets (filled circles) win at ratios as low as

1.1

. Synthetic Gaussian (open squares) requires ratio

>30

. year_prediction is the sole outlier at ratio

0.04

Dataset	NMSE (d=4)	LinTS-RQ wins?
skin_seg.	0.005	Yes ( $-96\%$ )
kddcup99	0.018	Yes ( $-16\%$ )
airlines	0.078	Yes ( $-3.6\%$ )
miniboone	0.111	Yes ( $-10\%$ )
covertype	0.185	Yes ( $-15\%$ )
beer_rev.	0.186	Yes ( $-2.9\%$ )

Appendix

Table 10: Normalized reconstruction error at depth 4.

Figure 13: Effective rank vs. regret rate. SGD-LinTS (squares) fails severely when eff. rank

<1.5

(red zone), performing worse than context-free TS (diamonds). TS-RQ (circles) is unaffected.

Dataset	LinTS-RQ	D-XGB	Gap	Call
kddcup99	0.2	9.0	$-8.8$	Win
bng_letter	46.6	56.1	$-9.5$	Win
covertype	23.9	28.9	$-5.0$	Win
beer_reviews	83.1	85.7	$-2.6$	Win
skin_seg	0.3	2.9	$-2.6$	Win
susy	20.3	21.7	$-1.4$	Win

Appendix

Table 11: Per-dataset regret rate (%): LinTS-RQ vs. Doubling-XGB. Negative gap = LinTS-RQ has lower regret.

Dataset	LinTS-RQ	D-NeuralCB	Gap	Call
beer_reviews	83.0	87.7	$-4.6$	Win
bng_letter	46.6	52.8	$-6.2$	Win
covertype	23.9	27.6	$-3.7$	Win
kddcup99	0.2	1.4	$-1.2$	Win
skin_seg	0.3	1.6	$-1.3$	Win
airlines_delay	39.8	40.2	$-0.4$	Tie

Appendix

Table 12: Per-dataset regret rate (%): LinTS-RQ vs. Doubling-NeuralCB (30 LinTS-RQ seeds, 10 NeuralCB seeds). Negative gap = LinTS-RQ has lower regret.

Figure 14: Per-dataset regret rate: LinTS-RQ vs. Doubling-XGB.

Figure 15: Per-dataset regret rate: LinTS-RQ vs. TabNet-LinTS.

Figure 16: Wall-clock

\mu

s/step per method at adapted DRQm depths. TS-RQ is consistently the cheapest RQ variant (

50{-}750\,\mu

s). LinTS-RQ on year_prediction (

d{=}90

K{=}84

) reaches

28

ms.

Contextual Bandits

Contextual bandits learn which action to take from features of each situation when only the chosen action’s reward is observed.

See all 151 papers

Provable Anytime Ensemble Sampling Algorithms in Nonlinear Contextual Bandits
Oct 12, 2025Jiazheng Sun, Weixin Wang, Pan XuEnsemble Learning Contextual Bandits
Duke University
A Practical Algorithm for Feature-Rich, Non-Stationary Bandit Problems
Mar 17, 2026Wei Min Loh, Sajib Kumer Sinha, Ankur Agarwal +1Bandit Learning Multi-Armed Bandits
University of Waterloo, Vector Institute · Manulife Financial
Batched Stochastic Linear Bandits with 1-Bit Communication Constraints
May 29, 2026Ivan Lau, Daniel McMorrow, Kevin Jamieson +1Multi-Armed Bandits Batched Bandits
National University of Singapore · University of Washington

Oct 12, 2025J/K move · Enter open · S save

Provable Anytime Ensemble Sampling Algorithms in Nonlinear Contextual Bandits

Jiazheng Sun, Weixin Wang, Pan Xu

Duke University

Ensemble Learning Contextual Bandits Generalized Linear Models Generalized Linear Bandits

Open paper AskPDF arXiv

We provide a unified algorithmic framework for ensemble sampling in nonlinear contextual bandits and develop corresponding regret bounds for two most common nonlinear contextual bandit settings: Generalized Linear Model Ensemble Sampling (GLM-ES) for generalized linear contextual bandits and Neural Ensemble Sampling (Neural-ES) for neural contextual bandits. Both methods maintain multiple estimators for the reward model parameters via maximum likelihood estimation on randomly perturbed data. We prove high-probability frequentist regret bounds of

\widetilde{\mathcal{O}}(d^{3/2} \sqrt{T} + d^{4})

for GLM-ES and

\widetilde{\mathcal{O}}(\widetilde{d}^{3/2} \sqrt{T})

for Neural-ES, where

d

is the dimension of feature vectors,

\widetilde{d}

is the effective dimension of a neural tangent kernel (NTK) matrix and

T

is the number of rounds. The regret bound of GLM-ES matches the state-of-the-art result of randomized exploration algorithms in generalized linear bandit setting. In the theoretical analysis, we introduce techniques that address challenges specific to nonlinear models. Practically, we remove fixed-time horizon assumption by developing anytime versions of our algorithms, suitable when

T

is unknown. Finally, we empirically evaluate GLM-ES, Neural-ES and their anytime variants, demonstrating strong performance. Overall, our results establish ensemble sampling as a provable and practical randomized exploration approach for nonlinear contextual bandits.

Abstract

Contextual Bandits

Explore similar work

Provable Anytime Ensemble Sampling Algorithms in Nonlinear Contextual Bandits

A Practical Algorithm for Feature-Rich, Non-Stationary Bandit Problems

Batched Stochastic Linear Bandits with 1-Bit Communication Constraints