In-Context Reward Adaptation for Robust Preference Modeling Zhenyu Sun1 , Zheng Xu2 , and Ermin Wei1 1
arXiv:2605.30323v1 [cs.LG] 28 May 2026
2
Northwestern University Meta Superintelligence Labs
Abstract Reinforcement Learning from Human Feedback (RLHF) typically relies on static reward models to align Large Language Models with human preferences. However, human values are inherently diverse and heterogeneous, and a single reward model often lacks the robustness required to generalize to unseen preference domains. While existing multi-reward frameworks attempt to address this, they are often restricted to a fixed set of known domains and fail to adapt to unseen human distributions without costly retraining. In this work, we propose In-Context Reward Adaptation, a transformer-based framework designed to model diverse and unseen human preferences on the fly. By leveraging the in-context learning capabilities of transformers, our approach adaptively infers the underlying reward structure from a small set of preference demonstrations. We demonstrate that while a standard transformer architecture is insufficient for this task by characterizing an asymptotic bias to the ground-truth, incorporating human response time as an auxiliary input signal enables the model to successfully adapt to preferences from previously unseen domains. Our findings show that this approach provides a more robust foundation for preference modeling, allowing for the representation of heterogeneous rewards and preference distribution shift, and offering a scalable path toward more flexible human-AI alignment.
1
Introduction
Reinforcement Learning from Human Feedback (RLHF) has become a central paradigm for aligning large language models (LLMs) with human intent, enabling models to be fine-tuned using reward functions learned from human preference data Korbak et al. [2023], Ouyang et al. [2022], Shaikh et al. [2024]. In most existing RLHF pipelines, a single reward model is trained on aggregated preference data and then used to guide policy optimization Ouyang et al. [2022], Wang et al. [2023]. While this approach has proven effective at scale, it implicitly assumes that human preferences are sufficiently stable and homogeneous to be represented by a single global objective. In practice, however, human preferences are neither homogeneous nor static. On the one hand, preferences can be dynamic even for the same individual: as users gain experience, encounter new contexts, or adapt their goals over time, their judgments and decision criteria may shift. On the other hand, preferences are highly heterogeneous across individuals Casper et al. [2023], Lambert and Calandra [2023], Li et al. [2024]. Differences in cultural background, personal values, domain expertise, and situational context can lead different humans to evaluate the same model output in drastically different ways. As a result, preference data collected from a population reflects a mixture of diverse and evolving reward structures rather than a single coherent objective. To better capture preference heterogeneity, recent work has proposed training multiple reward models, each corresponding to a predefined domain, annotator group, or objective Chakraborty et al. [2024], Ovadya [2023], Park et al. [2024], Rame et al. [2023]. While this strategy improves expressivity within known settings, it introduces several new challenges. First, the set of reward models must be specified in advance, making it difficult to accommodate dynamic changes in human preferences over time. Second, when new preference data are collected from previously unseen humans or domains, a new reward model typically needs to be trained or existing models must be retrained Jang et al. [2023], Park et al. [2024], Singh et al. [2025]. This process is computationally costly, requires repeated data collection, and does not scale well as the number of distinct preference types grows. As a result, multi-reward modeling remains brittle under task shift and offers limited flexibility in open-ended, real-world deployments. 1
In this work, we explore a different paradigm motivated by the in-context learning capabilities of transformer-based models to overcome the above-mentioned issues of multi-reward modeling. Incontext learning allows a pretrained transformer to adapt to new tasks at inference time by conditioning on a small number of examples, without requiring parameter updates Bai et al. [2023], Min et al. [2022], Zhang et al. [2024]. This property suggests a promising alternative for preference modeling, i.e., rather than training or retraining reward models whenever preferences change, a model could infer an intrinsic reward function on the fly from in-context preference demonstrations, which we refer to as in-context reward adaptation. However, a central finding of this paper is that naively applying in-context learning to pairwise preference data is insufficient. We show that when only binary comparative labels are available, incontext adaptation to unseen human reward models is fundamentally impossible, even with expressive transformer architectures and unlimited data. To resolve this problem, we introduce human response time as an auxiliary behavioral signal. Drawing inspiration from cognitive decision-making models Berlinghieri et al. [2023], Ratcliff and McKoon [2008], Wagenmakers et al. [2007], we show that response time encodes information about the strength of preferences, complementing the directional information provided by binary comparisons. We demonstrate theoretically that incorporating response time restores identifiability of reward parameters and enables correct in-context adaptation to previously unseen human preferences, which shows a promising approach for robust reward modeling under preference distribution shift. Experiments on both synthetic and real-world human decisionmaking datasets further validate our analysis, showing substantial improvements in robustness of our proposed in-context reward adaptation under preference distribution shift. In general, our results highlight fundamental limitations of existing reward modeling of RLHF pipelines based on static reward modeling and binary preferences. More broadly, this work suggests that scalable and robust human–AI alignment requires both adaptive learning mechanisms and richer forms of human feedback beyond binary comparisons.
Related Work Reward modeling with heterogeneous preferences. Aligning AI models to match diverse human preference data has been comprehensively studied in recent years. Chakraborty et al. [2024] shows a single reward model is insufficient to align with heterogeneous preferences, which inspires a min-max framework to train a compromised policy. Park et al. [2024] proposes a multi-reward learning approach based on representation learning and clustering and a preference aggregation approach to learn a single reward. Sorensen et al. [2024] considers a multi-objective reward modeling to balance several distinct objectives induced by diverse preferences. In Singh et al. [2025] a meta-learning approach is proposed to fit personalized reward models with few-shot preference demonstrations. Synthetic persona-guided methods are introduced in Ryan et al. [2025], Zhang [2024] in order to adapt to heterogeneous preferences. These works only consider reward adaptation and personalization to existing human preference data, or to unseen preferences but with asymptotic biases, while our method allows zero-bias adaptation to any preference data either seen or unseen. In-context learning. Transformer-based models have been observed to exhibit in-context learning abilities in natural language processing Dasgupta et al. [2022], Wei et al. [2022], Zhang et al. [2022]. Garg et al. [2022] initiates the study of in-context learning under a mathematical framework, showing transformers can in-contextly learn linear regression, two-larger ReLU networks, decision trees, etc. A line of works focuses on the construction of transformers, showing that the constructed transformers can perform as gradient descent Bai et al. [2023] and second-order algorithms Fu et al. [2024], Giannou et al. [2024]. Another line of research focuses on training dynamics and optimization landscape of transformers, where Zhang et al. [2024] and proves a single linear-attention layer can incontextly learn linear regression and is robust to task shift. Huang et al. [2023] shows similar results by further consider the additional softmax module. Recently, Shen et al. [2024] shows the in-context learning ability of linear attention transformers to classification problems by assuming mixture-ofGaussian posteriors of the feature space, while our results are not restricted to such assumptions on the feature space.
2
2
In-Context Learning Framework for RLHF with Heterogeneous Preferences
We study reward modeling of Reinforcement Learning from Human Feedback (RLHF) in a setting where human preferences are heterogeneous and drawn from an unknown population. Unlike classical formulations that assume a single latent reward function shared across all annotators, we explicitly model variability across persons and investigate whether a single in-context learner can adaptively infer a new person’s reward function from a small number of preference demonstrations.
2.1
Human Preference Model
Let x ∈ X denote a prompt, and let y0 , y1 ∈ Y(x) be two candidate responses generated by a pretrained language model. We assume the existence of a population of human types1 indexed by i, each characterized by an intrinsic reward function ri : X × Y → R. For a given human type i, preferences between candidate responses follow a Bradley–Terry (BT) model Bradley and Terry [1952]: p∗i (yw ≻ yl | x) = σ(ri (x, yw ) − ri (x, yl )) ,
(1)
where σ(t) = (1 + e−t )−1 is the sigmoid function, and yw , yl denote the preferred and unpreferred responses, respectively. Throughout the paper, we adopt a linear reward parameterization ri (x, y) = φ(x, y)⊤ θi∗ , where φ(x, y) ∈ Rd is a shared, known feature representation, and θi∗ ∈ Rd is an unknown parameter vector associated with human type i. The parameters θi∗ are drawn from an unknown population distribution H, which may have continuous support. For a pairwise comparison (x, y0 , y1 ), define a binary preference variable ( +1, if y1 ≻ y0 , zi = −1, if y0 ≻ y1 . Then the BT model can be equivalently written as ⊤ P(zi = 1 | x, y0 , y1 ) = σ φ(x, y1 ) − φ(x, y0 ) θi∗ .
(2)
For notational convenience, we define the difference feature
φ̃(x, y0 , y1 ) := φ(x, y1 ) − φ(x, y0 ) and hence
2.2
P(zi = 1 | x, y0 , y1 ) = σ φ̃(x, y0 , y1 )T θi∗ .
Limits of Multi-Reward Modeling
A common strategy for handling heterogeneous preferences is to train multiple reward models, each corresponding to a known human type or domain. While effective in-distribution, such approaches fundamentally assume access to a fixed and finite set of preference distributions during training Park et al. [2024], Singh et al. [2025]. In practice, however, the population distribution H may evolve over time due to changes in cultural context, task requirements, expertise, or prolonged interaction with the system. Or there might be new human annotators exhibit preferences that are totally unseen when training the multi-reward models. These circumstances introduce parameters θnew ∈ / H in practice. In such cases, static reward models suffer from out-of-distribution (OOD) degradation, and performance can only be recovered by retraining or fine-tuning on new preference data, which is costly, slow, and incompatible with rapid human-in-the-loop alignment. This motivates a different paradigm, i.e., instead of learning a fixed mapping from (x, y) to reward, we seek a model that is robust to OOD degradation θnew ∈ H and can adapt its effective reward function in context, conditioned on a small set of observed preferences. 1 The human type means a group of humans sharing the same reward model.
3
2.3
In-Context Reward Adaptation via Transformers
We leverage the in-context learning capabilities of transformers to perform on-the-fly inference of human reward parameters from preference demonstrations. Preference demonstrations. Fix a human type i. Consider a set of N prompt–response pairs l l {(xli , yi,0 , yi,1 )}N l=1 , for which preference data are collected. We additionally consider a query instance q q q (xi , yi,0 , yi,1 ), for which the goal is to predict the preference of the same human type i using only the in-context demonstrations. Prompt matrix construction. We encode the demonstration data into a prompt matrix 1 φi,0 φ2i,0 · · · φN φqi,0 i,0 φqi,1 , Ei = φ1i,1 φ2i,1 · · · φN i,1 1 2 N zi zi · · · zi ∗
l l where φli,0 := φ(xli , yi,0 ), φli,1 := φ(xli , yi,1 ), and the final entry ∗ corresponds to the unknown label for the query pair. To isolate preference-relevant structure, we apply a fixed linear transformation −Id Id 0 Ẽi = E, 0⊤ 0⊤ 1 i d d
so that each column of Ẽi contains the difference feature φ̃li = φli,1 − φli,0 together with its associated preference signal, i.e., after transformation 1 φ̃qi φ̃i φ̃2i · · · φ̃N i . Ẽi = 1 zi zi2 · · · ziN ∗ Transformer architecture. Following the linear-attention transformer formulation of Huang et al. [2023], Shen et al. [2024], Zhang et al. [2024], we consider a single-layer transformer of the form F (W V , W KQ ; Ẽi ) = Ẽi + W V Ẽi
Ẽi⊤ W KQ Ẽi . N
(3)
Similar to Ahn et al. [2023], Wu et al. [2023] we restrict the parameter matrices to 0 0d U 0d KQ W V = d×d , W = , 0⊤ 1 0⊤ 0 d d with U ∈ Rd×d trainable. Note that this structure simplifies the our theoretical analysis while preserving the critical properties of the transformer. Induced prediction. Under this architecture, the predicted preference for the query instance takes the form ! N 1 X l l ⊤ q q q q q P(ẑi = 1 | x , y0 , y1 ) = σ zi (φ̃ ) U φ̃ , (4) N l=1
where φ̃q = φ(xq , y1q ) − φ(xq , y0q ). otherwise.
2.4
We set ẑiq = 1 if the probability in (4) exceeds 1/2, and ẑiq = −1
Training Procedure
We assume that difference features are drawn from a distribution φ̃ ∼ Pφ̃ , independently of θi∗ ∼ H. Following Shen et al. [2024], Zhang et al. [2024], we evaluate performance in expectation over both human types and query instances, and minimize the expected cross-entropy loss h i 1 (5) LN (U ) = − Eθ∗ ∼H,φ̃∼P (1 + ziq ) log P(ẑiq = 1) + (1 − ziq ) log P(ẑiq = −1) , φ̃ 2 i where ziq is generated according to the ground-truth BT model (2), and P(ẑiq = ±1) is given by (4). 4
3
Impossibility of In-Context Reward Adaptation from Binary Preferences
In this section, we study the fundamental limits of in-context reward adaptation when the transformer is provided only with binary comparative preference labels. Despite the apparent flexibility of the in-context learning framework introduced in the previous section, we show that preference data alone are insufficient for robustly adapting to unseen human reward parameters. This negative result holds even under idealized conditions with infinite data and perfect optimization, revealing a structural obstruction rather than a finite-sample or algorithmic limitation. Our analysis proceeds in two stages. We first characterize the training dynamics of the transformer and establish convergence to a well-defined population-level objective. We then analyze inference on a previously unseen human and show that asymptotically correct prediction is generically impossible. For human type i we define the preference moment i h µi := Ezi ,φ̃ zi φ̃ | θi∗ .
We first analyze the optimization landscape induced by the expected training objective. See Appendix A for the proof. Theorem 1 (Asymptotic Optimality). Assume the following conditions hold: ⊤ 1. Ei∼H (µi µ⊤ i ) and Eφ̃∼P (φ̃φ̃ ) are full rank; φ̃
2. the feature difference satisfies kφ̃k ≤ B almost surely; 3. optimization is restricted to the bounded parameter set U = U ∈ Rd×d | kU kF ≤ R .
Then the following statements hold:
i. The training objective LN (U ) (5) is strongly convex with respect to U . ii. The population objective
is also strongly convex.
1 h q L̄(U ) := − E (1 + ziq ) log σ(µ⊤ i U φ̃i ) 2 i q + (1 − ziq ) log 1 − σ(µ⊤ U φ̃ ) i i
(6)
iii. The uniform convergence bound LN (U ) − L̄(U ) = O(N −1 ),
∀U ∈ Rd×d ,
holds. ∗ ∗ iv. Let UN and Ū ∗ denote the unique minimizers of LN and L̄, respectively. UN , Ū ∗ are unique and ∗ kUN − Ū ∗ kF = O(N −1 ).
Theorem 1 shows that the training problem is well behaved: the objective is strongly convex, admits a unique minimizer, and empirical risk minimization consistently recovers the population optimum. Consequently, standard optimization methods converge efficiently, and any failure of in-context adaptation cannot be attributed to optimization instability or finite-sample effects.
5
3.1
Inference on An Unseen Human
We now analyze inference for a previously unseen human. Suppose the transformer has been trained ∗ to obtain UN . We are given M preference samples collected from a new human with reward parameter θnew , without assuming θnew ∼ H. The transformer predicts the preference on a query instance (xq , y0q , y1q ) according to M X 1 ∗ q q z j (φ̃j )⊤ UN φ̃ , P(ẑnew = 1 | xq , y0q , y1q ) = σ M j=1 new
while the true preference probability is
q P(znew = 1 | xq , y0q , y1q ) = σ (φ̃q )⊤ θnew .
Define the empirical preference moment
M
µ̂new := By the law of large numbers,
1 X j z φ̃j . M j=1 new
p
→ µnew µ̂new − where
as M → ∞,
h i µnew = Eφ̃ tanh 21 φ̃⊤ θnew φ̃ .
∗ Together with the convergence UN → Ū ∗ from Theorem 1, we obtain
∀φ̃q ,
p
∗ q ∗ q → µ⊤ µ̂⊤ new Ū φ̃ , new UN φ̃ −
as N, M → ∞.
Since the sigmoid function is strictly monotone, asymptotically correct prediction for all φ̃q requires θnew = (Ū ∗ )⊤ µnew ,
∀θnew .
(7)
However, µnew is a nonlinear function of θnew , involving both the hyperbolic tangent and an expectation over the feature distribution. Condition (7) therefore requires a single linear operator (Ū ∗ )⊤ to invert this nonlinear mapping for all possible reward parameters, which is generically impossible.
3.2
Impossibility Result and Geometric Interpretation
We formalize this intuition in the following theorem (see proof in Appendix B). Theorem 2 (Failure of In-Context Reward Adaptation). There exist distributions Pφ̃ and H such that, for some θnew ∈ supp(H), lim dTV P̂new , Pnew | φ̃q > 0, ∀φ̃q , N,M→∞
where
q q dTV P̂new , Pnew | φ̃q = TV P(ẑnew | φ̃q ), P(znew | φ̃q ) .
Theorem 2 shows that even with infinitely many in-context demonstrations and perfect optimization, the transformer fails to correctly predict preferences for certain unseen human reward parameters. This impossibility result can be understood geometrically by examining how binary preference data embed reward parameters into the representation space used for in-context inference. Each human reward parameter θ ∈ Rd induces, through the Bradley–Terry model and the feature distribution Pφ̃ , a mapping defined by h i µ(θ) = Eφ̃ tanh 12 φ̃⊤ θ φ̃ .
This mapping θ 7→ µ(θ) summarizes all information about θ that can be recovered from infinitely many binary comparisons. Importantly, this mapping is generally nonlinear and compressive. 6
From a geometric perspective, as θ varies over Rd , the vectors µ(θ) do not fill the Rd space uniformly. Instead, they lie on a curved, typically low-dimensional manifold M = µ(θ) : θ ∈ Rd ⊂ Rd . Binary preference labels therefore embed the space of reward parameters into this nonlinear manifold, collapsing different directions of variation in θ into possibly indistinguishable representations. The in-context transformer observes a new human only through an empirical estimate of µ(θ) and applies a fixed linear decoding map (Ū ∗ )⊤ learned during training. Perfect in-context adaptation would require this linear map to act as a global inverse of the nonlinear embedding θ 7→ µ(θ), i.e., (Ū ∗ )⊤ µ(θ) = θ
for all θ.
Geometrically, this would require a single linear operator to decode a distorted, nonlinear embedding over the entire reward parameter space. Such an inversion is generically impossible unless µ(θ) is linear in θ or the set of admissible reward parameters is severely restricted. Theorem 2 reveals a fundamental limitation of in-context learning for RLHF, i.e., although transformers can aggregate and generalize patterns across demonstrations, they are constrained by the statistical sufficiency of the signals they receive. Binary comparative feedback collapses continuous reward parameters into nonlinear summary statistics that cannot be universally decoded by a fixed linear transformation. As a result, standard in-context learning architectures cannot robustly adapt to heterogeneous or previously unseen human preferences using preference labels alone. Motivation for introducing auxiliary signals. The geometric perspective also suggests a plausible solution to resolve the issue. The failure arises because binary comparisons encode only the sign of φ̃⊤ θ, discarding magnitude information. Auxiliary signals that correlate with the strength of preference can enrich the embedding θ 7→ µ(θ) and restore identifiability. In the next section, we show that incorporating response time effectively linearizes the embedding by recovering information about |φ̃⊤ θ|, thereby enabling robust in-context reward adaptation for unseen human preferences.
4
In-Context Reward Adaptation with Response Time
As established in Theorem 2, in-context learning based solely on binary preference labels cannot universally adapt to unseen human reward parameters. In this section, we show that incorporating human response time as an auxiliary signal resolves this obstruction and enables correct in-context reward adaptation. We model human decision-making using a standard drift–diffusion process following Berlinghieri et al. [2023], Li et al. [2024], Wagenmakers et al. [2007]. For a human of type i comparing responses y0 and y1 given prompt x, the response time is defined as ti (x, y0 , y1 ) := inf τ > 0 Si (τ ) ∈ {− 21 , 12 } , where the latent decision variable evolves according to
Si (τ ) = ri (x, y1 ) − ri (x, y0 ) τ + B(τ ),
where B(τ ) is a standard Brownian motion. Here ± 12 denote absorbing decision boundaries2 . Following Palmer et al. [2005], a direct calculation yields the conditional expectation 1 ⊤ ∗ 1 tanh φ̃ θi , φ̃⊤ θi∗ 6= 0, ⊤ θ∗ 2 2 φ̃ i E[ti | φ̃] = 1 , φ̃⊤ θi∗ = 0. 4 Recalling that E[zi | φ̃] = tanh 21 φ̃⊤ θi∗ , we obtain the key identity φ̃⊤ θi∗ =
1 E[zi | φ̃] . 2 E[ti | φ̃]
(8)
2 In Berlinghieri et al. [2023], Li et al. [2024] absorbing decision boundary level is set by a instead of 1/2 considered here. However, we note that we can always incorporate 2a into θi∗ , forming a new θ˜i∗ which is of interest, to make the setting consistent with ours.
7
Equation (8) reveals the central insight of this section: combining preference labels with response time recovers a linear signal in the reward parameter. In contrast to binary labels alone, which encode only the sign of φ̃⊤ θi∗ , the ratio zi /ti provides a continuous, magnitude-sensitive measurement (i.e., φ̃T θi∗ ) that eliminates the geometric ambiguity discussed in Section 3. Prompt construction with response time. Assume that for each human type i there are K independent annotators sharing the same reward parameter θi∗ . We construct the augmented prompt 1 φi,0 φ2i,0 · · · φN φqi,0 i,0 φ1i,1 φ2i,1 · · · φN φqi,1 i,1 , Ei,k = 1 2 N ti,k ti,k · · · ti,k ∗ 1 2 N zi,k zi,k · · · zi,k ∗ where (i, k) indexes the k-th annotator of type i. Aggregating over annotators yields K
zil :=
K
1 X l zi,k , K
tli :=
k=1
and the averaged prompt
1 φi,0 φ1 i,1 Ei = t1 i zi1
1 X l ti,k , K k=1
φ2i,0 φ2i,1 t2i zi2
··· ··· ··· ···
φN i,0 φN i,1 tN i ziN
φqi,0 φqi,1 . ∗ ∗
We construct the prompt matrix using the feature difference and ratio zi /ti as follows 1 φ̃qi φ̃i φ̃2i · · · φ̃N i , ziN Ẽi = zi1 zi2 ∗ · · · 1 2 N ti ti ti
where φ̃li = φli,1 − φli,0 .
Training objective and its asymptotic behavior. transformer architecture yields the prediction
Feeding Ẽi into the same linear-attention
N
ôqi =
1 X zil l ⊤ q (φ̃ ) U φ̃i . N tli i l=1
We train the model using the squared regression loss in expectation LN,K (U ) =
1 q q q 2 q E ∗ 1:N 1:N (ôi − (zi /ti )) 1:N 2 θi ∼H,φ̃i ,φ̃i ∼Pφ̃ ,zi ,ti
(9)
zil l q q T ∗ l=1 tli φ̃i , then ôi = ŝi U φ̃i . Conditioned on θi , according to the Law of Large Numbers (and noting i.i.d. φ̃li and independence from θi∗ ),
Defining ŝi = N1
PN
p
ŝi → Eφ̃ [2φ̃T θi∗ φ̃] = 2Σφ̃ θi∗ , as N, K → ∞ where Σφ̃ = E[φ̃φ̃T ]. Also, for the query, ziq p → 2(φ̃qi )T θi∗ as K → ∞. tqi Thus, for large N and K, LN,K converges to L∞ , where L∞ (U ) :=
lim LN,K (U ) i 1 h = E ((2Σφ̃ θi∗ )T U φ̃q − 2(φ̃q )T θi∗ )2 2 h i = 2E ((θi∗ )T (Σφ̃ U − I)φ̃q )2 N,K→∞
8
(10)
One can easily observe that a minimizer of L∞ is given by U ∗ = Σφ̃−1 (which is in fact also unique as shown in Theorem) if assuming Σφ̃ is full-rank, which should match the minimizer of LN,K as N, K → ∞. This is formally presented in the following theorem (see Appendix C for proof). Theorem 3. Consider LN,K (U ) and L∞ (U ) defined by (9) and (10), respectively. Assuming that Σφ̃ := E[φ̃φ̃T ] and E[θi∗ (θi∗ )T ] are full-rank; kφ̃k ≤ B, kU kF ≤ R, then the following hold:
(i). LN,K and L∞ are strongly-convex;
(ii). U ∗ = Σφ̃−1 is the unique minimizer of L∞ ; √ √ ∗ ∗ (iii). kUN,K − U ∗ kF ≤ O(1/ N + 1/ K), where UN,K is the unique minimizer of LN,K . Correct in-context adaptation for unseen humans. Now let us consider a new human type with arbitrary θnew which is not necessarily from H. Given M in-context samples drawn from that ∗ human type with K human users, as UN,K → U ∗ = Σφ̃−1 with N, K → ∞, p
T φ̃q ôqnew → (2Σφ̃ θnew )T Σφ̃−1 φ̃q = 2θnew p
q T which matches the true target as znew /tqnew → 2θnew φ̃q , K → ∞. Therefore, the transformer is able to in-contextly adapt to human θnew ’s preference model without retraining. The above discussion is formally characterized by the following statement (see Appendix D for proof).
Corollary 1 (Correct In-Context Adaptation). Given arbitrary θnew , let ôqnew be the prediction of the ∗ transformer after complete training, i.e., with UN,K being the minimizer of LN,K , M
ôqnew :=
l 1 X znew ∗ (φ̃l )T UN,K φ̃q l M tnew l=1
PK PK l l , tlnew = (1/K) k=1 tlnew,k are averages over K humans drawn from where znew = (1/K) k=1 znew,k the human type θnew . Then, T φ̃q )2 ≤ O(M −1 + N −1 + K −1 + (N K)−1/2 ). E(ôqnew − 2θnew
In summary, these results show that response time resolves the limitation of in-context learning with binary preference data in Section 3. By enriching binary comparisons with a continuous signal proportional to preference strength, the transformer can recover a linear embedding of reward parameters and achieve correct in-context adaptation to unseen human preferences.
5
Experiments
We evaluate the proposed in-context reward adaptation framework on both a controlled synthetic dataset and a real-world human preference dataset with recorded response times. The synthetic experiments are designed to closely mirror the theoretical setting and allow precise validation under known ground-truth reward models, while the real-world experiments demonstrate the applicability of response-time–augmented in-context learning on behavioral data collected from human subjects.
5.1
Synthetic Dataset
For the synthetic setting, we construct a preference learning environment that follows the setting of our problem. For each query, we independently sample two feature vectors φ0 , φ1 ∈ Rd from some fixed distributions, and define the feature difference φ̃ := φ1 − φ0 . Human reward parameters θi∗ are sampled from a mixture of two Gaussian distributions, inducing heterogeneity across human types. To in-context learning ability to unseen preferences, we additionally sample a new human type θnew from a third Gaussian distribution that is disjoint from the training mixture, thereby creating an explicit out-of-distribution (OOD) test setting. To obtain response times, we simulate a drift–diffusion process as described in Section 4. Specifically, for each human type, we simulate K independent annotators by running the drift–diffusion process K times per query. 9
Although our theoretical results focus on a linear-attention transformer for analytical tractability, we also implement the in-context learning framework using a GPT-2 model with 124M parameters. Both linear attention model and GPT-2 model are trained on the same synthesized preference data, using either binary preference labels alone or the response-time–augmented targets introduced in Section 4. Table 1: Test accuracy under in-distribution (ID) and out-of-distribution (OOD) preference data for linear attention and GPT2 Setting
LinearAttn
GPT2
w/o resp (ID) w/o resp (OOD)
0.936 0.783
0.925 0.694
w/ resp (ID) w/ resp (OOD)
0.878 0.891
0.905 0.875
Numerical results are presented in Table 1, where ”resp” represents ”response time”; ”ID” or ”OOD” means that the new human parameter is drawn from the same distribution as used for training or from a different distribution. When trained without response time, both linear attention and GPT2 perform well on in-distribution preference data but exhibit a pronounced degradation under OOD preferences. This behavior is consistent with our theoretical analysis, which predicts that binary preference labels alone are insufficient for in-context adaptation to unseen reward parameters. Incorporating response time substantially improves OOD performance for both architectures, bringing test accuracy close to in-distribution levels. This supports the central claim of the paper: response time resolves the information bottleneck inherent in binary comparisons and enables robust in-context reward adaptation to unseen human preferences. Notably, the same qualitative trend appears for both linear attention and GPT-2, indicating that the failure of binary-only in-context learning and the benefits of response time are not artifacts of limited model capacity, but stem from fundamental properties of preference information.
5.2
Real-World Dataset: Food-Risk Preferences
To evaluate our approach on real human data, we adopt the food-risk dataset introduced by Smith and Krajbich [2018] and used in prior work. The dataset consists of binary choices and response times collected from 42 participants, each responding to between 60 and 200 queries. Each query presents two arms, where each arm contains two food items. By selecting an arm, participants receive one of the two food items uniformly at random. In addition to choices and response times, participants’ eye movements were recorded during the experiment. Table 2: Test accuracy under in-distribution (ID) and out-of-distribution (OOD) food-risk preference data for GPT2 across different inference lengths Setting
M =4
M =8
M = 16
w/o resp (ID) w/o resp (OOD)
0.606 0.581
0.679 0.625
0.675 0.631
w/ resp (ID) w/ resp (OOD)
0.633 0.605
0.681 0.667
0.710 0.705
Prior to the choice task, participants rated all food items on a discrete scale from −10 to 10. Following prior work Li et al. [2024], we define each arm’s feature vector as the participant’s ratings of the two food items in that arm, augmented with second-order polynomial features. This yields a feature representation in R5 .
10
75
Accuracy (%)
70
Mean Accuracy w/ response time (M=4) w/o response time (M=4) w/ response time (M=8) w/o response time (M=8) w/ response time (M=16) w/o response time (M=16)
65 60 55 ID OOD (M=4) (M=4)
ID OOD (M=8) (M=8)
ID OOD (M=16) (M=16)
Figure 1: Inference accuracy (mean ± std) across different M Table 2 and Figure 1 show performance on the food-risk dataset when training with GPT-2, where a small subset of participants is grouped and treated as a new human type, and the remaining participants are used for training. We vary the inference length M , corresponding to the number of in-context preference demonstrations available at test time. Similar to the results of synthetic data, incorporating response time generally narrows the performance gap between in-distribution and out-of-distribution inferences, indicating effective in-context adaptation to unseen human preferences. This phenomenon becomes more significant and clear as inference length M increases, which is consistent with our theory, as longer demonstration sequences improve in-context learning ability of the trained transformer.
6
Limitations
We briefly discuss the limitations of this work. For theoretical tractability, our analysis focuses on a linear-attention transformer, which is a simplified abstraction of the architectures commonly used in practice. Although our experiments with GPT-2 demonstrate that the observed phenomena extend to standard transformer models, extending the theoretical guarantees to more complex architectures remains an important problem for future work. In addition, while response time serves as an effective auxiliary signal for resolving the failure of in-context adaptation under binary preference labels, it may be difficult to reliably measure or collect in practical settings. Exploring alternative auxiliary signals that are easier to obtain yet similarly informative for preference modeling is therefore a promising direction.
7
Conclusion
In this paper, We study in-context reward adaptation for preference modeling in RLHF and showed that in-context learning based only on binary preference labels is fundamentally insufficient for adapting to unseen human preferences. This limitation stems from an intrinsic information bottleneck, where binary labels collapse reward parameters into nonlinear representations that cannot be universally decoded by linear attentions. We address this issue by incorporating human response time as an auxiliary signal. Our theoretical analysis shows that response time recovers a linear signal of preference strength and enables correct in-context adaptation to unseen reward models. Experiments on both synthetic and real-world datasets support our theory, demonstrating improved robustness to distribution shift and more effective use of in-context demonstrations.
11
References K. Ahn, X. Cheng, H. Daneshmand, and S. Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36: 45614–45650, 2023. Y. Bai, F. Chen, H. Wang, C. Xiong, and S. Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36: 57125–57211, 2023. R. Berlinghieri, I. Krajbich, F. Maccheroni, M. Marinacci, and M. Pirazzini. Measuring utility with diffusion models. Science Advances, 9(34):eadf1665, 2023. R. A. Bradley and M. E. Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):324–345, 1952. S. Casper, X. Davies, C. Shi, T. K. Gilbert, J. Scheurer, J. Rando, R. Freedman, T. Korbak, D. Lindner, P. Freire, et al. Open problems and fundamental limitations of reinforcement learning from human feedback. arXiv preprint arXiv:2307.15217, 2023. S. Chakraborty, J. Qiu, H. Yuan, A. Koppel, F. Huang, D. Manocha, A. S. Bedi, and M. Wang. Maxmin-rlhf: Alignment with diverse human preferences. arXiv preprint arXiv:2402.08925, 2024. I. Dasgupta, A. K. Lampinen, S. C. Chan, H. R. Sheahan, A. Creswell, D. Kumaran, J. L. McClelland, and F. Hill. Language models show human-like content effects on reasoning tasks. arXiv preprint arXiv:2207.07051, 2022. D. Fu, T.-Q. Chen, R. Jia, and V. Sharan. Transformers learn to achieve second-order convergence rates for in-context linear regression. Advances in Neural Information Processing Systems, 37:98675– 98716, 2024. S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in neural information processing systems, 35:30583–30598, 2022. A. Giannou, L. Yang, T. Wang, D. Papailiopoulos, and J. D. Lee. How well can transformers emulate in-context newton’s method? arXiv preprint arXiv:2403.03183, 2024. Y. Huang, Y. Cheng, and Y. Liang. arXiv:2310.05249, 2023.
In-context convergence of transformers.
arXiv preprint
J. Jang, S. Kim, B. Y. Lin, Y. Wang, J. Hessel, L. Zettlemoyer, H. Hajishirzi, Y. Choi, and P. Ammanabrolu. Personalized soups: Personalized large language model alignment via post-hoc parameter merging. arXiv preprint arXiv:2310.11564, 2023. T. Korbak, K. Shi, A. Chen, R. V. Bhalerao, C. Buckley, J. Phang, S. R. Bowman, and E. Perez. Pretraining language models with human preferences. In International Conference on Machine Learning, pages 17506–17533. PMLR, 2023. N. Lambert and R. Calandra. The alignment ceiling: Objective mismatch in reinforcement learning from human feedback. arXiv preprint arXiv:2311.00168, 2023. S. Li, Y. Zhang, Z. Ren, C. Liang, N. Li, and J. A. Shah. Enhancing preference-based linear bandits via human response time. Advances in Neural Information Processing Systems, 37:16852–16893, 2024. S. Min, X. Lyu, A. Holtzman, M. Artetxe, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837, 2022. L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022.
12
A. Ovadya. ’generative ci’through collective response systems. arXiv preprint arXiv:2302.00672, 2023. J. Palmer, A. C. Huk, and M. N. Shadlen. The effect of stimulus strength on the speed and accuracy of a perceptual decision. Journal of vision, 5(5):1–1, 2005. C. Park, M. Liu, D. Kong, K. Zhang, and A. Ozdaglar. Rlhf from heterogeneous feedback via personalization and preference aggregation. arXiv preprint arXiv:2405.00254, 2024. A. Rame, G. Couairon, C. Dancette, J.-B. Gaya, M. Shukor, L. Soulier, and M. Cord. Rewarded soups: towards pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards. Advances in Neural Information Processing Systems, 36:71095–71134, 2023. R. Ratcliff and G. McKoon. The diffusion decision model: theory and data for two-choice decision tasks. Neural computation, 20(4):873–922, 2008. M. J. Ryan, O. Shaikh, A. Bhagirath, D. Frees, W. B. Held, and D. Yang. Synthesizeme! inducing persona-guided prompts for personalized reward models in llms. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8045– 8078, 2025. O. Shaikh, M. S. Lam, J. Hejna, Y. Shao, H. Cho, M. S. Bernstein, and D. Yang. Aligning language models with demonstrated feedback. arXiv preprint arXiv:2406.00888, 2024. W. Shen, R. Zhou, J. Yang, and C. Shen. On the training convergence of transformers for in-context classification of gaussian mixtures. arXiv preprint arXiv:2410.11778, 2024. A. Singh, S. Hsu, K. Hsu, E. Mitchell, S. Ermon, T. Hashimoto, A. Sharma, and C. Finn. Fspo: Few-shot preference optimization of synthetic preference data in llms elicits effective personalization to real users. arXiv preprint arXiv:2502.19312, 2025. S. M. Smith and I. Krajbich. Attention and choice across domains. Journal of Experimental Psychology: General, 147(12):1810, 2018. T. Sorensen, J. Moore, J. Fisher, M. Gordon, N. Mireshghallah, C. M. Rytting, A. Ye, L. Jiang, X. Lu, N. Dziri, et al. A roadmap to pluralistic alignment. arXiv preprint arXiv:2402.05070, 2024. E.-J. Wagenmakers, H. L. Van Der Maas, and R. P. Grasman. An ez-diffusion model for response time and accuracy. Psychonomic bulletin & review, 14(1):3–22, 2007. Y. Wang, W. Zhong, L. Li, F. Mi, X. Zeng, W. Huang, L. Shang, X. Jiang, and Q. Liu. Aligning large language models with human: A survey. arXiv preprint arXiv:2307.12966, 2023. J. Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682, 2022. J. Wu, D. Zou, Z. Chen, V. Braverman, Q. Gu, and P. L. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.08391, 2023. J. Zhang. Guided profile generation improves personalization with llms. arXiv:2409.13093, 2024.
arXiv preprint
R. Zhang, S. Frei, and P. L. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1–55, 2024. S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
13
A
Proof of Theorem 1
Proof. To show Parts (i) and (ii), define µ̂i = N1
PN
1 T ∗ ∗ l l l=1 zi φ̃i and µi = E[zi φ̃ | θi ] = Eφ̃∼Pφ̃ [tanh( 2 φ̃ θi )φ̃]. 2
Let ŝi (U ) = µ̂Ti U φ̃qi and si (U ) = µTi U φ̃qi . Denote w = vec(U ) ∈ Rd by vectorizing U . Then a simple calculation gives ŝi (U ) = tr(U φ̃qi µ̂Ti ) = tr(U T µ̂i (φ̃qi )T ) = wT (φ̃qi ⊗ µ̂i ).
Denoting lN (w) = LN (U ) and letting v = φ̃qi ⊗ µ̂i and y = (1 + ziq )/2, calculating its gradient and Hessian as follows ∇lN (w) = E[(σ(wT v) − y)v] ∇2 lN (w) = E[σ(wT v)(1 − σ(wT v))vv T ].
Since kU kF ≤ R, kφ̃k ≤ B, with φ̃ ∼ Pφ̃ almost surely by assumptions, we obtain |wT x| ≤ kµ̂i kkU kF kφ̃q k ≤ RB 2 < ∞ and hence almost surely there exists some strictly positive constant λ > 0 such that ∇2 l(w) ≥ λE[vv T ] ≻ 0
by noticing E[vv T ] = E[φ̃qi (φ̃qi )T ] ⊗ E[µ̂i µ̂Ti ] = E[φ̃qi (φ̃qi )T ] ⊗ E[µi µTi ] ≻ 0 where the first equality follows the independence of φ̃qi and µ̂i . This concludes the strong convexity of LN (U ). Similarly, by replacing µ̂i by µi , following the same procedure concludes the strong convexity of L̄(U ). To show Part (iii), conditioned on θi∗ , φ̃qi and ziq , defining gU (u) = (1 + ziq ) log σ(uT U φ̃qi ) + (1 − q zi ) log(1 − σ(uT U φ̃qi )), then it is obvious that LN (U ) = E[gU (µ̂i )] and L̄(U ) = E[gU (µi )]. For any fixed U , denoting δ = µ̂i − µi , we have 1 gU (µ̂i ) − g(µi ) = ∇gU (µi )δ + δ T ∇2 gU (µi + tδ)δ 2
for some t ∈ (0, 1). Taking expectations on both sides yields 1 LN (U ) − L̄(U ) = E δ T ∇2 g(µi + tδ)δ 2 where we use the fact E[δ | i] = 0. Note that
∇2u gU (u) = σ(uT U φ̃qi )(1 − σ(uT U φ̃qi ))U φ̃qi (U φ̃qi )T
Under σ(x)(1 − σ(x)) ≤ 1/4, ∀x and boundedness kφ̃k ≤ B, kU kF ≤ R, k∇2 g(·)k ≤
1 1 kU φ̃qi k2 ≤ R2 B 2 , 4 4
which indicates 1 2 2 R B Ekδk2 8 1 2 2 R B (2B)2 = 8N R2 B 4 . = 2N
|LN (U ) − L̄(U )| ≤
To show Part (iv), define hU (u) = (σ(uT U φ̃qi ) − (1 + ziq )/2)u. Then, ∇LN (U ) = E[hU (µ̂i )(φ̃qi )T ], ˙ ≤ 1/4, σ ′′ (·) ≤ 1/6, we have ∇L̄(U ) = E[hU (µi )(φ̃qi )T ]. After careful calculations and noting |σ ′ ()| k∇2u hU (u)kop ≤ c1 kU φ̃qi k + c2 kU φ̃qi k2 kuk for some positive constants c1 , c2 > 03 . Since E[∇hU (µi )δ | i, φ̃qi , ziq ] = ∇hU (µi )E[δ] = 0, we have k∇LN (U ) − ∇L̄(U )kF = kE[E[hU (µ̂i ) − hU (µi ) | i, φ̃qi , ziq ]φ̃qi ]k 1 ≤ E[k∇2 hU (µi + tδ)kop kδk2 kφ̃qi k]. 2 3 Note that ∇2 h (·) is a tensor lying in Rd×d×d . u U
14
By the almost sure boundedness of φ̃, kµi + tδk ≤ kµ̂i k + 2kµi k ≤ 3B almost surely, hence indicating k∇2 hU (µi + tδ)kop ≤ c1 RB + 3c2 R2 B 3 almost surely. Therefore, 1 (c1 RB 2 + 3c2 R2 B 4 )Ekδk2 2 2(c1 + 3c2 RB 2 )B 4 . ≤ N
k∇LN (U ) − ∇L̄(U )kF ≤
According to strong convexity of L̄(·), there is some positive constant α such that 1 ∗ k∇L̄(UN ) − ∇L̄(Ū ∗ )kF α ∗ ∗ = k∇L̄(UN ) − ∇LN (UN )kF ≤ O(1/N )
∗ kUN − Ū ∗ kF ≤
∗ where the second equality follows ∇L̄(Ū ∗ ) = ∇LN (UN ) = 0. This concludes the proof.
B
Proof of Theorem 2
Proof. We consider the simplest case where d = 1 and let Pφ̃ be the following: P (φ̃ = +1) = P (φ̃ = −1) = 1/2. Then, 1 1 1 1 1 1 µnew (θ) = Eφ̃ [tanh( φ̃T θ)φ̃] = · tanh( θ) + · tanh(− θ) · (−1) = tanh( θ). 2 2 2 2 2 2 Considering the asymptotic case, i.e., M = ∞, N = ∞, we have M X 1 ∗ q q z j (φ̃j )T UN φ̃ Pnew (ẑnew = 1 | φ̃) = lim σ M→∞,N →∞ M j=1 new M X 1 z j (φ̃j )T Ū ∗ φ̃q = lim σ M→∞ M j=1 new = σ µnew (θnew )T Ū ∗ φ̃q 1 = σ tanh( θnew )Ū ∗ φ̃q . 2 Suppose a correct prediction is achieved asymptotically, i.e., we need to find some Ū ∗ ∈ R such that 1 tanh( θnew )Ū ∗ = θnew 2 which implies that Ū ∗ =
θnew . tanh(0.5θnew )
Moreover, note that Ū ∗ satisfies ∇L̄(Ū ∗ ) = 0, which is equivalent to 0 = Ei∼H [(σ(µi Ū ∗ ) − σ(θi∗ ))µi ] 1 = Ei [(tanh(0.5µi Ū ∗ ) − tanh(0.5θi∗ ))µi ] 2 1 = Ei [(tanh(0.25 tanh(0.5θi∗ )Ū ∗ ) − tanh(0.5θi∗ )) tanh(0.5θi∗ )] 4 where we use µi = 21 tanh( 12 θi∗ ) and σ(x) = 21 tanh( 12 x) + 21 . Let m = tanh(0.5θi∗ ) ∈ (−1, 1). Due to monotonicity and invertibility of tanh(·), there exists some distribution θi∗ ∼ H whose support is R 15
such that m ∼ Unif[−1, 1] where Unif[·, ·] denotes the uniform distribution. Thus, it is equivalent that Ū ∗ satisfies 0 = Em [m(tanh(0.25mŪ ∗ ) − m)] 1 = Em [m(tanh(0.25mŪ ∗ )] − 3 Z 1 1 1 ∗ = m(tanh(0.25mŪ )dm − 2 −1 3 Define
1 I(u) = 2
Z 1
m(tanh(0.25mŪ ∗)dm.
−1
It is obvious that I ′ (u) > 0, ∀u which indicates I(u) is monotonically increasing. Noticing that I(0) < θnew in order for asymptotically 1/3 I(6) > 1/3, it implies Ū ∗ ∈ (0, 6). Further, as Ū ∗ = tanh(0.5θ new ) ∗ correct prediction, considering ∀θnew > 6, we have Ū > 6, which contradicts to the previous claim Ū ∗ ∈ (0, 6). Therefore, we conclude the proof as we find distributions Pφ̃ , H and θnew such that q q P(ẑnew = · | φ̃q ) 6= P(znew = · | φ̃q ).
C
Proof of Theorem 3
Proof. Define ŝi = N1
PN
l l l ∗ l l l l l=1 (zi /ti )φ̃i , si = 2Σφ̃ θi and µzil = E[zi | φ̃i ], µtli = E[ti | φ̃i ]. Then, N
1 X ŝi − si = N l=1
µz l zil − i l µtli ti
!
N
=
N
φ̃li +
1 X N l=1
µzil µtli
φ̃li − 2Σφ̃ θi∗
!
l N zi 1 X l T ∗ l ∗ l l T ∗ φ̃ − 2Σ θ 2( φ̃ ) θ φ̃ + − 2( φ̃ ) θ i i i i i i φ̃ i . N tli l=1 l=1
1 X N
Taking expectation of its norm square yields # " 2 N i 2 X 2 h zil 2 l 2 l T ∗ ∗ l T ∗ l Ekŝi − si k ≤ + k φ̃ k − 2( φ̃ ) θ φ̃ k | θ ) E Var(k2( φ̃ ) θ E i i i i i i i N N tli l=1 ≤ O(1/K) + O(1/N ).
where we use Lemma 1 in the last inequality. It is straightforward that i h vec(∇LN,K (U )) = E (ŝTi U φ̃qi − ziq /tqi )φ̃qi ⊗ ŝi i h vec(∇L∞ (U )) = E (sTi U φ̃qi − 2(φ̃qi )T θi∗ )φ̃qi ⊗ si
and hence
i h ∇LN,K (U ) − ∇L∞ (U ) = E ((ŝi − si )T U φ̃qi − (ziq /tqi − 2(φ̃qi )T θi∗ ))(ŝi − si )(φ̃qi )T | {z } e1
i h + E (sTi U φ̃qi − 2(φ̃qi )T θi∗ )(ŝi − si )(φ̃qi )T {z } | e2
i h + E (((ŝi − si )T U φ̃qi ) − (ziq /tqi − 2(φ̃qi )T θi∗ ))si (φ̃qi )T . {z } | e3
We first bound ke1 kF :
h i 2 ke1 kF ≤ RB 2 E kŝi − si k + BE kziq /tqi − 2(φ̃qi )T θi∗ kkŝi − si k q p 2 ≤ RB 2 E kŝi − si k + B Ekziq /tqi − 2(φ̃qi )T θi∗ k2 · Ekŝi − si k2 √ ≤ O(1/N + 1/K + 1/ N K) 16
where we use Lemma 1 in the last inequality. For e2 we have q p ke2 kF ≤ B Ek(θi∗ )T (Σφ̃ U − I)φ̃qi k2 · Ekŝi − si k2 q p = B E[(φ̃qi )T (Σφ̃ U − I)T θi∗ (θi∗ )T (Σφ̃ U − I)φ̃qi ] · Ekŝi − si k2 √ √ ≤ B 2 (Rtr(Σφ̃ ) + d)tr(E[θi∗ (θi∗ )T ]) · O(1/ N + 1/ K) √ √ = O(1/ N + 1/ K).
√ √ √ N + 1/ K). Thus, k∇L (U ) − ∇L (U )k ≤ O(1/ N+ Similarly, we can show ke k ≤ O(1/ N,K ∞ F 3 F √ 1/ K). Next, we show that L∞ is strongly-convex. Note that ∇2vec(U) L∞ (U ) = E[φ̃qi (φ̃qi )T ⊗ si sTi ]
= 4E[φ̃qi (φ̃qi )T ⊗ Σφ̃ θi∗ (θi∗ )T Σφ̃ ]
= 4E[φ̃qi (φ̃qi )T ] ⊗ [Σφ̃ E(θi∗ (θi∗ )T )Σφ̃ ].
Since both E[φ̃qi (φ̃qi )T ] and E[θi∗ (θi∗ )T ] are full-rank, ∇2vec(U) L∞ (U ) ≻ 0, implying L∞ is ᾱ-stronglyconvex for some α > 0. Similarly, one can show LN,K is strongly-convex. As U ∗ = Σφ̃−1 is a minimizer of L∞ (which is obviously obtained by ∇L∞ (Σφ̃−1 ) = 0), strong convexity yields U ∗ is unique. By strong convexity of L∞ , 1 k∇L∞ (UN,K ) − ∇L∞ (U ∗ )kF ᾱ 1 ∗ ∗ ) − ∇LN,K (UN,K )kF = k∇L∞ (UN,K ᾱ √ √ ≤ O(1/ N + 1/ K),
∗ kUN,K − U ∗ kF ≤
∗ where the second equality follows UN,K and U ∗ are minimizers of LN,K and L∞ by definition. This completes the proof.
D
Proof of Corollary 1
1 Proof. Let ŝnew = M
l znew l l=1 tlnew φ̃ , snew = 2Σφ̃ θnew . Similar to the proof of Theorem 3, we can show
PM
Ekŝnew − snew k2 ≤ O(1/M + 1/K).
Then, T ∗ E(ôqnew − 2θnew φ̃q )2 = E(ŝTnew UN,K φ̃q − sTnew U ∗ φ̃q )2
≤ 3E[(ŝnew − snew )T (UN,K − U ∗ )φ̃q ]2 + 3E[sT (UN,K − U ∗ )φ̃q ]2 + 3E[(ŝnew − snew )T U ∗ φ̃q ]2
√ ≤ O(1/M + 1/N + 1/K + 1/ N K) which completes the proof.
E
A Useful Lemma
1 N Lemma 1. For i.i.d. {Xi }N i=1 and i.i.d. {Yi }i=1 , define X̄ = N µX = E[X1 ], µY = E[Y1 ]. Then 2 µX X̄ = O(N −1 ). E − µY Ȳ
17
PN
i=1 Xi ,
Ȳ = N1
PN
i=1 Yi
and
Proof. Let δX = X̄ − µX , δY = Ȳ − µY . −1 µX + δ X X̄ δY = 1+ µY µY Ȳ δY δY2 µX + δ X 1− + 2 + higher-order terms = µY µY µY µX δX µX δ Y = + − + higher-order terms. µY µY µ2Y − µµX . We have Let ∆ = X̄ Ȳ Y ∆2 = =
δX µX δ Y − µY µ2Y
2
2 µ2 δ 2 µX δ X δ Y δX + X4 Y − + higher-order terms. 2 µY µY µ3Y
2 Since E[δX ] = Var(X) = O(1/N ), E[δY2 ] = Var(Y ) = O(1/N ), E[δX δY ] = Cov(X, Y ) = O(1/N ), we conclude that E[∆2 ] = O(1/N ).
18