Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Latency-Aware Contextual Bandit: Application to Cryo-EM Data Collection

Lai Wei, Ambuj Tewari, Michael A. Cianfrocco

TL;DR

Latency-aware contextual bandits extend the contextual bandit framework to account for action latency and multi-arm decisions by formulating the problem as a semi-Markov decision process and deriving a Bellman-optimality characterization. The COAF algorithm combines stochastic approximation with UCB-based confidence sets to estimate the optimal average reward $Γ^*$ and filter arms accordingly, achieving sublinear regret. Theoretical results give regret bounds for both oracle and learned reward settings across linear and general function classes, matching standard contextual bandit rates. Empirical evaluations on MovieLens and cryo-EM data collection demonstrate improved time-efficient reward accumulation and promise for automated, high-throughput experimental workflows.

Abstract

We introduce a latency-aware contextual bandit framework that generalizes the standard contextual bandit problem, where the learner adaptively selects arms and switches decision sets under action delays. In this setting, the learner observes the context and may select multiple arms from a decision set, with the total time determined by the selected subset. The problem can be framed as a special case of semi-Markov decision processes (SMDPs), where contexts and latencies are drawn from an unknown distribution. Leveraging the Bellman optimality equation, we design the contextual online arm filtering (COAF) algorithm, which balances exploration, exploitation, and action latency to minimize regret relative to the optimal average-reward policy. We analyze the algorithm and show that its regret upper bounds match established results in the contextual bandit literature. In numerical experiments on a movie recommendation dataset and cryogenic electron microscopy (cryo-EM) data, we demonstrate that our approach efficiently maximizes cumulative reward over time.

Latency-Aware Contextual Bandit: Application to Cryo-EM Data Collection

TL;DR

Latency-aware contextual bandits extend the contextual bandit framework to account for action latency and multi-arm decisions by formulating the problem as a semi-Markov decision process and deriving a Bellman-optimality characterization. The COAF algorithm combines stochastic approximation with UCB-based confidence sets to estimate the optimal average reward and filter arms accordingly, achieving sublinear regret. Theoretical results give regret bounds for both oracle and learned reward settings across linear and general function classes, matching standard contextual bandit rates. Empirical evaluations on MovieLens and cryo-EM data collection demonstrate improved time-efficient reward accumulation and promise for automated, high-throughput experimental workflows.

Abstract

We introduce a latency-aware contextual bandit framework that generalizes the standard contextual bandit problem, where the learner adaptively selects arms and switches decision sets under action delays. In this setting, the learner observes the context and may select multiple arms from a decision set, with the total time determined by the selected subset. The problem can be framed as a special case of semi-Markov decision processes (SMDPs), where contexts and latencies are drawn from an unknown distribution. Leveraging the Bellman optimality equation, we design the contextual online arm filtering (COAF) algorithm, which balances exploration, exploitation, and action latency to minimize regret relative to the optimal average-reward policy. We analyze the algorithm and show that its regret upper bounds match established results in the contextual bandit literature. In numerical experiments on a movie recommendation dataset and cryogenic electron microscopy (cryo-EM) data, we demonstrate that our approach efficiently maximizes cumulative reward over time.

Paper Structure

This paper contains 31 sections, 13 theorems, 99 equations, 3 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Problem Settings
  4. Latency-Aware Contextual Bandits
  5. Relationship with Existing Works
  6. Cryo-EM Data Collection
  7. Maximization of Average Reward
  8. Optimal Average Reward
  9. Step 1:
  10. Step 2:
  11. Algorithm Regret
  12. Contextual Online Arm Filtering Algorithm
  13. A Generic Algorithm Framework
  14. Worst-case Regret Analysis for COAF
  15. Regret Upper Bound with Linear $\mathcal{F}$
  16. ...and 16 more sections

Key Result

Theorem 1

For the latency-aware contextual bandit problem $\mathcal{M} = (P_{\mathrm{env}}, \psi_*)$, let $({\mathbb{X}}, {\mathbb{A}}, {\textnormal{l}}) \sim P_{\mathrm{env}}$ and let $\boldsymbol{\mu} = \{\mu_i\}_{i=1}^{{\textnormal{n}}}$, where $\mu_i = \psi_*({\mathbf{x}}_i)$ for each ${\mathbf{x}}_{i} \i

Figures (3)

  • Figure 1: Cryo-EM data collection at multiple magnifications: (i) grid-level shows the entire grid at low magnification, (ii) square-level captures individual squares at medium magnification to assess ice quality within holes, and (iii) full exposures are high-magnification images of selected holes.
  • Figure 2: Regrets of COAF and $3$ baselines: (i) Threshold selects movies with rating $\geq1.75$, (ii) COAF-ORC uses the true mean reward, and (iii) COAF-TS employs Thompson Sampling approach.
  • Figure 3: Setup and results of the cryo-EM data collection experiment.

Theorems & Definitions (27)

  • Remark 1
  • Theorem 1
  • proof
  • Remark 2
  • Lemma 2
  • Remark 3
  • Theorem 3
  • Definition 1
  • Definition 2
  • Theorem 4
  • ...and 17 more