Learning Personalized Memory Gating from Delayed One-Sided Feedback

Preprint from Salesforce AI Research

Learning Personalized Memory Gating from Delayed One-Sided Feedback Kaixuan Ji1,2 , Zixiang Chen1 , Wenting Zhao1 , Cheng Qian1,3 , Heng Wang1,3 , Weizhi Zhang1,4 , Jielin Qiu1 , Liangwei Yang1 , Quanquan Gu2 , Silvio Savarese1 , Shelby Heinecke1 , Huan Wang1 1 Salesforce AI Research 2 University of California, Los Angeles 3 University of Illinois at Urbana-Champaign 4 University of Illinois Chicago

arXiv:2610.23812v1 [cs.LG] 16 Oct 2026

Abstract Long-term memory enables language-model agents to personalize future interactions, but a finite memory bank cannot preserve every extracted fact. Existing memory gating policies are typically shared across users and therefore fail to account for user-specific patterns of future reuse. In this paper, we study memory gating at the fact level. We first propose a retrieval-based labeling strategy that labels each fact by whether it is later retrieved and incorporated into responses. We then formulate online memory gating as contextual apple tasting with delayed one-sided feedback, and propose Warm-Start UCB, which combines offline learning from existing data with online adaptation √ establishes a high-probability √ to the current user. Our theoretical analysis e(𝑇𝜖 ∧ 𝑑 𝑇) when the delay window is at most 𝑂 ( 𝑇), where 𝑇 is the total number regret upper bound of 𝑂 e(·) hides constant and of rounds, 𝜖 is the initialization error, 𝑑 is the dimension of the feature map, and 𝑂 poly-logarithmic factors. The results of our experiment demonstrate that our method outperforms all baselines under different memory budgets and history lengths, establishing it as a promising personalized memory gating policy.

1

Introduction

Memory systems, which extract and store important information during human-AI interaction, have been widely deployed in large language model (LLM)-based agents (Packer et al., 2023; Zhong et al., 2024; Chhikara et al., 2025). Typically, a memory system extracts candidate facts from each dialogue, stores a subset in a memory bank, and retrieves relevant entries for later queries. Because storage and context budgets are finite, the system must decide which facts are sufficiently valuable to preserve. Existing systems generally make this decision using criteria shared across users, such as heuristics like whether a fact expresses a stable preference, profile attribute, or ongoing goal (Park et al., 2023; Zhong et al., 2024), or using policies learned from long-context training data (Yan et al., 2026; Yu et al., 2026; Sun et al., 2026). Nevertheless, such universal criteria overlook the heterogeneity across different users. In et al. (2026) formalized this issue at the session level and introduced PerMem-Bench, together with several LLM-based storage-gating baselines. Their results show substantial headroom between perfect personalized gating and existing causal baselines, motivating methods that learn directly from user-specific interaction outcomes. To learn directly from user-specific interaction outcomes, we first need to identify a feedback signal grounded in the interaction that reflects the utility of each fact. While a deployed system does not observe whether a discarded fact would have been useful, for a stored fact, later retrieval provides a natural signal of future utility: a fact is utilized if it is hit by the retrieval and incorporated into the response in some later conversation. This implies that whether a fact gets retrieved and confirmed by the LLM agent can serve as a flag of utility, and therefore provides natural and reliable guidance for learning a memory gating policy. The retrieval-based feedback has two outstanding features: the feedback is only available on those added facts, and it becomes meaningful only after some unknown delay. To capture this nature, we formally frame the personalized memory gating problem as an apple tasting with one-sided delayed feedback problem (Helmbold et al., 1992; Vernade et al., 2020). We further propose a two-stage learning strategy, Warm-Start UCB. The first stage learns the utility 1

Preprint from Salesforce AI Research pattern from data of existing users. The second stage adapts to user-specific patterns during online interaction via an upper confidence bound (UCB) algorithm. Our theoretical analysis and empirical results verify the effectiveness of Warm-Start UCB as a gating policy. Our contributions are summarized as follows: • We propose a labeling pipeline for learning personalized memory gating, which derives utility labels from downstream retrieval outcomes, refined through LLM-judge verification. This labeling strategy provides a direct feedback for memory gating learning, without relying on heuristic or downstream reward signals. • We frame personalized memory gating as an apple tasting with one-sided delayed feedback problem. To solve this problem, we propose a two-stage learning strategy, Warm-Start UCB. Our theoretical analysis implies that Warm-Start UCB effectively leverages the offline data while evolving towards user-specific patterns through online interaction. • We empirically evaluate Warm-Start UCB on our benchmark across a range of memory budgets and interaction lengths. Experiment results show that Warm-Start UCB consistently outperforms previous LLM-based methods and the naive admit-all baseline, suggesting that Warm-Start UCB serves as an effective gating policy for personalized memory. Notation. We use [𝑛] to denote the set {1, . . . , 𝑛}. In this work, we use lowercase letters to represent scalars, and 𝑑 denote vectors and matrices by lower and uppercase boldface √ letters respectively. Given a vector x ∈ R , we denote the vector’s 𝐿 2 -norm by ∥x∥ 2 . We further define ∥x∥ 𝚺 = x⊤ 𝚺x given a positive semidefinite matrix 𝚺 ∈ R𝑑×𝑑 . For two semidefinite matrices 𝚺, 𝚲 ∈ R𝑑×𝑑 , we use 𝚺 ⪰ 𝚲 if x⊤ (𝚺 − 𝚲)x ≥ 0 for any x ∈ R𝑑 . We reserve I for the identity matrix. For nonnegative sequences {𝑥 𝑛 } and {𝑦 𝑛 }, we write 𝑥 𝑛 = 𝑂 (𝑦 𝑛 ) interchangeably with 𝑥 𝑛 ≲ 𝑦 𝑛 if lim sup𝑛→∞ 𝑥 𝑛 /𝑦 𝑛 < ∞, 𝑦 𝑛 = Ω(𝑥 𝑛 ) or 𝑥 𝑛 ≳ 𝑦 𝑛 if 𝑥 𝑛 = 𝑂 (𝑦 𝑛 ), and 𝑦 𝑛 = Θ(𝑥 𝑛 ), or alternatively, 𝑦 𝑛 ∼ 𝑥 𝑛 , if e e to hide polylog factors. We use 1{·} to denote e(·), Ω(·), 𝑥 𝑛 = 𝑂 (𝑦 𝑛 ) and 𝑥 𝑛 = Ω(𝑦 𝑛 ). We further employ 𝑂 and Θ(·) the indicator function. For any 𝑥, 𝑦 ∈ R, we denote 𝑥 ∧ 𝑦 = min{𝑥, 𝑦} and 𝑥 ∨ 𝑦 = max{𝑥, 𝑦}.

2

Related Work

Personalized Memory System. Memory systems are widely used in various applications, including embodied systems (Zhu et al., 2023; Wang et al., 2024b), personal assistants (Wang et al., 2025; Zheng et al., 2024) and recommender systems (Wang et al., 2024a), incentivizing a long line of work in improving the effectiveness and efficiency. One direction mainly focuses on learning a policy that selectively extracts and persists crucial information from the interaction history (Chhikara et al., 2025; Zhou et al., 2026; Yan et al., 2026; Tan et al., 2025). Another direction focuses on leveraging various data structures to capture relation between memory units to improve subsequent retrieval (Xu et al., 2026; Yang et al., 2026; Hu et al., 2025; Rezazadeh et al., 2025; Lu et al., 2026). Despite these advances, most existing systems employ a universal storage policy and therefore fail to fully account for the heterogeneity across different users. In et al. (2026) first formalize this limitation and introduce PerMem-Bench for evaluating personalized memory systems. They further propose several baseline approaches that operate the gating at session level, with LLM-based learning-free methods as the backbone. Compared to In et al. (2026), we consider fact-level online adaptation from retrieval-derived feedback, which eliminates the reliance on human or LLM-based heuristic. This results in an apple tasting with one-sided delayed feedback formulation. Learning with Apple Tasting Feedback. Apple tasting, introduced by Helmbold et al. (1992; 2000) is one of the canonical models for one-sided feedback, where the binary outcome is only revealed under the positive action. Apple tasting has been studied under several distinct comparators and modeling assumptions. In online binary classification against a hypothesis class, Raman et al. (2024) characterize agnostic learnability through Littlestone dimension and establish a realizable-case trichotomy; Chase & Mehalel (2025) study the additional restriction to deterministic learners. Apple tasting also appears as a canonical partial-monitoring or feedback-graph problem when regret is measured against a fixed action (Bartók et al., 2014; Alon et al., 2015; Cesari & Colomboni, 2026), with rates depending on the precise loss, feedback, and adversary model. Contextual variants have been studied through partial monitoring with side information, Bayesian logistic models, feedback-graph reductions, and log-loss (Bartók & Szepesvári, 2012; Grant & Leslie, 2021; Zhang et al., 2023; An et al., 2026). Extension has also been made to other settings like fairnessaware (Bechavod et al., 2019; Bechavod & Roth, 2023) and strategically manipulated context (Harris et al., 2023). 2

Preprint from Salesforce AI Research Learning with Delayed Feedback. There are long lines of work studying decision making problem with nonimmediate feedback, the most canonical among which is that the delay is generated stochastically or adversarially independently of the feedback. In this setting, a regret comparable to the non-delay counterpart up to only additive terms of the delay is provably achievable under both multi-armed setting (Joulani et al., 2013; Cesa-Bianchi et al., 2019; Thune et al., 2019; Bistritz et al., 2022; Zimmert & Seldin, 2020; Gyorgy & Joulani, 2021; Masoudian et al., 2022; Schlisselberg et al., 2026a) and linear setting (Zhou et al., 2019; Howson et al., 2023; Van Der Hoeven et al., 2023; Blanchet et al., 2024; Schlisselberg et al., 2026b). Another line of work considers delayed feedback that is entangled with the feedback structure (Lancewicki et al., 2021; Tang et al., 2024; Schlisselberg et al., 2025), among which the most relevant to us are Vernade et al. (2017; 2020) that studied the running-zero-flipping-one delay model √ e(𝑑 𝑇 + 𝑚𝑑) regret upper bound (Vernade et al., 2020), where 𝑑 is the dimension of feature map and 𝑚 and achieved 𝑂 is a high-probability upper bound of the stochastic delay. In this paper, we adopt the running-zero-flipping-one delay model with apple-tasting feedback to account for the specific structure in agentic memory system, where feedback is only available at added memory entries.

3

Dataset Construction

In this section, we introduce a retrieval-based labeling pipeline for learning memory gating policy, and provide a quick analysis for how the utility patterns are shared across users and how they differ. We defer the meta-evaluation of label validity to Appendix A.2 and statistics of the constructed benchmark to Appendix A.3. 3.1

Benchmark Construction

Our construction begins with any corpus of user-LLM interaction history between user 𝑢 and the LLM, denoted by {S1(𝑢) , S2(𝑢) , . . .}, such that each session is composed of S𝑡(𝑢) = {𝑥1𝑡 , 𝑎 1𝑡 , · · · , 𝑥 𝑡𝑘 , 𝑎 𝑡𝑘 }, where 𝑥𝑖𝑡 is the user’s 𝑖-th prompt and 𝑎 𝑖𝑡 the agent’s corresponding response in the 𝑡-th session of the entire interaction history between user 𝑢 and the agent. In this paper, we consider the corpus generated via the strategy in In et al. (2026). Specifically, we adopt the pre-generated PerMem-Bench𝑠 (PerMemBench-S), and additionally synthesize a longer dataset (PerMemBench-L), to better capture the long-horizon nature of human-AI interaction. Fact Extraction. For each user 𝑢, we process the sequence of session, {S𝑡 }𝑡 in chronological order by feeding them into an off-the-shelf memory system1 . Starting from an empty memory bank M0 = ∅, for each session S𝑡 , the memory system extracts a collection of candidate atomic facts F𝑡 = { 𝑓𝑡 ,1 , 𝑓𝑡 ,2 , . . .} and then merges these candidate facts into the memory bank M𝑡 = M𝑡 −1 ∪ F𝑡 . In plain words, M𝑡 is the memory bank before the 𝑡 + 1-th session, if we add all extracted facts and the memory bank has infinite budget. Retrieval-Based Labeling. After constructing the pool of candidate facts {M𝑡 }𝑡 , we now assign each candidate fact 𝑓𝑡 ,𝑖 a binary utility label indicating its usefulness. Recall that a fact can only be useful if it gets retrieved at some later point. Formally, fix a fact 𝑓𝑡 ,𝑖 , for every user turn 𝑥 𝑡 ′ ,𝑖 ′ with 𝑡 ′ > 𝑡, an LLM first decides whether answering that turn requires querying memory at all. For turns judged to require memory, the same LLM formulates a natural-language query 𝑞 𝑡 ′ and invokes the retriever of the memory system with query 𝑞 𝑡 ′ against the memory bank state M𝑡 ′ −1 . 𝑓𝑡 ,𝑖 is labeled with “required for 𝑥 𝑡 ′ ,𝑖 ′ ” only if 𝑓𝑡 ,𝑖 appears among the top-𝑘 retrieval hits. LLM-judge Verification. A retrieval hit alone is a noisy relevance signal: a fact may be topically related to 𝑞 𝑡 ′ without actually being needed to answer it (e.g., a temporal mismatch). We therefore filter the obtained retrieval hits with an LLM judge. For each pair of (𝑥 𝑡 ′ ,𝑖 ′ , 𝑓𝑡 ,𝑖 ) such that 𝑓𝑡 ,𝑖 hit by the retriever for 𝑥 𝑡 ′ ,𝑖 ′ , we ask an LLM to judge whether the presence of 𝑓𝑡 ,𝑖 in the context is necessary for the answer to 𝑥 𝑡 ′ ,𝑖 ′ , or substantially improves the answer quality. We denote all hits verified by this procedure with H0 = {( 𝑓𝑡 ,𝑖 , 𝑥 𝑡 ′ ,𝑖 ′ ) : 𝑓𝑡 ,𝑖 hit by the retriever for 𝑥 𝑡 ′ ,𝑖 ′ and verified by the LLM judge}, as all valid retrieval hits. Finally, to mitigate potential neglect of the retrieval system, for all 𝑥 𝑡 ′ ,𝑖 ′ where the retrieval results in no strong match (i.e., the scores of all retrieved facts lie below a threshold), we employ an LLM to inspect whether there is a subset of facts from the top retrievals that collectively help with answering the user’s prompt. We 1 Here we use mem0 (Chhikara et al., 2025), though the pipeline is agnostic to the underlying extractor.

3

Preprint from Salesforce AI Research denote these prompt-fact pairs as H1 = {( 𝑓𝑡 ,𝑖 , 𝑥 𝑡 ′ ,𝑖 ′ )} and merge  them with H0 to obtain the final set of valid retrieval hits H = H0 ∪ H1 . The final utility label is given by 𝑦 𝑡 ,𝑖 = 1 ∃ (𝑡 ′ , 𝑖 ′ ), ( 𝑓𝑡 ,𝑖 , 𝑥 𝑡 ′ ,𝑖 ′ ) ∈ H . Evaluation Protocol. During evaluation, we consider a memory budget with a fixed budget size 𝐵. For each distinct fact 𝑓 , the gating decides whether it should be added to the memory bank. If the decision is “add” and the bank is full, then the earliest entry in the bank will be evicted. Our main metric is retrieval-hit retention rate (RHR), which is the proportion of valid hits in H that are realized under the fixed budget. In particular, H ′ := {( 𝑓𝑡 ,𝑖 , 𝑥 𝑡 ′ ,𝑖 ′ ) ∈ H : 𝑓𝑡 ,𝑖 is still in the memory bank at 𝑡 ′ }, RHR := |H ′ |/|H |. At a high level, the RHR captures how well a policy preserves facts until the turn where they actually make a difference, given that the memory budget is limited. A miss of fact 𝑓𝑡 ,𝑖 at 𝑡 ′ stems from one of two reasons: the gating policy drops 𝑓𝑡 ,𝑖 , or it adds too many non-required facts so that 𝑓𝑡 ,𝑖 gets evicted too early. Hence, a policy that attains a high RHR usually adds the required facts to the memory bank and drops unnecessary ones. Besides this, we also consider the F1 score of a gating policy, by treating the storage decision itself as a binary classification problem against the ground truth labels. 3.2

Meta Analysis

In this section, we provide a preliminary study of the heterogeneity and homogeneity among different users. To do this, for each fact, we predict the binary utility label using the TF-IDF representation of the fact content and a logistic regression classifier. We compare two evaluation regimes per user: (i) Self, a stratified 𝑘-fold cross-validation using only 𝑢’s own facts, and (ii) leave-one-user-out (LOUO), a single shared-vocabulary model trained on the pooled facts of the other 19 users and evaluated on 𝑢.

Table 1: Self (within-user CV) vs. LOUO (fixed global model) TF-IDF + logistic regression AUC for the fact-gating label, averaged over 20 users. Dataset PerMemBench-S PerMemBench-L

Self AUC

LOUO AUC

Gap

0.756 ± 0.047 0.724 ± 0.029

0.622 ± 0.049 0.665 ± 0.024

0.134 0.059

The results are summarized in Table 1. First, we observe that the LOUO model itself is far above random (AUC 0.622 and 0.665, respectively, vs. 0.5), meaning that a classifier that has never seen the held-out user’s facts still recovers a considerable share of the signal: part of what makes a fact useful for future interactions is shared among users. Despite this, the self AUC exceeds the LOUO AUC prominently, indicating that a substantial, user-specific component of the gating rule does not transfer under a single fixed global classifier. This coexistence of homogeneity and heterogeneity indicates that, to make the best use of the data, a policy should both learn from other users and adapt to the current one. This motivates the proposed methods, introduced in the next section.

4

Algorithm and Analysis

In this section, we introduce an online learning model that captures the structure of retrieval-based feedback, along with an algorithm for learning a memory-gating policy from this feedback. 4.1

Problem Setup

We model the personalized memory gating problem as a contextual bandit with apple-tasting, delayed, one-sided feedback problem and describe the formal setup in this section. For simplicity, we fix a single user and omit the user subscript.

4

Preprint from Salesforce AI Research 4.1.1

Interaction Protocol

At each round 𝑡 = 1, . . . , 𝑇, the learner observes a context x𝑡 ∈ R𝑑 and chooses an action 𝑎 𝑡 ∈ {add, skip}. Here, the context x𝑡 is a feature representation2 of the candidate memory fact extracted at timestamp 𝑡, and the action 𝑎 𝑡 is the memory system’s gating decision: add commits the fact to the memory bank, while skip discards it. The learner receives different feedback depending on the action taken. If 𝑎 𝑡 = skip, the learner discards the fact and observes nothing at any future time. If 𝑎 𝑡 = add, the environment simultaneously samples a delay 𝐷 𝑡 > 0 from a fixed distribution D, independently, and at each subsequent time 𝑡 ′ > 𝑡, the learner observes the censored conversion indicator  𝑧 𝑡 ,𝑡 ′ = 𝑦 𝑡 · 1 𝐷 𝑡 ≤ 𝑡 ′ − 𝑡 , where 𝑦 𝑡 ∈ {0, 1} is the actual label such that 𝑦 𝑡 = 1 indicates that the fact is utilized by some future user turn, and 𝑦 𝑡 = 0 indicates that the fact is never used. Moreover, 𝐷 𝑡 is the number of rounds until the utilization is confirmed (e.g., the fact is retrieved for a later turn). The observation begins at 𝑧 𝑡 ,𝑡 = 0 and, if and only if 𝑦 𝑡 = 1, flips to 1 once the delay elapses. A persistent 0 is ambiguous between a transient fact and a required fact whose confirming retrieval has not yet occurred. Remark 4.1. This one-sided feedback structure, where the outcome of a decision is only ever revealed when the learner chooses the positive action, is classically termed apple tasting3 . Decision-making problems with such a feedback structure have been widely studied in the previous literature (Helmbold et al., 1992; 2000). Our problem can be seen as a combination of the apple-tasting feedback and the delayed running-zero-flipping-one feedback considered in Vernade et al. (2020). We assume the learner has some knowledge of the delay distribution. Assumption 4.2. There exist some known 𝑚 and 𝜅 𝑚 = Ω(1) such that P[𝐷 𝑡 ≤ 𝑚] = 𝜅 𝑚 , for any 𝑡 ∈ [𝑇]. Furthermore, we assume 𝑦 𝑡 ∼ Bernoulli(𝜂(x𝑡 )), where 𝜂 : X → [0, 1] is a user-specific function and 𝜂(x𝑡 ) is the utility score. A larger 𝜂(𝑥) means a higher probability that the information in 𝑥 is used in the future. We assume that there exists some unknown 𝜽 ∗ ∈ R𝑑 , such that 𝜂(x𝑡 ) = ⟨𝜽 ∗ , x𝑡 ⟩ ∈ [0, 1], where the features admit the following regularity conditions. Assumption 4.3. There exists some known constant 𝐿 > 0 such that ∥x𝑡 ∥ 2 ≤ 𝐿 for all 𝑡. 4.1.2

Learning Objective

With store-gain 𝑏 > 0 and asymmetric error costs 𝑐 miss , 𝑐 store > 0, the reward is

𝑟 (x𝑡 , 𝑎 𝑡 , 𝑦 𝑡 ) =

 +𝑏     −𝑐  

store

 −𝑐 miss    0 

𝑎 𝑡 = add, 𝑦 𝑡 = 1, 𝑎 𝑡 = add, 𝑦 𝑡 = 0, 𝑎 𝑡 = skip, 𝑦 𝑡 = 1, 𝑎 𝑡 = skip, 𝑦 𝑡 = 0.

Without loss of generality, we assume 𝑏 + 𝑐 store + 𝑐 miss = 1. Taking expectations over 𝑦 𝑡 ∼ Bernoulli(𝜂(x𝑡 )), the action add is preferred if and only if 𝜂(x𝑡 ) ≥ 𝜏, with 𝜏 = 𝑐 store , so the oracle is the threshold policy 𝜋★ (x) = 1{⟨𝜽 ★, x⟩ ≥ 𝜏}. The learner’s goal is to minimize the pseudo-regret against 𝜋★, Regret(𝑇) =

𝑇  ∑︁

  E 𝑦 𝑟 (x𝑡 , 𝜋★ (x𝑡 ), 𝑦) − 𝑟 (x𝑡 , 𝑎 𝑡 , 𝑦) .

𝑡=1

Remark 4.4. While the interaction protocol captures two key features of feedback from retrieval-based labeling, apple-tasting and one-sided delay, we notice that our online learning model still differs from practice in several ways. For example, within all retrieved and verified facts, their utility still varies: some are rarely retrieved and incorporated, whereas others are used more frequently. Also, the cost 𝑐 store for adding a non-required fact may vary over time: it will be higher when the memory bank is nearly full and lower when it has more capacity. Extending the model to account for these variations, and developing algorithms for the resulting settings, are promising directions for future work. 2 For example, an embedding of the fact together with its surrounding session context. 3 Consider a shopper deciding whether to buy an apple from a bin: they only learn whether it was good or bad by tasting (buying) it, never by leaving it on the shelf.

5

Preprint from Salesforce AI Research Algorithm 1 Warm-Started Ridge UCB with One-Sided Delayed Feedback (Warm-Start UCB) Require: Threshold 𝜏, warm-start center 𝝁0 , prior shape 𝚺0 ⪰ I, regularization parameter 𝜆 ≥ 𝐿 2 , confidence level 𝛿, delay parameters (𝑚, 𝜅 𝑚 ), horizon 𝑇 b0 ← 𝜅 𝑚 𝝁0 , 𝛼0 ← 𝛽𝜆, 𝛿 . 1: Initialize 𝑽0 ← 𝜆𝚺 0 , 𝒃 0 ← 𝜆𝚺 0 (𝜅 𝑚 𝝁0 ), 𝜽 2: for 𝑡 = 1, 2, . . . , 𝑇 do 3: Observe context x𝑡 ∈ R𝑑 4: Compute 𝜂e𝑡 ← ⟨𝜽b𝑡 −1 , x𝑡 ⟩ + 𝛼𝑡 −1 ∥x𝑡 ∥𝑽 −1 𝑡 −1 5: Play 𝑎 𝑡 ← add if 𝜂e𝑡 ≥ 𝜅 𝑚 𝜏. Otherwise play 𝑎 𝑡 ← skip 6: Observe all delayed feedback available at the end of round 𝑡 Í Í −1 b 7: Set V𝑡 ← 𝜆𝚺0 + Í𝑠≤𝑡 1{𝑎 𝑠 = add}x𝑠 x⊤ 𝑠 , 𝒃 𝑡 ← 𝜆𝚺 0 (𝜅 𝑚 𝝁0 ) + 𝑠≤𝑡 1{𝑎 𝑠 = add }x𝑠 𝑧¯ 𝑠,𝑡 , and 𝜽 𝑡 ← V𝑡 𝒃 𝑡 8: Set 𝛼𝑡 ← 𝛽𝜆, 𝛿 + 𝑠∈ P𝑡 ∥x𝑠 ∥ V𝑡−1 9: end for

4.2

Training from Other Profiles

The meta-analysis in Section 3.2 suggests that users’ storage-preference parameters contain both a shared component and a user-specific component. This motivates using data from other users to construct an informative warm start for the current user. We formalize the quality of this warm start through the following high-probability assumption. Assumption 4.5. The learner is given a prior center 𝝁0 ∈ R𝑑 , a positive-definite matrix 𝚺0 ⪰ I, and a known radius 𝜖 > 0 such that
P ∥𝜽 ★ − 𝝁0 ∥ 𝚺0 ≤ 𝜖 ≥ 1 − 𝛿0 . Here, the probability is taken over both the current user’s parameter 𝜽 ★ and the randomness in the pretraining data used to construct 𝝁0 and 𝚺0 . Overall, Assumption 4.5 states that the learner, through techniques like linear regression on other existing users, has access to a sufficiently accurate estimate of the center of the prior distribution of the per-user ground truth 𝜽 ∗ , which is close to any realized 𝜽 ∗ with high probability. In fact, Assumption 4.5 is plausible given the meta-analysis in Section 3.2 and standard analyses of random-design linear regression (Hsu et al., 2014). In particular, let 𝜽 denote a population-level parameter shared across users. The transferability observed in Section 3.2 suggests that the current user’s parameter is concentrated around this common component. Therefore, it is reasonable to make the assumption that, with probability at least 1 − 𝛿pop , ∥𝜽 ★ − 𝜽 ∥ 𝚺0 ≤ 𝜖pop .

(4.1)

At the same time, the cross-user predictability demonstrated in Section 3.2 suggests that data from other users can be used to estimate the shared component. Standard random-design linear-regression results, such as Hsu et al. (2014, Theorem 2), can be used to establish that, with probability at least 1 − 𝛿est , ∥ 𝝁0 − 𝜽 ∥ 𝚺0 ≤ 𝜖est .

(4.2)

Combining (4.1) and (4.2) with the triangle inequality gives ∥𝜽 ★ − 𝝁0 ∥ 𝚺0 ≤ 𝜖 pop + 𝜖est with probability at least 1 − 𝛿pop − 𝛿est . Thus, Assumption 4.5 holds with 𝜖 = 𝜖pop + 𝜖 est and 𝛿0 = 𝛿pop + 𝛿est . We refer to Appendix D for a detailed discussion. 4.3

Online Adaptation

While a parameter pre-trained on other profiles might already serve as a good proxy for most individual profiles, it does not account for the heterogeneity revealed in Section 3.2, and therefore leads to linear regret as the time horizon 𝑇 → ∞. To address this issue, we introduce an online adaptation strategy that starts from the pre-trained parameter 𝝁0 but still adjusts towards the user’s specific parameter during the online interaction. The strategy is summarized in Algorithm 1. Similar warm-start strategies have also been applied in previous bandit papers such as Cella et al. (2020); Peleg et al. (2022). 6

Preprint from Salesforce AI Research Algorithm 1 extends previous algorithms for one-sided delayed feedback (Vernade et al., 2020) to the apple-tasting feedback setting. At each round, Algorithm 1 observes the context x𝑡 and computes an index 𝜂e𝑡 that is optimistic for 𝜅 𝑚 𝜂(x𝑡 ) with high probability, inspired by previous work on linear contextual bandits (Abbasi-Yadkori et al., 2011). Then the algorithm plays add if and only if the index exceeds the threshold 𝜅 𝑚 𝜏. Finally, the algorithm updates its parameter estimate based on the most recently available feedback. The key challenge arises from the estimation of the parameter. For every added sample 𝑠 and current time step 𝑡 ≥ 𝑠, a running-zero label is ambiguous: it might be either that the fact is not useful (i.e., 𝑦 𝑠 = 0) or that it is still under the delay (i.e., 𝐷 𝑠 > 𝑡 − 𝑠). To address this issue, we follow the approach in Vernade et al. (2020) and truncate the observation by a window of length 𝑚: for each added sample 𝑠, we wait for 𝑚 rounds and treat its label as 1 if we observe a flip, and 0 otherwise, even if it flips to 1 later. Formally, define the truncated observation  𝑧¯𝑠,𝑡 := 𝑦 𝑠 1 𝐷 𝑠 ≤ min{𝑡 − 𝑠, 𝑚} . Our next step is to treat the truncated observation 𝑧¯𝑠,𝑡 as the regression label. Notice that E[ 𝑧¯𝑠,𝑠+𝑚 | x𝑠 ] = P[𝐷 𝑠 ≤ 𝑚]𝜂(x𝑠 ) = ⟨𝜅 𝑚 𝜽 ★, x𝑠 ⟩. −𝑚 Consequently, the regression on the feature-label pairs {(x𝑠 , 𝑧¯𝑠,𝑡 )}𝑡𝑠=1 targets 𝜅 𝑚 𝜽 ★. Since our regression target is −1 when comparing it to the 𝜅 𝑚 𝜽 ★ rather than 𝜽 ★, we compensate for the scaling by multiplying our estimate 𝜂e𝑡 by 𝜅 𝑚 threshold 𝜏. Based on the regression, we compute the confidence radius of the estimate √︁ √ 𝛽𝜆, 𝛿 = 𝜎 𝑑ℓ𝜆 + 2 log(1/𝛿) + 𝜅 𝑚 𝜆𝜖,

where ℓ𝜆 := log(1 + 𝑇 𝐿 2 /(𝜆𝑑)), and 𝜎 is a sub-Gaussian parameter of the noise, which can be taken to be 1/2 in our setting. Finally, because observations from the most recent 𝑚 added samples may not yet have reached their final truncated values, the confidence radius 𝛼𝑡 contains an additional term corresponding to the pending-observation rounds P𝑡 = {𝑠 ≤ 𝑡 : 𝑎 𝑠 = add, 𝑡 − 𝑠 < 𝑚} ∑︁ ∥x𝑠 ∥ V𝑡−1 . 𝛼𝑡 = 𝛽𝜆, 𝛿 + 𝑠∈ P𝑡

Remark 4.6. Unlike the original algorithm of Vernade et al. (2020), our decision rule compares the estimated utility 𝜂(x) against a fixed, exogenous threshold 𝜏. In contrast, comparisons among contemporaneous arms are invariant to a common multiplicative factor. Consequently, our algorithm must know 𝜅 𝑚 = P(𝐷 ≤ 𝑚), or have a sufficiently accurate estimate of it, in order to rescale the decision threshold.

Table 2: Results on the PerMemBench-S benchmark. We use bold for the highest score and underline for the second highest. For LLM-based methods, “(Ours)” indicates the corresponding method with our fine-grained prompt. The results show that Warm-Start UCB outperforms the Universal baseline and LLM-based methods uniformly across all budget sizes. Metric Budget

60

Micro Average 120 250

500

60

Macro Average 120 250

500

Oracle

69.34

92.99

100.00

100.00

77.28

95.22

100.00

100.00

Universal Greedy Greedy (Ours) Context Context (Ours)

17.44 23.01 23.94 23.39 23.34

35.94 32.95 35.99 32.28 36.37

56.93 44.49 53.70 42.96 52.11

80.60 47.34 68.13 45.28 66.32

19.34 24.86 26.11 25.33 24.72

39.29 34.16 38.29 33.20 37.75

60.76 45.08 55.00 42.95 53.09

84.91 47.56 67.87 45.21 66.29

Cold-Start UCB Pre-trained Scorer Warm-Start UCB

20.55 24.35 24.10

37.63 38.62 38.64

56.62 56.90 58.73

82.13 79.64 84.48

23.29 26.52 26.40

41.49 41.72 41.98

60.95 59.73 61.84

83.93 81.06 86.94

7

Preprint from Salesforce AI Research Theorem 4.7. Suppose Assumptions 4.2, 4.3 and 4.5 hold. Fix any 𝛿 ∈ (0, 1). Then, with a proper choice of 𝜆 that depends on 𝑇, 𝑑, 𝑚, 𝜅 𝑚 , 𝜖, for sufficiently small 𝜖, with probability at least 1 − 𝛿0 − 𝛿, Algorithm 1 satisfies √    𝑑 𝑇 + 𝑚𝑑 e min 𝑇𝜖, Regret(𝑇) ≤ 𝑂 , 𝜅𝑚 e(·) hides all logarithmic factors and the dependency on 𝐿. where 𝑂 Theorem 4.7 shows that, when the pre-trained estimate 𝝁0 is accurate enough, the per-step regret is roughly bounded by 𝜖, so the algorithm only incurs 𝑇𝜖 regret when 𝑇 is small. On the other hand, when 𝑇 is large, the algorithm enjoys √ e(𝑑 𝑇) regret, indicating that the proposed strategy effectively learns the user-specific parameter. a sublinear 𝑂

5

Experiments

Table 3: Results on the PerMemBench-L benchmark. We use bold for the highest score and underline for the second highest. For LLM-based methods, “(Ours)” indicates the corresponding method with our fine-grained prompt. The results show that Warm-Start UCB significantly outperforms the Universal baseline and LLM-based methods uniformly across all budget sizes.

5.1

Metric Budget Size

120

250

Micro Average 500 800

Oracle

54.30

78.61

97.51

Universal Greedy Greedy (Ours) Context Context (Ours)

20.17 19.65 18.64 17.01 18.76

35.28 29.19 29.87 25.43 28.87

Cold-Start UCB Pre-trained Scorer Warm-Start UCB

20.75 24.72 25.22

35.46 38.08 37.77

Macro Average 500 800

1000

120

250

1000

100.00

100.00

58.68

82.62

98.33

100.00

100.00

50.77 40.59 45.55 36.30 45.13

65.73 47.42 58.12 41.56 57.33

73.77 49.19 62.87 43.49 63.49

22.03 20.38 19.54 17.70 19.79

37.13 30.02 31.32 25.93 30.13

52.55 40.53 46.59 36.50 46.76

67.50 46.50 58.56 40.61 58.54

75.71 47.68 63.10 42.52 64.40

51.22 55.21 56.14

66.82 66.13 69.21

74.42 75.02 76.46

22.12 26.44 26.73

37.33 39.98 39.87

53.09 56.88 57.92

68.67 67.02 69.93

76.43 76.05 77.52

Baselines and Methods

Model-Free Baselines. specific model

We follow In et al. (2026) and consider the following two baselines that do not operate on a

• Oracle: The memory system is provided with the ground truth labels and only adds the required facts, serving as the skyline of the performances. • Universal: The memory system adds every extracted fact into the memory bank, representing current memory systems. LLM-Based Methods. We also consider the causal LLM-based baselines introduced in In et al. (2026). For both LLM-based baselines, we use gpt-5-mini as the backbone and evaluate two prompting strategies. One follows the prompt of In et al. (2026), while the other, denoted by “Ours”, uses a more fine-grained rubric that additionally accounts for the remaining capacity of the memory bank. The complete prompts are provided in Appendix B.1. • LLM-Greedy. For each extracted fact, an LLM predicts whether the fact captures crucial information that is likely to be useful in a future session. The prediction is based only on the current fact, without access to the user’s previous interaction history. • LLM-Context-Aware. After each session, the dialogue is summarized in one or two sentences. For each subsequent fact, the LLM is provided with the current fact, the current session, and a sliding window containing the most recent 𝐾 session summaries, and is asked to determine whether the fact, given the context, might be required in the future. 8

UCB (Cold) Pre-trained UCB (Warm)

0

1

2

3 4 5 6 7 session-position decile

8

9

UCB (Cold) Pre-trained UCB (Warm)

0.10

F1 uplift over admit_all

0.14 0.12 0.10 0.08 0.06 0.04 0.02 0.00 −0.02

F1 uplift over admit_all

F1 uplift over admit_all

Preprint from Salesforce AI Research

0.08 0.06 0.04 0.02 0.00

0

1

2

3 4 5 6 session-position decile

(a) 120

7

8

0.07 0.06 0.05 0.04 0.03 0.02 0.01 0.00

9

UCB (Cold) Pre-trained UCB (Warm)

0

1

(b) 500

2

3 4 5 6 session-position decile

7

8

9

(c) 1000

Figure 1: Learning dynamics on PerMemBench-L: F1 gain over the Universal baseline in each session-position decile under memory budgets 𝐵 ∈ {120, 500, 1000}. Shaded regions denote the 95% CI over all users. The gating quality improves in later deciles for online methods Cold-Start UCB and Warm-Start UCB.

0.25 0.24 0.23 1 2

5

10

Pretraining Pool Size (a) 120

micro macro 15

0.58 0.57 0.56 0.55 0.54 0.53 0.52 0.51

0.79 0.78

Best Hit Rate

0.26

Best Hit Rate

Best Hit Rate

0.27

1 2

5

10

Pretraining Pool Size (b) 500

micro macro 15

0.77 0.76 0.75 0.74

1 2

5

10

Pretraining Pool Size

micro macro 15

(c) 1000

Figure 2: Effect of the pre-training pool size on PerMemBench-L: best micro- and macro-averaged RHRs under 𝐵 ∈ {120, 500, 1000}. The results show that the performance continues to improve as the pretraining pool size grows. Linear Gating Methods. We consider three linear gating methods, all of which use the same linear feature representation of the facts. Besides the Warm-Start UCB introduced in Algorithm 1, we consider its two variants as its ablations. We refer to Appendix B.2 for the implementation details. • Cold-Start UCB. Algorithm 1 run with 𝝁0 = 0 and no access to other users’ data. The gating rule is learned using only the current user’s feedback, which isolates the contribution of online adaptation. • Pre-trained Scorer. A linear scorer 𝝁0 fitted offline on the labeled facts of other users and kept frozen throughout the interaction, without adaptation to the current user’s feedback, which isolates the contribution of cross-user pre-training. 5.2

Experiment Results

The results are summarized in Table 2 and Table 3. We report the RHR as a percentage under both micro and macro averaging. Micro averaging pools valid retrieval hits across users, whereas macro averaging first computes the RHR separately for each user and then averages across users. Results on PerMemBench-S. As shown in Table 2, Warm-Start UCB consistently outperforms the Universal baseline and all LLM-based gating methods across every memory budget under both micro and macro averaging. For example, at 𝐵 = 60, our method performs on par with the best LLM-based methods and improves over the Universal baseline by 6.66 micro and 7.06 macro percentage points. On the other hand, when the memory budget is large, LLM-based methods underperform the Universal baseline, while our method still improves over it by 1.80 micro and 1.08 macro points at 𝐵 = 250, and by 3.88 micro and 2.03 macro points at 𝐵 = 500. To summarize, the results indicate that our method consistently outperforms the considered baselines. Meanwhile, the comparison among the three linear methods further reveals the complementary roles of pre-training and online adaptation. Under the smallest budget, the feedback signal is sparse, so the improvement mostly comes from 9

Preprint from Salesforce AI Research Pre-trained Scorer. On the other hand, with a larger budget, more online feedback becomes available, which improves the performance of UCB. Since our proposed method combines the two strategies, it yields a stable improvement across all budget sizes. Results on PerMemBench-L. Table 3 shows a consistent pattern on the substantially longer interaction trajectories of PerMemBench-L. Our method, Warm-Start UCB, achieves the best non-oracle result in eight of the ten evaluation columns and the second-best result in the remaining two, where it is narrowly below Pre-trained Scorer at 𝐵 = 250. Compared with the Universal baseline, our method improves the micro RHR by 2.49–5.37 percentage points and the macro RHR by 1.81–5.37 points across the five evaluated budgets. Moreover, a longer interaction history means that both pre-training and online adaptation are exposed to more data and feedback. Nevertheless, the results show that, compared to the Universal baseline, Pre-trained Scorer leads to a considerable improvement, while UCB without initialization performs only marginally above the Universal baseline. Despite this, we still observe a consistent improvement brought by online adaptation when the memory budget is large. This result indicates that a population-level gating rule transfers useful information to a new user, but cannot fully capture the user’s individual storage preferences. Updating the rule with user-specific feedback is therefore essential for maintaining performance over long interaction horizons. Remark 5.1. Across all experimental settings, we see that LLM-based methods are comparable to the Universal baseline under small budgets. However, they become inferior to the Universal baseline under larger budgets, and the gap increases with the memory budget size. This observation is consistent with the results in In et al. (2026). We also observe that a more fine-grained, capacity-aware prompt only mitigates this inferiority but hardly resolves it. 5.3

Ablation Studies

Learning Dynamics. To examine whether the online methods genuinely adapt to the current user, we partition each user’s interaction trajectory into ten session-position deciles and report the F1-score improvement over the Universal baseline within each decile. Figure 1 presents the results on PerMemBench-L for representative memory budgets; the remaining budgets and the corresponding PerMemBench-S results are provided in Appendix C. As suggested by the plots, Cold-Start UCB begins with little or even negative improvement over the Universal baseline because it has not yet received enough one-sided, delayed feedback to learn the current user’s gating rule. Its performance gradually improves in later deciles, confirming that the online updates extract useful information from the user’s interaction history. In contrast, Pre-trained Scorer provides an immediate positive uplift but exhibits comparatively limited improvement as the interaction proceeds because its parameter remains fixed. Warm-Start UCB combines the strengths of both methods. It begins with the transferable knowledge supplied by pre-training and then increasingly separates from the frozen policy after user-specific feedback becomes available, also verifying the conceptual superiority of Warm-Start UCB. Effect of the Pre-training Pool Size. We next vary the number of user profiles included in the pre-training pool while keeping the held-out test user fixed. Figure 2 reports the best micro and macro RHRs on PerMemBench-L for pre-training pool sizes in {1, 2, 5, 10, 15}. Performance improves monotonically with the number of pre-training profiles under every displayed budget and both averaging schemes. Notably, the curves continue to improve at a pool size of 15, especially for 𝐵 = 500 and 𝐵 = 1000, suggesting that the pre-trained scorer has not yet saturated and may benefit from a larger and more diverse collection of user profiles.

6

Conclusion and Future Work

In this paper, we study the personalized memory gating problem. We first provide a retrieval-based utility labeling pipeline, based on which we construct a dataset for memory-gating evaluation. We further frame the problem as an apple tasting delayed feedback problem, and propose an algorithm for this problem that achieves an √ with one-sided
e 𝑇𝜖 ∧ (𝑑 𝑇 + 𝑚𝑑) regret guarantee, where 𝑚 is the delay window„ as demonstrated by our theoretical analysis. 𝑂 Experiments on our constructed dataset show that our method consistently outperforms all baselines, confirming its effectiveness. Currently, our method assumes a stable parameter for each user. Accounting for the evolution of the utility pattern remains an interesting direction for future work.

10

Preprint from Salesforce AI Research

References Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. Advances in neural information processing systems, 24, 2011. Noga Alon, Nicolo Cesa-Bianchi, Ofer Dekel, and Tomer Koren. Online learning with feedback graphs: Beyond bandits. In Conference on Learning Theory, pp. 23–35. PMLR, 2015. Byeongwoo An, Kapilan Balagopalan, Hyowon Wi, Sehwa Jeong, Jiun Jeong, Kyoungseok Jang, Noseong Park, GiSoo Kim, and Kwang-Sung Jun. Efficient algorithms for contextual apple tasting with log-loss. In Decision-Making from Offline Datasets to Online Adaptation: Black-Box Optimization to Reinforcement Learning, 2026. Gábor Bartók and Csaba Szepesvári. Partial monitoring with side information. In International Conference on Algorithmic Learning Theory, pp. 305–319. Springer, 2012. Gábor Bartók, Dean P Foster, Dávid Pál, Alexander Rakhlin, and Csaba Szepesvári. Partial monitoring—classification, regret bounds, and algorithms. Mathematics of Operations Research, 39(4):967–997, 2014. Yahav Bechavod and Aaron Roth. Individually fair learning with one-sided feedback. In International Conference on Machine Learning, pp. 1954–1977. PMLR, 2023. Yahav Bechavod, Katrina Ligett, Aaron Roth, Bo Waggoner, and Steven Wu. Equal opportunity in online classification with partial feedback. Advances in Neural Information Processing Systems, 32, 2019. Ilai Bistritz, Zhengyuan Zhou, Xi Chen, Nicholas Bambos, and Jose Blanchet. No weighted-regret learning in adversarial bandits with delays. Journal of Machine Learning Research, 23(139):1–43, 2022. Jose Blanchet, Renyuan Xu, and Zhengyuan Zhou. Delay-adaptive learning in generalized linear contextual bandits. Mathematics of Operations Research, 49(1):326–345, 2024. Leonardo Cella, Alessandro Lazaric, and Massimiliano Pontil. Meta-learning with stochastic linear bandits. In International Conference on Machine Learning, pp. 1360–1370. PMLR, 2020. Nicolò Cesa-Bianchi, Claudio Gentile, and Yishay Mansour. Delay and cooperation in nonstochastic bandits. Journal of Machine Learning Research, 20(17):1–38, 2019. Tommaso Cesari and Roberto Colomboni. arXiv:2606.03851, 2026.

Two-action apple tasting with switching costs.

arXiv preprint

Zachary Chase and Idan Mehalel. Deterministic apple tasting. In Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pp. 888–923. PMLR, 30 Jun–04 Jul 2025. Prateek Chhikara, Dev Khant, Saket Aryan, Taranjeet Singh, and Deshraj Yadav. Mem0: Building production-ready ai agents with scalable long-term memory. arXiv preprint arXiv:2504.19413, 2025. James A Grant and David S Leslie. Apple tasting revisited: Bayesian approaches to partially monitored online binary classification. arXiv preprint arXiv:2109.14412, 2021. Andras Gyorgy and Pooria Joulani. Adapting to delays and data in adversarial multi-armed bandits. In International Conference on Machine Learning, pp. 3988–3997. PMLR, 2021. Keegan Harris, Chara Podimata, and Steven Z Wu. Strategic apple tasting. Advances in Neural Information Processing Systems, 36:79918–79945, 2023. David P Helmbold, Nick Littlestone, and Philip M Long. Apple tasting and nearly one-sided learning. In Proceedings., 33rd Annual Symposium on Foundations of Computer Science, pp. 493–502. IEEE, 1992. David P Helmbold, Nicholas Littlestone, and Philip M Long. Apple tasting. Information and Computation, 161(2): 85–139, 2000. Benjamin Howson, Ciara Pike-Burke, and Sarah Filippi. Delayed feedback in generalised linear bandits revisited. In International Conference on Artificial Intelligence and Statistics, pp. 6095–6119. PMLR, 2023. 11

Preprint from Salesforce AI Research Daniel Hsu, Sham M Kakade, and Tong Zhang. Random design analysis of ridge regression. Foundations of Computational Mathematics, 14(3):569–600, 2014. Yuyang Hu, Shichun Liu, Yanwei Yue, Guibin Zhang, Boyang Liu, Fangyi Zhu, Jiahang Lin, Honglin Guo, Shihan Dou, Zhiheng Xi, et al. Memory in the age of ai agents. arXiv preprint arXiv:2512.13564, 2025. Yeonjun In, Wonjoong Kim, Sangwu Park, Kanghoon Yoon, and Chanyoung Park. Personalize-then-store: Benchmarking and learning personalized memory for long-horizon agents. arXiv preprint arXiv:2605.25535, 2026. Pooria Joulani, Andras Gyorgy, and Csaba Szepesvári. Online learning under delayed feedback. In International conference on machine learning, pp. 1453–1461. PMLR, 2013. Tal Lancewicki, Shahar Segal, Tomer Koren, and Yishay Mansour. Stochastic multi-armed bandits with unrestricted delay distributions. In International Conference on Machine Learning, pp. 5969–5978. PMLR, 2021. Mingfei Lu, Mengjia Wu, Feng Liu, Jiawei Xu, Weikai Li, Haoyang Wang, Zhengdong Hu, Ying Ding, Yizhou Sun, Jie Lu, et al. Choosing how to remember: Adaptive memory structures for llm agents. arXiv preprint arXiv:2602.14038, 2026. Saeed Masoudian, Julian Zimmert, and Yevgeny Seldin. A best-of-both-worlds algorithm for bandits with delayed feedback. Advances in Neural Information Processing Systems, 35:11752–11762, 2022. Charles Packer, Sarah Wooders, Kevin Lin, Vivian Fang, Shishir G Patil, Ion Stoica, and Joseph E Gonzalez. Memgpt: Towards llms as operating systems. arXiv preprint arXiv:2310.08560, 2023. Joon Sung Park, Joseph O’Brien, Carrie Jun Cai, Meredith Ringel Morris, Percy Liang, and Michael S Bernstein. Generative agents: Interactive simulacra of human behavior. In Proceedings of the 36th annual acm symposium on user interface software and technology, pp. 1–22, 2023. Amit Peleg, Naama Pearl, and Ron Meir. Metalearning linear bandits by prior update. In International Conference on Artificial Intelligence and Statistics, pp. 2885–2926. PMLR, 2022. Vinod Raman, Unique Subedi, Ananth Raman, and Ambuj Tewari. Apple tasting: Combinatorial dimensions and minimax rates. In The Thirty Seventh Annual Conference on Learning Theory, pp. 4358–4380. PMLR, 2024. Alireza Rezazadeh, Zichao Li, Wei Wei, and Yujia Bao. From isolated conversations to hierarchical schemas: Dynamic tree memory representation for llms. In International Conference on Learning Representations, volume 2025, pp. 990–1023, 2025. Ofir Schlisselberg, Ido Cohen, Tal Lancewicki, and Yishay Mansour. Delay as payoff in mab. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 20310–20317, 2025. Ofir Schlisselberg, Tal Lancewicki, Peter Auer, and Yishay Mansour. Improved best-of-both-worlds regret for bandits with delayed feedback. Advances in Neural Information Processing Systems, 38:2035–2072, 2026a. Ofir Schlisselberg, Mengxiao Zhang, and Yishay Mansour. Near-optimal stochastic linear bandits with delay. arXiv preprint arXiv:2606.16656, 2026b. Xiaolong Sun, Qichao Wang, Hangyu Li, and Liang Chen. Verifiable memory: Learning unified memory management with local and global verifiers for large language model agents. arXiv preprint arXiv:2608.03137, 2026. Zhen Tan, Jun Yan, I-Hung Hsu, Rujun Han, Zifeng Wang, Long Le, Yiwen Song, Yanfei Chen, Hamid Palangi, George Lee, et al. In prospect and retrospect: Reflective memory management for long-term personalized dialogue agents. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 8416–8439, 2025. Yifu Tang, Yingfei Wang, and Zeyu Zheng. Stochastic multi-armed bandits with strongly reward-dependent delays. In International Conference on Artificial Intelligence and Statistics, pp. 3043–3051. PMLR, 2024. Tobias Sommer Thune, Nicolò Cesa-Bianchi, and Yevgeny Seldin. Nonstochastic multiarmed bandits with unrestricted delays. Advances in Neural Information Processing Systems, 32, 2019. 12

Preprint from Salesforce AI Research Dirk Van Der Hoeven, Lukas Zierahn, Tal Lancewicki, Aviv Rosenberg, and Nicolò Cesa-Bianchi. A unified analysis of nonstochastic delayed feedback for combinatorial semi-bandits, linear bandits, and mdps. In The Thirty Sixth Annual Conference on Learning Theory, pp. 1285–1321. PMLR, 2023. Claire Vernade, Olivier Cappé, and Vianney Perchet. Stochastic bandit models for delayed conversions. In UAI, 2017. Claire Vernade, Alexandra Carpentier, Tor Lattimore, Giovanni Zappella, Beyza Ermis, and Michael Brueckner. Linear bandits with stochastic delayed feedback. In International Conference on Machine Learning, pp. 9712–9721. PMLR, 2020. Yancheng Wang, Ziyan Jiang, Zheng Chen, Fan Yang, Yingxue Zhou, Eunah Cho, Xing Fan, Yanbin Lu, Xiaojiang Huang, and Yingzhen Yang. Recmind: Large language model powered agent for recommendation. In Findings of the Association for Computational Linguistics: NAACL 2024, pp. 4351–4364, 2024a. Zihao Wang, Shaofei Cai, Anji Liu, Yonggang Jin, Jinbing Hou, Bowei Zhang, Haowei Lin, Zhaofeng He, Zilong Zheng, Yaodong Yang, et al. Jarvis-1: Open-world multi-task agents with memory-augmented multimodal language models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 47(3):1894–1907, 2024b. Zora Zhiruo Wang, Jiayuan Mao, Daniel Fried, and Graham Neubig. Agent workflow memory. In Forty-second International Conference on Machine Learning, 2025. Wujiang Xu, Zujie Liang, Kai Mei, Hang Gao, Juntao Tan, and Yongfeng Zhang. A-mem: Agentic memory for llm agents. Advances in Neural Information Processing Systems, 38:17577–17604, 2026. Sikuan Yan, Xiufeng Yang, Zuchao Huang, Ercong Nie, Zifeng Ding, Zonggen Li, Xiaowen Ma, Jinhe Bi, Kristian Kersting, Jeff Z Pan, et al. Memory-r1: Enhancing large language model agents to manage and utilize memories via reinforcement learning. In Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 12805–12825, 2026. Ke Yang, Zixi Chen, Xuan He, Jize Jiang, Michel Galley, Chenglong Wang, Jianfeng Gao, Jiawei Han, and ChengXiang Zhai. Plugmem: A task-agnostic plugin memory module for LLM agents. In Forty-third International Conference on Machine Learning, 2026. Yi Yu, Liuyi Yao, Yuexiang Xie, Qingquan Tan, Jiaqi Feng, Yaliang Li, and Libing Wu. Agentic memory: Learning unified long-term and short-term memory management for large language model agents. In Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 21457–21483, 2026. Mengxiao Zhang, Yuheng Zhang, Olga Vrousgou, Haipeng Luo, and Paul Mineiro. Practical contextual bandits with feedback graphs. Advances in Neural Information Processing Systems, 36:30592–30617, 2023. Longtao Zheng, Rundong Wang, Xinrun Wang, and Bo An. Synapse: Trajectory-as-exemplar prompting with memory for computer control. In International Conference on Learning Representations, volume 2024, pp. 19036–19066, 2024. Wanjun Zhong, Lianghong Guo, Qiqi Gao, He Ye, and Yanlin Wang. Memorybank: Enhancing large language models with long-term memory. In Proceedings of the AAAI conference on artificial intelligence, volume 38, pp. 19724–19731, 2024. Zhengyuan Zhou, Renyuan Xu, and Jose Blanchet. Learning in generalized linear contextual bandits with stochastic delays. Advances in Neural Information Processing Systems, 32, 2019. Zijian Zhou, Ao Qu, Zhaoxuan Wu, Sunghwan Kim, Alok Prakash, Daniela Rus, Bryan Kian Hsiang Low, and Paul Liang. Mem1: Learning to synergize memory and reasoning for efficient long-horizon agents. In International Conference on Learning Representations, volume 2026, pp. 58413–58438, 2026. Xizhou Zhu, Yuntao Chen, Hao Tian, Chenxin Tao, Weijie Su, Chenyu Yang, Gao Huang, Bin Li, Lewei Lu, Xiaogang Wang, et al. Ghost in the minecraft: Generally capable agents for open-world environments via large language models with text-based knowledge and memory. arXiv preprint arXiv:2305.17144, 2023. Julian Zimmert and Yevgeny Seldin. An optimal algorithm for adversarial bandits with arbitrary delays. In International Conference on Artificial Intelligence and Statistics, pp. 3285–3294. PMLR, 2020. 13

Preprint from Salesforce AI Research

A

Benchmark Construction Details

A.1

Prompts in Dataset Construction

In this section, we provide the detailed prompts used in the construction of the dataset. For the stages up to and including dialogue generation, we use the same prompts as In et al. (2026) and refer readers to that paper for details. Retrieval Trigger Prompt You are deciding, on behalf of a memory-augmented assistant, whether responding well to a user’s message requires recalling something from an earlier conversation with this same user. User’s message: {user_content} First, decide: does responding well to this message require recalling a specific fact, preference, plan, or detail the user has shared before – as opposed to being answerable from this message alone or from general knowledge? If a memory lookup is needed, compose a short, self-contained search query capturing specifically what needs to be recalled – not a copy of the whole message, just the retrieval-relevant part. Return strict JSON only, no other text: {{"query_needed": true or false, "query": "<query text, or empty string if not needed>"}}

Retrieval Verification Prompt You are judging whether retrieved memory facts are potentially helpful to respond to a user’s message in an ongoing conversation. User message: {query_text} Candidate fact: {candidate} For each candidate, decide whether it is potentially helpful to respond well to this message. A fact is necessary if omitting it would make the response worse in any of these ways: - Factually wrong or incomplete given what the user has told you - Generic or one-size-fits-all where a tailored response is possible - Inconsistent with a preference, constraint, or decision the user has already established - Missing personalization the user would reasonably expect, given they already shared this A fact is NOT necessary if it is merely on the same topic or in the same domain as the message, but doesn’t actually change what a good response would say – e.g. it’s superseded, tangential, or the message doesn’t call for that level of detail. Being topically related is not sufficient on its own, but genuine personalization value (not just factual correctness) does count as necessary. If a candidate is a close call, prefer to keep it – this check does not need to be the final word. Return strict JSON only, no other text, with the result and your reasoning: {{"helpful": <yes|no>, "reasoning": "<your reasoning>"}} (empty list if none are necessary)

Retrieval Recall Prompt You are judging retrieved memory facts against a user’s message in an ongoing conversation. None of these facts individually scored as a strong match, but they may still matter in combination. User message: {query_text} Candidate fact: {format_candidates(candidates)} Is there any subset of these facts that, used together, would be useful or necessary to respond to the user’s message? A single fact by itself not being a strong match does not rule out a combination of two or more being relevant. Return strict JSON only, no other text: {{"useful_entry_ids": ["<id>", ...]}} (empty list if no combination is useful)

Quality Assessment Prompt You are checking whether a specific piece of background information was actually used in an assistant’s reply within an ongoing conversation. Conversation so far: {session_prefix} Background fact under test: {fact_content} Assistant’s reply: {response}

14

Preprint from Salesforce AI Research

Classify how much the reply draws on that specific fact: - contains info from fact: the reply includes specific content, details, or claims that come from this fact (not just a generic reply that happens to be compatible with it). - weakly related: the reply touches the same general topic as the fact but doesn’t actually use its specific content. - irrelevant: the reply shows no sign of drawing on this fact at all. Return strict JSON only, no other text: {{"incorporation": "<contains_info_from_fact|weakly_related|irrelevant>"}}

A.2

Meta Evaluation

In this section, we conduct experiments to verify that the labels produced by the introduced pipeline genuinely reflect the actual utility during later interaction, rather than merely being an artifact of the selected memory system and LLM-judge. Method. We sample collection of probes (i.e., user’s prompt) with valid retrieval hits. For each sample, we sample one judge-confirmed fact and one unconfirmed candidate, either not retrieved or rejected by the LLM-judge. For each sampled fact, we place it in the context of an LLM (we used gpt-5-mini here) and draw 𝑘 completions at non-zero temperature. For each generated response, we employ gpt-5-mini as a separate judge model to decide if the generated response related to the candidate fact: contains the fact, weakly related or irrelevant, and mapping to a score within {1, 0.5, 0} and average over the 𝑘 responses. For the incorporation level test, we conduct a Wilcoxon signed-rank test over all pairs of (probe, average score of confirmed fact, average score of unconfirmed fact), and report the 𝑝-value, effect size 𝑟, and the 95% confidence interval of the scores. Table 4: Per-user fact-incorporation scores (mean, with 95% bootstrap CI) for confirmed vs. unconfirmed target facts, and the paired Wilcoxon test comparing them within each user. Reported as a robustness check on the pooled result, not as independent significance tests. User

𝑛

Confirmed

Unconfirmed

𝑟

𝑝

0045d22a 0097794e 009850d8 00aefb8e 00e85c7f

12 29 18 32 22

0.600 [0.42, 0.78] 0.510 [0.36, 0.66] 0.372 [0.20, 0.55] 0.456 [0.32, 0.60] 0.482 [0.35, 0.62]

0.167 [0.03, 0.33] 0.155 [0.07, 0.26] 0.144 [0.02, 0.31] 0.091 [0.03, 0.18] 0.077 [0.01, 0.15]

1.00 0.75 1.00 0.90 0.95

1.95 × 10−3 1.04 × 10−3 4.98 × 10−3 1.11 × 10−4 4.07 × 10−4

Pooled

113

0.477 [0.40, 0.55]

0.121 [0.08, 0.17]

0.88

6.95 × 10−13

Results. We sampled 5 users and pooled all 𝑛=113 matched confirmed and unconfirmed pairs across the five evaluated users, of which the results are presented in Table 4. The pooled results show that, compared to unconfirmed facts, confirmed facts are significantly more incorporated into the responses. Moreover, a paired Wilcoxon signed-rank test on the matched per-probe scores rejects the null of no difference (𝑊=220.5, 𝑝 ≈ 7 × 10−13 ) with a large matched-pairs rank-biserial effect size (𝑟=0.88), indicating a statistical significance. Moreover, when breaking down to the per-user results, we find that all of these results agree with the pooled results. For every user, confirmed facts are incorporated into responses much more often than unconfirmed facts. Mean scores range from 0.372 to 0.600 for confirmed facts and from 0.077 to 0.167 for unconfirmed facts, leading to a per-user gap of 0.23 to 0.43. Moreover, the effect size is at least 0.75 and 𝑝 < 0.005 for all users, indicating that the difference stays significant across five users, showing that the pooled result is not driven by any single user: facts with positive labels are consistently used in the responses. A.3

Statistics of the Constructed Benchmark

Tables 5 and 6 summarize some basic statistics of the constructed dataset. We highlight several notable properties that shape the gating problem. Sparsity of Utilized Facts. On average, only 14.2% of a user’s facts are utilized in PerMemBench-S and 19.9% in PerMemBench-L. Even the user with the highest ratio reaches only 23.6% and 25.6%, respectively. In other words, for typical user, more than three quarters of the extracted facts are never used to answer a later query. Consequently, a 15

Preprint from Salesforce AI Research

Table 5: Per-user statistics of PerMemBench-S, computed over the 20 users. Delay-to-first-hit is measured in sessions from extraction to the first valid retrieval hit. Metric

Mean

Median

Min

Max

Std

Sessions per uuid Total entries per uuid Entries with valid hits per uuid Hit rate (hit entries / total entries) Avg delay-to-first-hit (sessions)

54.25 856.5 124.5 0.142 13.47

57 894 122 0.140 13.02

36 474 45 0.064 6.95

70 1295 241 0.236 23.2

9.93 203.58 56.93 0.047 3.76

Table 6: Per-user statistics of PerMemBench-L, computed over the 20 users. Delay-to-first-hit is measured in sessions from extraction to the first valid retrieval hit. Metric

Mean

Median

Min

Max

Std

Sessions per uuid Total entries per uuid Entries with valid hits per uuid Hit rate (hit entries / total entries) Avg delay-to-first-hit (sessions)

150.35 2449.65 493.9 0.199 30.44

153 2435 518 0.203 28.64

94 1378 225 0.113 19.38

193 3341 780 0.256 41.08

30.84 550.39 170.45 0.043 5.86

policy that stores every fact fills a limited memory bank mostly with facts that are never used, and exhausts the bank capacity quickly. This leaves substantial room for a gating policy that learns to skip un-needed facts. Relatively Long Delay of Feedback. On average, a utilized fact receives its first valid hit 13.47 sessions after extraction in PerMemBench-S and 30.44 sessions in PerMemBench-L. These delays are roughly a quarter and a fifth of the average history length, respectively, and they grow with the length of the history rather than vanishing. With about 16 facts extracted per session on both benchmarks, this corresponds to several hundred subsequent facts, which exceeds some of the memory budgets we evaluate. This indicates that, feedback on an added fact arrives only after a substantial fraction of the interaction has passed, and might even get evicted before a meaningful signal is observed. Thus, the delay presents as an essential source of hardness for learning and motivates both the delay-aware estimator in Algorithm 1 and the warm start from other users’ data.

B

Experiment Details

B.1

Prompts of LLM-Based Methods

In this section, we provide the details of the prompts we used for LLM-greedy gating and LLM-context-aware gating. For each method, we consider two prompt templates: one directly adapted from In et al. (2026), and a fine-grained version that takes the remaining capacity of the memory bank into consideration. Greedy.

At each fact, the LLM receives the content of the current fact and outputs a binary decision.

LLM-Greedy Gating Prompt (In et al., 2026) You are a memory policy agent. Decide whether this fact should be stored in long-term memory. ## Fact {fact.get("content", "")} store = true: durable, worth-storing information about the user – part of an ongoing project, recurring goal, preference, or constraint. store = false: transient fact – trivial, one-off, or self-contained with no lasting relevance. Respond ONLY with: {{"store": <true|false>}}

16

Preprint from Salesforce AI Research LLM-Greedy Gating Prompt with Remaining Capacity You are the admission gate for a long-term memory bank in a memory-augmented assistant. The bank has a hard capacity limit – every fact you admit competes with every other fact for a slot. {capacity_block(history)} {GATING_RUBRIC} Fact: {fact.get("content", "")} First briefly note which KEEP or DROP category the fact best matches, then decide. Return strict JSON only, no other text: {{"category": "<short KEEP/DROP category label>", "store": true or false}}

Context-Aware. Before prediction, the current session is summarized in 1–2 sentences and appended to a running summary buffer. At each session, the LLM receives the current dialogue alongside the most recent K session summaries as historical context. LLM-Context-Aware Gating Prompt (In et al., 2026) You are a memory policy agent. Decide whether this fact should be stored in long-term memory. ## Recent Session Summaries {history_summaries} ## Current Session Dialogue {current_dialogue} ## Fact {fact.get("content", "")} store = true: durable, worth-storing information about the user – part of an ongoing project, recurring goal, preference, or constraint. store = false: transient fact – trivial, one-off, or self-contained with no lasting relevance. Respond ONLY with: {{"store": <true|false>}}

LLM-Context-Aware Gating Prompt with Remaining Capacity You are the admission gate for a long-term memory bank in a memory-augmented assistant. The bank has a hard capacity limit – every fact you admit competes with every other fact for a slot. {capacity_block(history)} {GATING_RUBRIC} ## Recent Session Summaries {history_summaries} ## Current Session Dialogue {current_dialogue} Fact: {fact.get("content", "")} First briefly note which KEEP or DROP category the fact best matches, then decide. Return strict JSON only, no other text: {{"category": "<short KEEP/DROP category label>", "store": true or false}}

The contents of the gating rubric and capacity block are given as follows. Gating Rubric KEEP (durable, user-specific – worth recalling in a future conversation): - Concrete preferences ("prefers X over Y", likes/dislikes) - Standing plans, goals, or decisions the user has committed to - Constraints or limitations (budget, health, schedule, skill level, access) - Relationships, roles, or people/pets recurring in the user’s life - Stable biographical or contextual facts (job, location, ongoing projects) DROP (low value – not worth the slot): - Small talk, pleasantries, filler, or meta-conversation - One-off transient state ("is tired right now", "it’s raining today")

17

Preprint from Salesforce AI Research

- Generic statements that aren’t specific to this user - Speculative or hypothetical musings the user hasn’t actually decided on - Facts that are near-duplicates of something already obviously established

Capacity Block If vacant Capacity check: the memory bank has plenty of free room right now, so bias toward KEEPING. Admit anything that plausibly fits a KEEP category below, even if it’s only moderately durable or specific – only DROP facts that clearly belong in the DROP list. If somewhat vacant Capacity check: the memory bank is filling up, so be selective. Admit only facts that clearly fit a KEEP category below with a reasonable chance of mattering later – DROP anything vague, borderline, or that fits the DROP list. If space limited Capacity check: the memory bank is nearly full, so be highly selective. Admit only facts you are confident are high-value KEEPs (firm decisions, explicit constraints or preferences the user is likely to reference again, major ongoing plans) – DROP anything minor, speculative, or merely plausible.

B.2

Hyper-parameters of Proposed Methods

In this section, we provide the implementation details and hyper-parameters of linear gating methods. We use OpenAI text-embedding-3-small to obtain an embedding of each fact and treat the embedding as the linear feature. For Warm-Start UCB and Pre-trained Scorer, for each given test user, we first fit a linear parameter on a training set composed of 15 other users randomly selected from the pool of 19 users. We then select the hyper-parameters by validating on another 4 users. We keep the routine of Algorithm 1 unchanged, while making the confidence radii 𝛼𝑡 and the gating threshold 𝜏 as hyper-parameters. We further simplify the initialization of the pre-trained covariance matrix 𝚺0 by setting it to identity matrix I, for Warm-Start UCB. For Cold-Start UCB, we consider the confidence radius 𝛼 ∈ {0.1, 0.2, 0.5, 1.0} and the gating threshold 𝜏 ∈ {0.05, 0.1, 0.2, 0.5} and 𝜆 ∈ {0.1, 1, 10}. For each fact, we wait for 𝑚 sessions to see if there is at least one valid retrieval. We consider 𝑚 ∈ {8, 10, 15} for PerMemBench-S and 𝑚 ∈ {25, 35, 45} for PerMemBench-L, to account for the longer delay. We report the best performance achieved by the combination of hyper-parameters that achieves the highest micro RHR. For Pre-trained Scorer, we sweep the gating threshold 𝜏 ∈ {0.05, 0.1, 0.2, 0.5}. For UCB with warm start, we adapt the gating threshold by considering 𝜏 ∈ {0.5𝜏pt , 𝜏pt , 2𝜏pt }, where 𝜏pt is the best validated value for Pre-trained Scorer. We further consider the confidence radius 𝛼 ∈ {0.1, 0.2, 0.5, 1.0}. For PerMemBench-L, we consider 𝜆 ∈ {200, 100, 50, 20} and 𝑚 ∈ {25, 35}. For PerMemBench-S, we account for the smaller amount of pre-training data and the shorter delay, and therefore consider 𝜆 ∈ {20, 10, 5, 2} and 𝑚 ∈ {10, 15}.

C

Additional Experiment Results

This section reports the learning-dynamics and pre-training pool-size results omitted from Section 5.3. Figures 4 and 5 show the per-decile F1 gain over the Universal baseline and the effect of the pre-training pool size on PerMemBench-S, respectively. The same set of results for the remaining budgets of PerMemBench-L (i.e., 250 and 800) is presented in Figure 3. On PerMemBench-L, the additional budgets reproduce the trends in Section 5.3. Pre-trained Scorer gives an immediate gain that shrinks in later deciles. Warm-Start UCB separates from it at later deciles. Cold-Start UCB remains close to the Universal baseline. RHRs also increase monotonically with the pool size under both budgets, improving by roughly 3–4 points from one to fifteen pre-training users. On PerMemBench-S, the same qualitative picture holds, and pre-training again gives a positive gain from the first decile. Warm-Start UCB and Pre-trained Scorer both outperform Cold-Start UCB for most of the interaction, and RHRs with fifteen pre-training users exceed those with a single user under every budget. The results are, however, noticeably noisier when the interaction history is short. In Figure 4, the confidence bands of Warm-Start UCB and Pre-trained Scorer largely overlap, and a clear separation between them appears only at 𝐵 = 60. Cold-start UCB fluctuates strongly and drops sharply in the ninth decile under all four budgets before recovering. In Figure 5, performance does not increase monotonically with the pool size: it dips at a pool size of five for 𝐵 = 60 and 𝐵 = 500 and stays nearly flat between 18

Preprint from Salesforce AI Research

0.06 0.04 0.02 0.00 0

1

2

3 4 5 6 session-position decile

7

8

0.06 0.04 0.02 0.00

9

0.40

0.71

0.39

Best Hit Rate

0.08

0.72

UCB (Cold) Pre-trained UCB (Warm)

0.08

Best Hit Rate

UCB (Cold) Pre-trained UCB (Warm)

0.10

F1 uplift over admit_all

F1 uplift over admit_all

0.12

0.38 0.37 0.36

0

1

2

(a) 250

3 4 5 6 session-position decile

7

8

1 2

9

(b) 800

5

10

Pretraining Pool Size

micro macro 15

0.70 0.69 0.68 0.67 1 2

5

10

Pretraining Pool Size

(c) 250

micro macro 15

(d) 800

0

1

2

3 4 5 6 7 session-position decile

(a) 60

8

9

0.05

UCB (Cold) Pre-trained UCB (Warm)

0

1

2

3 4 5 6 7 session-position decile

8

9

0.03

0.04 0.03 0.02 0.01 0.00

UCB (Cold) Pre-trained UCB (Warm)

−0.01 0

(b) 120

F1 uplift over admit_all

0.05 0.04 0.03 0.02 0.01 0.00 −0.01 −0.02

F1 uplift over admit_all

UCB (Cold) Pre-trained UCB (Warm)

0.10 0.08 0.06 0.04 0.02 0.00 −0.02 −0.04

F1 uplift over admit_all

F1 uplift over admit_all

Figure 3: Learning dynamics (3(a), 3(b)) and effect of the pre-training pool size (3(c), 3(d)) on PerMemBench-L for the remaining memory budgets 𝐵 ∈ {250, 800}. Shaded regions denote the 95% CI over all users.

1

2

3 4 5 6 7 session-position decile

(c) 250

8

9

0.02 0.01 0.00

−0.01

UCB (Cold) Pre-trained UCB (Warm)

−0.02 −0.03

0

1

2

3 4 5 6 7 session-position decile

8

9

(d) 500

Figure 4: Learning dynamics on PerMemBench-S: F1 gain over the Universal baseline in each session-position decile under memory budgets 𝐵 ∈ {60, 120, 250, 500}. Shaded regions denote the 95% CI over all users. two and ten users for 𝐵 = 120 and 𝐵 = 250. This is because, with a substantially shorter interaction history, the online estimator does not collect sufficient feedback to converge, and the offline pre-training pool is also much smaller. We therefore regard PerMemBench-L as the more reliable testbed for studying learning dynamics, and we report it in the main text.

D

Discussion of Assumption 4.5

In this section, we provide a two-step justification of Assumption 4.5. The first step controls the variation of the current user’s parameter around a population center, while the second controls the error incurred when estimating this population center from other users’ data. Let 𝜽 := E𝜽★∼𝜌 [𝜽 ★] denote the population-level parameter. Suppose first that the current user’s parameter is concentrated around 𝜽 in the geometry induced by 𝚺0 . In particular, assume that   P𝜽★∼𝜌 ∥𝜽 ★ − 𝜽 ∥ 𝚺0 ≤ 𝜖pop ≥ 1 − 𝛿pop . (D.1) The cross-user predictability observed in Section 3.2 provides empirical motivation for the existence of such a shared population component. At the same time, the gap between the within-user and leave-one-user-out performance suggests that 𝜖 pop need not be zero: it captures the irreducible heterogeneity between the current user and the population. We next justify that the population center 𝜽 can be estimated accurately from other users’ data. Suppose that the 𝑛 satisfies pretraining sample {(x𝑖 , 𝑦 𝑖 )}𝑖=1 𝑦 𝑖 = ⟨𝜽, x𝑖 ⟩ + 𝜉𝑖 ,

E[𝜉𝑖 | x𝑖 ] = 0,

where 𝜉𝑖 includes both observation noise and zero-mean variation among the pretraining users. Let 𝑛 ∑︁ b𝑛 := 1 𝚺 x𝑖 x⊤ 𝑖 , 𝑛 𝑖=1

𝚺 𝑥 := E[xx⊤ ],

and suppose 𝚺 𝑥 ⪰ 𝜇I for some 𝜇 > 0. Consider the ridge estimator with regularization 𝛾𝑛 to be determined later, b𝑛 + 𝛾𝑛 I 𝝁0 := 𝚺

𝑛
−1 1 ∑︁ x𝑖 𝑦 𝑖 . 𝑛 𝑖=1

19

0.28

0.43

0.27

0.42

0.63

0.87

0.25 0.24 1 2

5

10

Pretraining Pool Size

(a) 60

micro macro 15

0.41 0.40 0.39

1 2

5

10

Pretraining Pool Size

micro macro 15

Best Hit Rate

0.26

Best Hit Rate

0.62

Best Hit Rate

Best Hit Rate

Preprint from Salesforce AI Research

0.61 0.60 0.59 0.58 0.57

(b) 120

1 2

5

10

Pretraining Pool Size

micro macro 15

0.86 0.85 0.84 0.83

(c) 250

1 2

5

10

Pretraining Pool Size

micro macro 15

(d) 500

Figure 5: Effect of the pre-training pool size on PerMemBench-S: best micro- and macro-averaged RHRs under 𝐵 ∈ {60, 120, 250, 500}. Standard random-design ridge-regression results, such as Hsu et al. (2014, Theorem 2), imply that, when the sample size 𝑛 is sufficiently large4 , with probability at least 1 − 𝛿est , √︄
𝑑eff (𝛾𝑛 ) + log 1/𝛿est +𝜎 , (D.2) ∥ 𝝁0 − 𝜽 ∥ 𝚺 𝑥 ≲ 𝛾𝑛 ∥𝜽 ∥ 𝚺 −1 𝑥 𝑛 up to lower-order terms arising from estimating the covariance matrix, where 
−1  𝑑eff (𝛾𝑛 ) := tr 𝚺 𝑥 𝚺 𝑥 + 𝛾𝑛 I ≤ 𝑑. For the idealized choice 𝚺0 = 𝚺 𝑥 /𝜇, which satisfies 𝚺0 ⪰ I, (D.2) gives √︄
𝛾𝑛 𝜎 𝑑eff (𝛾𝑛 ) + log 1/𝛿est ∥ 𝝁0 − 𝜽 ∥ 𝚺0 ≲ √ ∥𝜽 ∥ 𝚺 −1 + = 𝜖est . √ 𝑥 𝜇 𝜇 𝑛

(D.3)

with probability at least 1 − 𝛿est . More generally, the same conclusion holds when 𝚺0 is an empirical covariance estimate that is spectrally equivalent to 𝚺 𝑥 /𝜇. Moreover, when 𝛾𝑛 is sufficiently small and √︁ 𝑛 sufficiently large, the e(𝜎 𝑑/(𝜇𝑛)). ridge bias is dominated by the stochastic error, leading to a sufficiently small error 𝜖 est = 𝑂 On the intersection of the events in (D.1) and (D.3), the triangle inequality gives ∥𝜽 ★ − 𝝁0 ∥ 𝚺0 ≤ ∥𝜽 ★ − 𝜽 ∥ 𝚺0 + ∥𝜽 − 𝝁0 ∥ 𝚺0 ≤ 𝜖pop + 𝜖 est . Therefore, we recover Assumption 4.5 by setting 𝜖 = 𝜖pop + 𝜖 est and 𝛿0 = 𝛿pop + 𝛿est :   P ∥𝜽 ★ − 𝝁0 ∥ 𝚺0 ≤ 𝜖pop + 𝜖 est ≥ 1 − 𝛿pop − 𝛿est . Remark D.1. While the estimation component 𝜖 est decreases as the amount of pretraining data grows, the population component 𝜖 pop captures genuine user heterogeneity and generally does not vanish with 𝑛.

E

Missing Proof from Section 4

In this section, we provide the proof of Theorem 4.7. We first define the following filtration. For 𝑡 = 1, 2, . . . , 𝑇, we consider F𝑡 generated by x1 , 𝑎 1 , 𝑦 1 , 𝐷 1 , · · · , x𝑡 −1 , 𝑎 𝑡 −1 , 𝑦 𝑡 −1 , 𝐷 𝑡 −1 , x𝑡 , 𝑎 𝑡 where 𝐷 𝑠 is defined as is if 𝑎 𝑠 = add and otherwise 0. For 𝑡 ≥ 1, we have x𝑡 ∈ F𝑡 , 𝑦 𝑡 , 𝐷 𝑡 ∈ F𝑡+1 , and 𝑧¯𝑠,𝑡 = 𝑦 𝑠 1{𝐷 𝑠 ≤ min{𝑡 − 𝑠, 𝑚}} ∈ F𝑠+1 . 4 Roughly speaking, we require 𝑛 ≳ 𝛾 −1 log(𝑑/ 𝛿 ). We refer to Hsu et al. (2014) for the exact burn-in cost. est 𝑛

20

Preprint from Salesforce AI Research E.1

Proof of Optimism

 Lemma E.1. Let the warm-start event E0 := ∥𝜽 ★ − 𝝁0 ∥ 𝚺0 ≤ 𝜖 . Then conditional on E0 , with probability at least 1 − 𝛿, simultaneously for all 𝑡 ∈ [𝑇], ∥ 𝜽b𝑡 − 𝜽 ★ 𝑚 ∥ V𝑡 ≤ 𝛼𝑡 , ★ where 𝜽 ★ 𝑚 := 𝜅 𝑚 𝜽 . Consequently, if P(E0 ) ≥ 1 − 𝛿0 , then the preceding event holds with probability at least 1 − 𝛿0 − 𝛿.

 Proof of Lemma E.1. Let S𝑡 := 𝑠 ≤ 𝑡 : 𝑎 𝑠 = add denote the set of rounds up to time 𝑡 in which the learner plays action add. We partition S𝑡 into the completed and pending sets   C𝑡 := 𝑠 ∈ S𝑡 : 𝑠 ≤ 𝑡 − 𝑚 , P𝑡 := 𝑠 ∈ S𝑡 : 𝑡 − 𝑚 < 𝑠 ≤ 𝑡 . For every 𝑠 ∈ C𝑡 , the 𝑚-round observation window has elapsed, and hence 𝑧¯𝑠,𝑡 = 𝑦 𝑠 1{𝐷 𝑠 ≤ 𝑚} =: 𝑧 𝑠(𝑚) . Because P(𝐷 𝑠 ≤ 𝑚) = 𝜅 𝑚 and 𝐷 𝑠 is conditionally independent of 𝑦 𝑠 given x𝑠 , we have     E 𝑧 𝑠(𝑚) | x𝑠 = 𝜅 𝑚 E 𝑦 𝑠 | x𝑠 = 𝜅 𝑚 𝜂(x𝑠 ) = ⟨𝜽 ★ 𝑚 , x𝑠 ⟩. Therefore, defining 𝜉 𝑠 := 𝑧 𝑠(𝑚) − ⟨𝜽 ★ 𝑚 , x𝑠 ⟩, we obtain a conditionally mean-zero and 𝜎-sub-Gaussian noise sequence. Now define the completed-sample design matrix and its corresponding ridge estimator ! ∑︁ ∑︁ −1 (𝑚) ⊤ x𝑠 𝑧 𝑠 . V𝑡 := 𝜆𝚺0 + x𝑠 x𝑠 , 𝜽 𝑡 := V𝑡 𝜆𝚺0 (𝜅 𝑚 𝝁0 ) + 𝑠∈ C𝑡

𝑠∈ C𝑡

Subtracting 𝜽 ★ 𝑚 gives ∑︁
★
V𝑡 𝜽 𝑡 − 𝜽 ★ x𝑠 𝜉 𝑠 . 𝑚 = 𝜆𝚺 0 𝜅 𝑚 𝝁0 − 𝜽 𝑚 + 𝑠∈ C𝑡 ★ Since 𝜽 ★ 𝑚 = 𝜅 𝑚 𝜽 , the prior-mismatch term satisfies, on E0 , √ √ ★ ∥𝜅 𝑚 𝝁0 − 𝜽 ★ 𝑚 ∥ 𝜆𝚺0 = 𝜅 𝑚 𝜆∥ 𝝁0 − 𝜽 ∥ 𝚺0 ≤ 𝜅 𝑚 𝜆𝜖 .

We next bound the determinant term in the self-normalized inequality. Since x𝑠 ∈ F𝑠 and 𝜉 𝑠 ∈ F𝑠+1 is 𝜎-sub-Gaussian with 𝜎 = 1/2 condition on F𝑠 , we can apply Lemma F.2, which leads to that with probability at least 1 − 𝛿, for all 𝑡 ≥ 0   ∑︁ 2 det(V𝑡 ) 1/2 det(𝜆𝚺0 ) −1/2 2 x𝑠 𝜉 𝑠 −1 ≤ 2𝜎 log V𝑡 𝛿 𝑠∈ C 𝑡

Since 𝚺0 ⪰ I and ∥x𝑠 ∥ 2 ≤ 𝐿, Lemma F.3 leads to !   det(V𝑡 ) |C𝑡 |𝐿 2  𝑇 𝐿2  log ≤ 𝑑 log 1 + ≤ 𝑑 log 1 + . det(𝜆𝚺0 ) 𝜆𝑑 𝜆𝑑 Combining these estimation therefore implies that, with probability at least 1 − 𝛿, simultaneously for all 𝑡 ∈ [𝑇], √︂  1 ∑︁ 𝑇 𝐿2  + 2 log . x𝑠 𝜉 𝑠 −1 ≤ 𝜎 𝑑 log 1 + V𝑡 𝜆𝑑 𝛿 𝑠∈ C 𝑡

Combining the preceding inequality with the prior-mismatch bound yields √︂  1 √ 𝑇 𝐿2  ★ ∥𝜽 𝑡 − 𝜽 𝑚 ∥ V𝑡 ≤ 𝛽𝜆, 𝛿 := 𝜎 𝑑 log 1 + + 2 log + 𝜅 𝑚 𝜆𝜖 𝜆𝑑 𝛿 21

Preprint from Salesforce AI Research simultaneously for all 𝑡 ∈ [𝑇]. It remains to translate 𝜽 𝑡 to 𝜽b𝑡 . By the definition of Algorithm 1, we know that 𝜽b𝑡 = V𝑡−1 𝒃 𝑡 , which leads to ★ V𝑡 ( 𝜽b𝑡 − 𝜽 ★ 𝑚 ) = 𝒃 𝑡 − V𝑡 𝜽 𝑚

= 𝜆𝚺0 (𝜅 𝑚 𝝁0 ) +

∑︁ 𝑠∈ C𝑡

−1

Since 𝜽 𝑡 = V𝑡

𝜆𝚺0 (𝜅 𝑚 𝝁0 ) +

Í

∑︁

x𝑠 𝑧¯𝑠,𝑡 +

  ∑︁ ∑︁ ⊤ ★ x𝑠 𝑧¯𝑠,𝑡 − 𝜆𝚺0 + x𝑠 x⊤ + x x 𝑠 𝑠 𝑠 𝜽𝑚.

𝑠∈ P𝑡

𝑠∈ C𝑡

𝑠∈ P𝑡

𝑠∈ C𝑡 x𝑠 𝑧¯ 𝑠,𝑡 , we know that

∑︁

★ V𝑡 ( 𝜽b𝑡 − 𝜽 ★ 𝑚 ) = V𝑡 (𝜽 𝑡 − 𝜽 𝑚 ) +

x𝑠 𝑧¯𝑠,𝑡 −

𝑠∈ P𝑡

∑︁

= V𝑡 (𝜽 𝑡 − 𝜽 ★ 𝑚) +

 ∑︁

 ★ x𝑠 x⊤ 𝑠 𝜽𝑚

𝑠∈ P𝑡

x𝑠 𝑧¯𝑠,𝑡 − ⟨𝜽 ★ 𝑚 , x𝑠 ⟩ .

𝑠∈ P𝑡

Therefore, the triangle inequality gives ★ ★
∥ 𝜽b𝑡 − 𝜽 ★ 𝑚 ∥ V𝑡 = ∥𝒃 𝑡 − V𝑡 𝜽 𝑚 ∥ V𝑡−1 ≤ ∥V𝑡 𝜽 𝑡 − 𝜽 𝑚 ∥ V𝑡−1 +

∑︁

  x𝑠 𝑧¯𝑠,𝑡 − ⟨𝜽 ★ , x ⟩ 𝑠 𝑚

𝑠∈ P𝑡

. V𝑡−1

−1

Because V𝑡 ⪰ V𝑡 , we have V𝑡−1 ⪯ V𝑡 . Consequently,
2 2 ★
⊤ −1 ★
★ 2 ∥V𝑡 𝜽 𝑡 − 𝜽 ★ ≤ 𝛽𝜆, 𝑚 ∥ V −1 = 𝜽 𝑡 − 𝜽 𝑚 V𝑡 V𝑡 V𝑡 𝜽 𝑡 − 𝜽 𝑚 ≤ ∥𝜽 𝑡 − 𝜽 𝑚 ∥ 𝛿. V𝑡

𝑡

For each 𝑠 ∈ P𝑡 , both 𝑧¯𝑠,𝑡 and ⟨𝜽 ★ 𝑚 , x𝑠 ⟩ = 𝜅 𝑚 𝜂(x𝑠 ) belong to [0, 1]. Therefore, 𝑧¯𝑠,𝑡 − ⟨𝜽 ★ 𝑚 , x𝑠 ⟩ ≤ 1. The triangle inequality then gives  ∑︁  x𝑠 𝑧¯𝑠,𝑡 − ⟨𝜽 ★ , x ⟩ 𝑚 𝑠 𝑠∈ P𝑡

V𝑡−1

≤

∑︁

∥x𝑠 ∥ V𝑡−1 𝑧¯𝑠,𝑡 − ⟨𝜽 ★ 𝑚 , x𝑠 ⟩ ≤

𝑠∈ P𝑡

∑︁

∥x𝑠 ∥ V𝑡−1 .

𝑠∈ P𝑡

Combining the last three displays, we conclude that, simultaneously for all 𝑡 ∈ [𝑇], ∑︁ ∥x𝑠 ∥ V𝑡−1 = 𝛼𝑡 . ∥ 𝜽b𝑡 − 𝜽 ★ 𝑚 ∥ V𝑡 ≤ 𝛽𝜆, 𝛿 + 𝑠∈ P𝑡

This proves the conditional statement. Finally, combining the online concentration event with P(E0 ) ≥ 1 − 𝛿0 proves the unconditional probability bound 1 − 𝛿0 − 𝛿. □ E.2

Main Theorem

Proof of Theorem 4.7. Throughout the proof, we work on the good event E from Lemma E.1. On this event, simultaneously for all 𝑡 ∈ [𝑇], ∑︁ ∥ 𝜽b𝑡 − 𝜽 ★ ∥x𝑠 ∥ V𝑡−1 , 𝑚 ∥ V𝑡 ≤ 𝛼𝑡 = 𝛽𝜆, 𝛿 + 𝑠∈ P𝑡 ★ −1 e is optimistic for 𝜂(x ). Fix any 𝑡 ∈ [𝑇]. By Cauchy–Schwarz where 𝜽 ★ 𝑡 𝑡 𝑚 = 𝜅 𝑚 𝜽 . We first show that the index 𝜅 𝑚 𝜂 and the definition of E, ★ b ⟨𝜽b𝑡 −1 − 𝜽 ★ 𝑚 , x𝑡 ⟩ ≥ −∥ 𝜽 𝑡 −1 − 𝜽 𝑚 ∥ V𝑡 −1 ∥x𝑡 ∥ V −1 ≥ −𝛼𝑡 −1 ∥x𝑡 ∥ V −1 . 𝑡 −1

𝑡 −1

Recalling the computation of 𝜂e𝑡 , it follows that 𝜂e𝑡 = ⟨𝜽b𝑡 −1 , x𝑡 ⟩ + 𝛼𝑡 −1 ∥x𝑡 ∥ V−1 ≥ ⟨𝜽 ★ 𝑚 , x𝑡 ⟩ = 𝜅 𝑚 𝜂(x𝑡 ). 𝑡 −1

22

(★)

Preprint from Salesforce AI Research Suppose first that 𝜂(x𝑡 ) ≥ 𝜏. By (★), 𝜂e𝑡 ≥ 𝜅 𝑚 𝜂(x𝑡 ) ≥ 𝜅 𝑚 𝜏, so Algorithm 1 chooses 𝑎 𝑡 = add, which agrees with the oracle policy. Thus, the algorithm never incorrectly skips a context that the oracle would add. Therefore, the only case in which the algorithm incurs nonzero regret is 𝜂(x𝑡 ) < 𝜏,

𝑎 𝑡 = add.

In this case, the decision rule implies 𝜂e𝑡 ≥ 𝜅 𝑚 𝜏, and the per-step regret is given by     −1 E 𝑦 𝑟 (x𝑡 , skip, 𝑦) − E 𝑦 𝑟 (x𝑡 , add, 𝑦) = 𝜏 − 𝜂(x𝑡 ) ≤ 𝜅 𝑚 𝜂e𝑡 − 𝜂(x𝑡 ) Moreover, we see that 𝜂e𝑡 − 𝜅 𝑚 𝜂(x𝑡 ) = ⟨𝜽b𝑡 −1 , x𝑡 ⟩ + 𝛼𝑡 −1 ∥x𝑡 ∥ V−1 − ⟨𝜽 ★ 𝑚 , x𝑡 ⟩ 𝑡 −1

= ⟨𝜽b𝑡 −1 − 𝜽 ★ 𝑚 , x𝑡 ⟩ + 𝛼𝑡 −1 ∥x𝑡 ∥ V −1

𝑡 −1

≤ 2𝛼𝑡 −1 ∥x𝑡 ∥ V−1 . 𝑡 −1

Consequently,     2 E 𝑦 𝑟 (x𝑡 , skip, 𝑦) − E 𝑦 𝑟 (x𝑡 , add, 𝑦) ≤ 𝛼𝑡 −1 ∥x𝑡 ∥ V−1 1{𝑎 𝑡 = 1}. 𝑡 −1 𝜅𝑚  Now let S𝑇 := 𝑡 ∈ [𝑇] : 𝑎 𝑡 = add denote the set of rounds with action add. Summing the preceding inequality gives Regret(𝑇) ≤

2 ∑︁ 𝛼𝑡 −1 ∥x𝑡 ∥ V−1 . 𝑡 −1 𝜅𝑚 𝑡 ∈ S𝑇

Substituting the definition of 𝛼𝑡 −1 , we decompose the sum as ! ∑︁

𝛼𝑡 −1 ∥x𝑡 ∥ V−1 = 𝛽𝜆, 𝛿

∑︁

𝑡 −1

𝑡 ∈ S𝑇

∥x𝑡 ∥ V−1 +

∑︁

𝑡 ∈ S𝑇

𝑠∈ P𝑡 −1

} |

{z (I)

∥x𝑠 ∥ V−1 ∥x𝑡 ∥ V−1 . 𝑡 −1

𝑡 −1

𝑡 ∈ S𝑇

|

∑︁

{z (II)

𝑡 −1

}

We first bound (I). Since 𝜆 ≥ 𝐿 2 and 𝚺0 ⪰ I, we see that ∥x𝑡 ∥ 2V−1 ≤ ∥x𝑡 ∥ 22 /𝜆 ≤ 1. This allows us to apply the 𝑡 −1

elliptical-potential lemma (Lemma F.1) to the subsequence of added contexts, which yields !  ∑︁ det(V𝑇 ) 𝑇 𝐿2  2 ∥x𝑡 ∥ V−1 ≤ 2 log ≤ 2𝑑 log 1 + , det(𝜆𝚺0 ) 𝜆𝑑 𝑡 −1 𝑡 ∈ S𝑇

where the last inequality holds due to Lemma F.3. Therefore, Cauchy–Schwarz implies √︂ √︄  ∑︁ ∑︁ 𝑇 𝐿2  ∥x𝑡 ∥ V−1 ≤ |S𝑇 | ∥x𝑡 ∥ 2 −1 ≤ 2𝑑𝑇 log 1 + . V𝑡 −1 𝑡 −1 𝜆𝑑 𝑡 ∈ S𝑇

𝑡 ∈ S𝑇

It follows that √︂ (I) ≤ 𝛽𝜆, 𝛿

 𝑇 𝐿2  2𝑑𝑇 log 1 + . 𝜆𝑑

We next bound (II). Applying 𝑎𝑏 ≤ (𝑎 2 + 𝑏 2 )/2 termwise yields (II) ≤


1 ∑︁ ∑︁  1 ∥x𝑠 ∥ 2V−1 + ∥x𝑡 ∥ 2V−1 =: (II) 𝑎 + (II) 𝑏 . 2 2 𝑡 −1 𝑡 −1 𝑠∈ P 𝑡 ∈ S𝑇

𝑡 −1

23

Preprint from Salesforce AI Research For every 𝑡, the pending set P𝑡 −1 contains at most 𝑚 elements. Hence, (II) 𝑏 ≤ 𝑚

∑︁ 𝑡 ∈ S𝑇

 𝑇 𝐿2  . ∥x𝑡 ∥ 2V−1 ≤ 2𝑚𝑑 log 1 + 𝜆𝑑 𝑡 −1

To bound (II) 𝑎 , we exchange the order of summation. Fix a round 𝑠 ∈ S𝑇 such that 𝑎 𝑠 = add. The term involving x𝑠 appears only for rounds 𝑡 satisfying 𝑠 ∈ P𝑡 −1 and 𝑎 𝑡 = add, which requires 𝑡 ∈ {𝑠 + 1, . . . , 𝑠 + 𝑚}. Thus, it appears at most 𝑚 times. Moreover, for every such 𝑡, we have 𝑡 − 1 ≥ 𝑠 and therefore V𝑡 −1 ⪰ V𝑠 ⪰ V𝑠−1 ⇒ ∥x𝑠 ∥ 2V−1 ≤ ∥x𝑠 ∥ 2V−1 . 𝑡 −1

𝑠−1

This immediately implies that (II) 𝑎 =

∑︁ ∑︁

∥x𝑠 ∥ 2V−1 ≤ 𝑚 𝑡 −1

𝑡 ∈ S𝑇 𝑠∈ P𝑡 −1

∑︁ 𝑠∈ S𝑇

 𝑇 𝐿2  ∥x𝑠 ∥ 2V−1 ≤ 2𝑚𝑑 log 1 + . 𝜆𝑑 𝑠−1

Combining the bounds for (II) 𝑎 and (II) 𝑏 gives  𝑇 𝐿2  . (II) ≤ 2𝑚𝑑 log 1 + 𝜆𝑑 Finally, combining the bounds for (I) and (II), we conclude that, on E, # " √︂   2 𝑇 𝐿2  𝑇 𝐿2  Regret(𝑇) ≤ + 2𝑚𝑑 log 1 + . 𝛽𝜆, 𝛿 2𝑑𝑇 log 1 + 𝜅𝑚 𝜆𝑑 𝜆𝑑 By Lemma E.1, the event E holds with probability at least 1 − 𝛿 conditional on the warm-start event E0 . Since P(E0 ) ≥ 1 − 𝛿0 , the claimed regret bound holds with probability at least 1 − 𝛿0 − 𝛿.
It remains to pick a proper 𝜆. For notational convenience, define ℓ𝜆 := log 1 + 𝑇 𝐿 2 /(𝜆𝑑) and 𝑞 := log(1/𝛿). We start with ( ) √  √︁ 𝑚𝑑 1  √︁ Regret(𝑇) ≲ inf 𝜎 𝑑ℓ𝜆 + 2𝑞 + 𝜅 𝑚 𝜆𝜖 𝑇 𝑑ℓ𝜆 + ℓ𝜆 . (E.1) 𝜅𝑚 𝜆≥ 𝐿 2 𝜅 𝑚 We first derive the warm-start bound. Using log(1 + 𝑢) ≤ 𝑢 gives ℓ𝜆 ≤ 𝑇 𝐿 2 /(𝜆𝑑). Consequently, the warm-start contribution satisfies √︂ √ √ √︁ 𝑇 2 𝐿2 𝜆𝜖 𝑇 𝑑ℓ𝜆 ≤ 𝜆𝜖 = 𝐿𝑇𝜖 . 𝜆 Meanwhile, as 𝜆 → ∞, we see that ℓ𝜆 → 0; thus, the remaining term vanishes: √︁ 𝜎 √︁ 𝑑ℓ𝜆 + 2𝑞 𝑇 𝑑ℓ𝜆 −→ 0, 𝜅𝑚

𝑚𝑑 ℓ𝜆 −→ 0. 𝜅𝑚

Taking the infimum over 𝜆 therefore yields Regret(𝑇) ≲ 𝐿𝑇𝜖 = 𝑂 (𝑇𝜖). We next derive the learning-based bound. Suppose 𝜖 > 0 and choose ( ) 𝑑 2 𝜆learn := max 𝐿 , 2 2 . 𝜅𝑚𝜖

24

Preprint from Salesforce AI Research √ −2 𝜖 −2 . For this choice, we see that When 𝜖 is sufficiently small such that 𝜖 ≲ 𝑑/(𝜅 𝑚 𝐿), this reduces to 𝜆learn = 𝑑𝜅 𝑚 √
√ 2 2 2 2 e 𝜅 𝑚 𝜆learn 𝜖 = 𝑑 and ℓ𝜆learn = log 1 + 𝜅 𝑚 𝜖 𝑇 𝐿 /𝑑 = 𝑂 (1). Substituting this choice into (E.1) gives √︁ √︁ 𝑑 𝑇ℓ𝜆learn 𝑚𝑑 𝜎 √︁ Regret(𝑇) ≲ + ℓ𝜆 𝑑ℓ𝜆learn + 2𝑞 𝑇 𝑑ℓ𝜆learn + 𝜅𝑚 𝜅𝑚 𝜅 𝑚 learn ! √ 𝑑 𝑇 + 𝑚𝑑 e , ≤𝑂 𝜅𝑚 where the last inequality suppresses logarithmic factors in 𝑇, 𝑑, 1/𝛿, and 1/𝜖. Combining this bound with the warm-start bound yields )! ( √ 𝑑 𝑇 + 𝑚𝑑 e min 𝑇𝜖, (E.2) Regret(𝑇) ≤ 𝑂 𝜅𝑚 which gives the desired bound. Finally, we characterize the phase transition in a principled way. The boundary between the √ so that √
𝜆 can be selected −1 𝑇 + 𝑚 . Letting 𝑢 = 𝑇, the boundary equation becomes 𝜅 𝑚 𝜖𝑢 2 = 𝑑𝑢 + 𝑑𝑚. two bounds is determined by 𝑇𝜖 ≍ 𝑑𝜅 𝑚 Solving this equation with respect to 𝑢 gives its positive solution √︁ 𝑑 + 𝑑 2 + 4𝜅 𝑚 𝜖 𝑑𝑚 . 𝑢= 2𝜅 𝑚 𝜖 Therefore, 𝑇crit ≍

√︂

𝑑

+

𝜅𝑚𝜖

!2

𝑑𝑚 𝜅𝑚𝜖

≍

𝑑2 𝜅 2𝑚 𝜖 2

+

𝑑𝑚 . 𝜅𝑚𝜖

−2 𝜖 −2 leads to the rate in (E.2). Therefore, setting 𝜆 → ∞ when 𝑇 ≤ 𝑇crit and 𝜆 = 𝑑𝜅 𝑚

F

□

Supplementary Lemmas

Lemma F.1 (Lemma 11, Abbasi-Yadkori et al. 2011). Let x1 , . . . , x𝑡 ∈ R𝑑 , and for 1 ≤ 𝑠 ≤ 𝑡 define V̄𝑠 = V +

𝑠 ∑︁

x𝑟 x𝑟⊤ ,

𝑟=1

where V ∈ R𝑑×𝑑 is a (deterministic) positive definite matrix. Then   𝑡 ∑︁ det( V̄𝑡 ) ≤ ∥x𝑠 ∥ 2V̄−1 . log det(V) 𝑠−1 𝑠=1 If, in addition, ∥x𝑠 ∥ 2 ≤ 𝐿 for all 𝑠 = 1, . . . , 𝑡, then 𝑡 ∑︁

  n o det( V̄𝑡 ) . min 1, ∥x𝑠 ∥ 2V̄−1 ≤ 2 log det(V) 𝑠−1 𝑠=1

∞ Lemma F.2 (Theorem 1, Abbasi-Yadkori et al. 2011). Let {F𝑡 }∞ 𝑡=0 be a filtration. Let {𝜂 𝑡 } 𝑡=1 be a real-valued stochastic process such that 𝜂𝑡 is F𝑡 -measurable and, for some 𝜎 ≥ 0, 𝜂𝑡 is conditionally 𝜎-sub-Gaussian given F𝑡 −1 , i.e.  𝜆2 𝜎 2    ∀𝜆 ∈ R, E 𝑒 𝜆𝜂𝑡 | F𝑡 −1 ≤ exp a.s. 2 𝑑 Let {x𝑡 }∞ 𝑡=1 be an R -valued stochastic process such that x𝑡 is F𝑡 −1 -measurable (i.e. x𝑡 may depend on the past but not on 𝜂𝑡 ). Let V ∈ R𝑑×𝑑 be any deterministic positive definite matrix. For any 𝑡 ≥ 0, define

V̄𝑡 = V +

𝑡 ∑︁

x𝑠 x⊤ 𝑠,

𝑠=1

𝑆𝑡 =

𝑡 ∑︁ 𝑠=1

25

𝜂 𝑠 x𝑠 .

Preprint from Salesforce AI Research Then, for any 𝛿 ∈ (0, 1), with probability at least 1 − 𝛿, for all 𝑡 ≥ 0 simultaneously,   det( V̄𝑡 ) 1/2 det(V) −1/2 2 2 𝑆𝑡 V̄−1 ≤ 2𝜎 log . 𝑡 𝛿 The following lemma provides the calculation of the log-determinant term appearing in Lemma F.1 and Lemma F.2. Lemma 𝚺0 ⪰ 𝜆I ≻ 0, and let x1 , . . . , x𝑇 ∈ R𝑑 satisfy ∥x𝑡 ∥ 2 ≤ 𝐿 for every 𝑡 ∈ [𝑇]. Define 𝚺𝑡 := Í𝑡 F.3. Let ⊤ 𝚺0 + 𝑠=1 x𝑠 x𝑠 . Then, for every 𝑡 ∈ [𝑇], 𝑡 𝐿2  𝑑  det(𝚺𝑡 )  𝑇 𝐿2  𝑑 ≤ 1+ ≤ 1+ . det(𝚺0 ) 𝑑𝜆 𝑑𝜆 Proof. Define M𝑡 := 𝚺0−1/2

𝑡 ∑︁

! x𝑠 x⊤ 𝑠

𝚺0−1/2 ⪰ 0.

𝑠=1

Since 𝚺0 ⪰ 𝜆I, we have 𝚺0−1 ⪯ 𝜆 −1 I, and hence tr(M𝑡 ) =

𝑡 ∑︁

1 −1 x⊤ 𝑠 𝚺 0 x𝑠 ≤

𝑡 ∑︁

𝜆 𝑠=1

𝑠=1

∥x𝑠 ∥ 22 ≤

𝑡 𝐿2 𝑇 𝐿2 ≤ . 𝜆 𝜆

Let 𝜇1 , . . . , 𝜇 𝑑 be the eigenvalues of M𝑡 , then we see that 𝑑 Ö det(𝚺𝑡 ) = det(I + M𝑡 ) = (1 + 𝜇𝑖 ). det(𝚺0 ) 𝑖=1

Applying AM–GM inequality,  ∑︁ 𝑑  𝑑 𝑑 1 tr(M𝑡 )  𝑑  𝑇 𝐿2  𝑑 det(𝚺𝑡 ) Ö (1 + 𝜇𝑖 ) ≤ (1 + 𝜇𝑖 ) = 1 + = ≤ 1+ , det(𝚺0 ) 𝑖=1 𝑑 𝑖=1 𝑑 𝑑𝜆 which finishes the proof.

□

26

Related documents

• Beyond Forgetting: Representation Misdirection Elicits Controllable Side Behaviors and Capabilities#1048566
• Heavy-Tailed Memory Traces in Long-Horizon Language Agents#1154633
• Scope Before You Persist: Preventing Cross-Family Interference in Agent Memory#1077073
• SkillApt: Learning When to Activate Agent Skills from Counterfactual Evidence#1063074
• On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models#1122567