Conceptio › Archive › arXiv CS

arXiv CSopen access

Augmented Hypothesis Testing with Persona-Based LLM Simulations

· arxiv_cs

arXiv CS · Papers · License: Open Access

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

Augmented Hypothesis Testing with Persona-Based LLM Simulations

arXiv:2609.24629v1 [cs.LG] 21 Sep 2026

Ziyad Benomar Amazon Luxembourg, Luxembourg [email protected]

Aymen Al Marjani Amazon Luxembourg, Luxembourg [email protected]

Paul Missault Amazon Luxembourg, Luxembourg [email protected]

Saab Mansour Amazon Barcelona, Spain [email protected]

Abstract A/B testing requires large sample sizes, long timelines, and significant costs. When auxiliary predictions of experimental outcomes are available from machine learning models, uncertain prediction quality precludes replacing human experiments entirely, yet these predictions may still contain useful signal. We propose a principled framework for learning-augmented hypothesis testing that leverages predictions of unknown quality to reduce sample sizes while maintaining statistical validity. Predictions naturally vary in granularity, from coarse aggregate signals to fine-grained individual-level estimates, and our framework addresses both ends of this spectrum: (1) for population-level directional predictions, where only a binary signal on the treatment effect sign is available, we use an asymmetric test and prove consistency and robustness bounds within the learning-augmented algorithms paradigm; (2) for individual-level predictions, we introduce Generalized PPI++ (GPPI), extending Prediction-Powered Inference to handle nonlinear prediction errors through higher-dimensional transformations. Both methods benefit from accurate predictions while remaining robust to inaccurate or adversarial ones. We validate our framework using persona-based LLM simulations, where AI agents equipped with user personas predict individual behavior, as a natural prediction source spanning both granularity levels. Experiments on four real-world datasets demonstrate that our methods, combined with persona-based predictions, substantially reduce experimental costs while preserving rigorous statistical validity.

Introduction

Hypothesis testing through A/B experiments remains the gold standard for behavioral data decision-making, with organizations conducting thousands of tests annually to optimize products and services (Kohavi et al., 2009; Xu et al., 2015). However, traditional A/B tests impose substantial costs: they require exposing different versions to distinct user groups, collecting behavioral data over extended periods, and hundreds of thousands or millions of interactions. These requirements translate directly into lost opportunity costs, delayed decisions, and significant infrastructure expenses (Deng & Shi, 2017). Recent advances in LLMs (Brown et al., 2020; Achiam et al., 2023; Touvron et al., 2023) have enabled a promising alternative: persona-based simulation, where AI agents equipped with detailed user personas predict experimental outcomes (Park et al., 2023; Horton et al., 2023; Argyle et al., 2023; Mansour et al., 2025). A persona encodes a user’s historical behavior, preferences, and characteristics in textual form. By providing this persona to an LLM and asking it to predict how the user would respond to new stimuli, researchers can 1

simulate behavior across a population of personas, potentially predicting A/B test outcomes before conducting costly human experiments (Aher et al., 2023; Manning et al., 2024). This approach has motivated growing research in agentic simulations for behavior modeling, with early results suggesting that LLM-based agents can capture meaningful patterns in human decision-making (Grossmann et al., 2023; Dillion et al., 2023). However, completely replacing human A/B tests with persona-based predictions faces fundamental challenges. The black-box nature of LLMs, sensitivity to prompt engineering (Zhao et al., 2021; Lu et al., 2022), distribution shifts, and the inherent complexity of human behavior (Skinner, 1965) make formal quality guarantees difficult. Moreover, persona-based simulations introduce a critical structural consideration: while a collection of personas may approximate the overall population distribution, establishing reliable one-to-one mappings between specific personas and specific users is often infeasible. This asymmetry means that persona simulations naturally produce different types of predictions depending on the simulation design—from coarse population-level directional signals (e.g., observing the sign of the treatment effect on the persona population) to fine-grained individual-level predictions when persona-user alignment is available. Rather than treating LLM predictions as perfect substitutes for human experiments, more principled approaches propose leveraging them to reduce the required human sample size while maintaining statistical validity (Angelopoulos et al., 2023a; Ji et al., 2026; Calderon et al., 2025). These approaches acknowledge prediction uncertainty while extracting value when predictions contain useful signal. The fundamental question becomes: given predictions of unknown quality from persona-based LLM simulations (or other sources), how can we design hypothesis testing procedures that provably benefit from accurate predictions while remaining robust to inaccurate ones? Learning-augmented algorithms. Our work contributes to the literature on learningaugmented algorithms (Lykouris & Vassilvitskii, 2018; Mitzenmacher & Vassilvitskii, 2022), which addresses how to leverage machine learning predictions lacking formal guarantees. The key criteria for designing such algorithms are (i) improved performance when predictions are accurate (consistency), (ii) guarantees comparable to prediction-free methods when predictions fail (robustness), (iii) and graceful degradation as prediction error increases (smoothness) (Purohit et al., 2018). This framework has been initially studied in the context of online algorithms such as caching (Lykouris & Vassilvitskii, 2018; Rohatgi, 2020), scheduling (Lattanzi et al., 2020; Mitzenmacher, 2020; Benomar & Perchet, 2024; Benomar et al., 2026), ski-rental (Purohit et al., 2018; Bamas et al., 2020; Benomar & Perchet, 2023) and online selection (Dütting et al., 2021; Antoniadis et al., 2020), it extended to designing data structures (Lin et al., 2022; McCauley et al., 2023; Zeynali et al.; Benomar & Coester, 2024), and improving machine learning algorithms and statistical methods (Bhaskara et al., 2021; Golowich & Moitra, 2022; Raman & Tewari, 2024). The typical objective is establishing a good tradeoff between consistency and robustness (Benomar & Perchet, 2025). Prediction-Powered Inference (PPI). PPI++ (Angelopoulos et al., 2023a;b) extends the learning-augmented setting to statistical inference, addressing hypothesis testing in semisupervised settings with n labeled samples, N unlabeled samples, and a predictor of unknown quality. PPI estimates systematic predictor bias using labeled data, then applies bias correction to predictions on unlabeled data. This provably improves inference quality for sufficiently large n, reducing confidence intervals and increasing power, with improvement magnitude depending on prediction-label correlation. When predictions are perfect, PPI achieves the performance of N additional labeled samples; when uninformative, it gracefully degrades to using only n labeled samples. PPI++ can be more generally extended to semi-supervised settings beyond inference, for example in online optimization (Shoham et al., 2025), multi-armed bandits (Ji et al., 2025), ranking (Chatzi et al., 2024), and inference with multiple predictors under budget constraints (Cowen-Breen et al., 2026) Different prediction types. Predicting AB test outcomes can be framed as regression or classification tasks. For instance, Li et al. (2015) develop a regression model to predict expected reward values and classify experiments as WIN, TIE, or LOSS when comparing 2

treatment to control variants, which gives a population-level prediction. Some works also explored individual-level behavior prediction from historical data (Sales et al., 2018) using classical ML methods. These different types of predictions can be used in the learningaugmented framework to improve AB testing. A more recent approach to obtain either population- or individual-level predictions is through persona-based LLM simulation (Hu & Collier, 2024; Chandrasegaran et al., 2025; Li et al., 2025). When persona-user alignment is unavailable (for e.g. synthetic personas), these simulations provide population-level predictions (Mansour et al., 2025; Castelo et al., 2026; Rieder et al., 2026). When persona-user mapping is possible (e.g., through demographic or behavioral matching), simulations may yield individual-level predictions, though with uncertain quality. This distinction is fundamental: the simulation architecture determines what type of prediction is available, which in turn dictates what statistical methods are appropriate. 1.1

Contributions and Organization

Theoretical contributions. (1) Population-level directional predictions: We consider hypothesis testing with binary prediction on the sign of the effect, i.e. the prediction only signals which treatment version is better. In this setting, we analyze hypothesis testing within the learning-augmented algorithms framework, characterizing the fundamental trade-off between leveraging predictions when accurate and maintaining validity when incorrect. For that, we used an asymmetric Z-test and prove consistency and robustness bounds. (2) Individual-level predictions: When granular predictions of individual engagement with both treatments are available, Prediction-Powered Inference (PPI/PPI++) (Angelopoulos et al., 2023b) can leverage these predictions to enhance hypothesis testing. We propose Generalized PPI (GPPI), extending to handle nonlinear prediction errors through higher dimension transformations. We also prove that, given a set of candidate transformations, a greedy selection approach can achieve the same performance as the best candidate, then we show how the performance scales with the dimension of such transformations. Persona-based simulation validation. A key contribution is applying PPI++ and our theoretical framework to persona-based LLM simulations. We define personas using users’ action histories and employ LLMs with these personas to predict individual ratings on new items. This can be translated into preference simulation for AB testing, where the objective is determining which item has statistically higher average ratings. Given a pair of items, we compute rating differences across personas. This validates our framework with realistic AI predictions exhibiting complex error patterns. We conduct experiments on four real-world datasets: MovieLens-25M (Harper & Konstan, 2015), Jester (Goldberg et al., 2001), Last.fm (Bertin-Mahieux et al., 2011), and Book-Crossing (Ziegler et al., 2005). We use three Mistral models to make predictions: Mistral-7B-Instruct-v0.2 (7B parameters), Mistral-Small-24B-Instruct-2501 (24B parameters), and Mistral-Large-Instruct-2407 (123B parameters), to which we refer respectively as mistral-7b, mistral-24b, and mistral-123b in the following. These models were chosen because they are open-source, which ensures reproducibility of the experiments, and because they have different parameter sizes, which demonstrates that our methods work across different model scales This comprehensive evaluation demonstrates that persona-based simulations, when combined with our statistical framework (either on individual or population level), can reduce experimental costs while maintaining rigorous validity guarantees.

Hypothesis testing with directional predictions

We investigate how population-level directional predictions can improve hypothesis testing in A/B testing scenarios. In many practical settings, the key decision is simply determining whether a treatment effect is positive or negative, for example whether a new feature increases or decreases user engagement. We consider the setting where an analyst has access to a binary prediction indicating the sign of the treatment effect. This represents the minimal type of prediction assistance, i.e. a single bit of directional information. 3

Problem setup. Consider i.i.d. observations X1 , . . . , Xn ∼ N (µ, σ2 ) with known σ, and let X̄ = n1 ∑in=1 Xi . We test H0 : µ = µ0 versus H1 : µ ̸= µ0 using the Z-test statistic √ Z = ( X̄ − µ0 )/(σ/ n). Under H0 , we have √ Z ∼ N (0, 1). Under H1 with µ = µ0 + δ, the statistic follows Z ∼ N (θ, 1) where θ = δ n/σ is the non-centrality parameter. The classical two-sided test at significance level α rejects H0 when | Z | > z1−α/2 , where z1−α/2 = Φ−1 (1 − α/2). This symmetric critical region allocates equal probability α/2 to each tail, reflecting no prior directional information. While we assume exact normality for analytical convenience, the Central Limit Theorem ensures asymptotic validity even with non-normal data. Test quality is measured by power Power(θ ), the probability of correctly rejecting H0 when θ ̸= 0. Asymmetric tests with predictions. Suppose we have a directional prediction ŝ ∈ {−1, +1} for sign(δ), potentially from domain expertise, preliminary analysis, or of by aggregating persona-based simulations. We perform an asymmetric test parameterized by λ ∈ [0, 1) controlling the level of asymmetry. Without loss of generality, assume ŝ = +1. We define asymmetric tail probabilities: α L = 1−2 λ α, and α R = 1+2 λ α, where α L + α R = α maintains the desired significance level. The critical values are c L (λ) = z(1−λ)α/2 and c R (λ) = z1−(1+λ)α/2 , where z p = Φ−1 ( p). We reject H0 if Z < c L (λ) or Z > c R (λ). By allocating more Type I error to the predicted tail (increasing α R when ŝ = +1), we lower the threshold c R (λ), facilitating rejection when the true effect aligns with the prediction. Conversely, we raise |c L (λ)|, making rejection harder when the effect contradicts the prediction. The parameter λ controls this bias strength, with λ = 0 recovering the symmetric classical test and λ → 1 concentrating nearly all error probability in the predicted direction, which corresponds to the one-sided Z-test for that direction. This embodies a fundamental consistency-robustness trade-off. Learning-augmented framework and performance guarantees. Our procedure falls into the learning-augmented algorithms paradigm: leveraging auxiliary predictions to improve performance when accurate while maintaining acceptable performance when predictions fail. We formalize this through approximation ratios comparing our test’s power to two natural benchmarks. When predictions are correct (sign(θ ) = ŝ), we measure consistency relative to the one-sided test (λ = 1): Power(θ, α, λ) Consistency ratio = inf . α,θ:sign(θ )=ŝ Power( θ, α, 1) When predictions are incorrect (sign(θ ) ̸= ŝ), we measure robustness relative to the prediction-agnostic two-sided baseline (λ = 0): Robustness ratio =

Power(θ, α, λ) . α,θ:sign(θ )̸=ŝ Power( θ, α, 0) inf

These ratios are worst-case guarantees. A consistency c means achieving at least c-fraction of the oracle’s power if the prediction is correct, while a robustness r means retaining at least r-fraction of baseline power regardless of the quality of the prediction. Our main theorem establishes that the asymmetric test has consistency and robustness guarantees that scale linearly with λ. Theorem 2.1 (Consistency-Robustness Trade-off). Assume without loss of generality that the predicted direction is ŝ = +1. Let λ ∈ [0, 1]. For the asymmetric test with parameter λ: Power(θ,α,λ)

(i) Consistency. For all α ∈ (0, 1) and θ > 0 (correct prediction): Power(θ,α,1) ≥ 1+2 λ Power(θ,α,λ)

(ii) Robustness. For all α ∈ (0, 1) and θ < 0 (incorrect prediction): Power(θ,α,0) ≥ 1 − λ Theorem 2.1 reveals a linear trade-off structure. Setting λ = 0 yields perfect robustness but only 50% consistency, as the symmetric test achieves half the power of the oracle onesided test in each direction. Conversely, setting λ → 1 yields perfect consistency but zero robustness. Intermediate values λ ∈ (0, 1) interpolate smoothly these extreme behaviors. 4

Hypothesis testing with individual-level predictions

Hypothesis testing for effect size estimation is directly related to mean estimation and confidence interval construction. In A/B testing, we typically estimate means for both treatments and construct confidence intervals to determine the effect size sign. This is why in this section we focus on mean estimation in the setting where individual-level predictions are available: rather than a single population-level prediction as in the previous section, we assume access to predictions for each individual. When agentic simulations or machine learning models generate individual-level outcome estimates, we can leverage substantially more information to improve statistical inference. This aligns with the Prediction-Powered Inference (PPI) framework (Angelopoulos et al., 2023a;b), which utilizes both labeled data (ground truth outcomes) and abundant unlabeled data (predictions only). Problem setup. We observe labeled data ( Xi , Yi ) for i ∈ [n], unlabeled data X̃ j for j ∈ [ N ] (typically N ≫ n), and have access to predictor f : X → R available for all observations. Our goal is to estimate the population mean θ ∗ = E[Y ] and construct a valid confidence interval leveraging the predictions. The classical estimator θ̂classical = n1 ∑in=1 Yi ignores unlabeled data entirely. PPI++ (Angelopoulos et al., 2023b) improves upon this by constructing N

λ 1 θbPPI = f ( X̃ j ) + ∑ (Yi − λ f ( Xi )) , N j∑ n =1 i =1

(1)

where λ ∈ R minimizes asymptotic variance. This estimator combines scaled predictions on unlabeled data with a correction term learned from labeled data that captures systematic deviation between Yi and λ f ( Xi ). The structure applies a multiplicative then additive correction to predictions, which is effective when Y ≈ λ∗ f ( X ) + b for constants λ∗ , b, up to noise. However, when the relationship is more complex, for example Y = G ( f ( X )) for a nonlinear function G, the linear correction is insufficient. 3.1

Generalized PPI++ (GPPI)

To overcome the limitation above, we extend PPI++ by transforming predictions through a transformation ϕ : R → Rd , for instance ϕ(z) = (z, z2 , . . . , zd ) T . Define sample averages Ȳn =

1 n Yi , n i∑ =1

The GPPI estimator is then

ϕ̄nL =

1 n ϕ( f ( Xi )), n i∑ =1

ϕ̄U N =

1 N ϕ( f ( X̃ j )). N j∑ =1

L θ̂GPPI = Ȳn + λ T (ϕ̄U N − ϕ̄n )

(2)

∈ Rd is a weight vector.

where λ Assuming n/N → r, the optimal weight minimiz−1 1 ing asymptotic variance is λ∗ = 1+ r Σϕ σY,ϕ , where Σϕ = Cov( ϕ ( f ( X ))) and σY,ϕ = Cov(Y, ϕ( f ( X ))). We estimate these using sample covariances: Σ̂ϕ = σ̂Y,ϕ =

N 1 U T (ϕ( f ( X̃ j )) − ϕ̄U ∑ N )( ϕ ( f ( X̃ j )) − ϕ̄ N ) N − 1 j =1

(3)

n 1 (Y − Ȳn )(ϕ( f ( Xi )) − ϕ̄nL ) ∑ n − 1 i =1 i

(4)

1 1 yielding the plug-in estimator λ̂ = 1+n/N Σ̂− ϕ σ̂Y,ϕ .

Our main result proves how to construct an asymptotically valid confidence interval using the predictions and their augmentation via transformation. Theorem 3.1. Assume that E[Y 2 ] < ∞, E[∥ϕ( f ( X ))∥2 ] < ∞, rn := n/N → r > 0, and Σϕ = Cov(ϕ( f ( X ))) is positive-definite. Then: −1 1 T (i) Let V ∗ = σY2 − 1+ r σY,ϕ Σϕ σY,ϕ , then we have

√

n(θ̂GPPI − θ ∗ ) − → N (0, V ∗ )

(ii) Define respectively the residual variance and the variance estimator σ̂R2 =

n 1 [(Y − Ȳn ) − λ̂ T (ϕ( f ( Xi )) − ϕ̄nL )]2 ∑ n − 1 i =1 i

and

V̂ = σ̂R2 + rn λ̂ T Σ̂ϕ λ̂ ,

Then V̂ − → V∗. (iii) CI1−α = [θ̂GPPI ± zα/2

V̂/n ] satisfies limn→∞ P(θ ∗ ∈ CI1−α ) = 1 − α.

The proof of Theorem 3.1 follows the general approach of Angelopoulos et al. (2023a), which establishes asymptotic normality for PPI estimators in the context of general Mestimators. Our result extends mean-estimation to d-dimensional transformations ϕ(·) and vector-valued weights λ ∈ Rd . 3.2

GPPI with greedy transformation Selection

While Theorem 3.1 establishes GPPI’s validity with a fixed transformation ϕ, the choice of this transformation critically affects performance. Although function classes like polynomials can approximate the optimal transformation under appropriate assumptions with sufficient dimension, this creates a fundamental trade-off: high-dimensional bases better approximate the relationship between f ( X ) and Y but yield noisy coefficient estimates with limited labeled data. The challenge is that decision-makers lack prior knowledge of the underlying relationship, hence cannot decide beforehand which transformation to consider. We propose a greedy selection algorithm from a candidate set of m transformations, relaxing the need to specify a single transformation a priori. Let ϕ1 , . . . , ϕm be candidate transforj mations where ϕ j : R → Rd with d j ≤ d for all j ∈ [m]. Each transformation induces an asymptotic variance Vj in the estimator (2). Since width of the confidence interval built by p GPPI scales with Vj , we seek j∗ = arg min j∈[m] Vj . Our greedy algorithm is (i) For each ϕ j , compute the plug-in variance estimator using Theorem 3.1(ii): 2 V̂j = σ̂R,j + rn λ̂ Tj Σ̂ϕ j λ̂ j

(5)

1 2 is the residual variance for transformation j. where λ̂ j = 1+1rn Σ̂− σ̂ j and σ̂R,j ϕ j Y,ϕ

(ii) Select ĵ = arg min j∈[m] V̂j . (iii) Construct the GPPI estimator using ϕ ĵ . The following theorem shows this procedure asymptotically achieves the optimal transformation variance. Theorem 3.2 (Greedy transformation Selection). Let ϕ1 , . . . , ϕm be candidate transformations with dimensions d j ≤ d, where d is a fixed constant. Assume for all j ∈ [m] that ∥ϕ j ( f ( X ))∥2 ≤ M and Y 2 ≤ KY almost surely, Σϕ j = Cov(ϕ j ( f ( X ))) is positive-definite, and there exists a unique minimizer j∗ = arg min j∈[m] Vj . Let ĵ be the selected transformation index. Then: p

(i) Consistent selection: ĵ − → j∗ as n → ∞. (ii) Asymptotic normality: The greedy GPPI estimator satisfies

√

ĵ

n(θ̂GPPI − θ ∗ ) − → N (0, Vj∗ ) .

(iii) Variance bound: For all δ ∈ (0, 1), we have with probability 1 − δ that r log(md/δ) 2 V̂ĵ ≤ min Vj + CKY M . n j∈[m] Note that Theorem 3.2 requires Y and ϕ( f ( X )) to be bounded almost surely, rather than finite second moments as in Theorem 3.1. This assumption enables high-probability bounds. 6

A natural question concerns coverage: does it deteriorate (in finite sample regime) as the number of candidate transformations increases? We prove in B.14 that the nominal coverage is also asymptotically preserved by greedy GPPI. Other aggregation strategies Theorem 3.2 (iii) demonstrates that the confidence interval q

width, proportional to

V̂ĵ , approaches the optimal asymptotic width among candidate

transformations. Alternatively, we could stack all transformations into an md-dimensional transformation and estimate the optimal linear combination from labeled data. While this may improve asymptotic performance, it creates a crucial trade-off: since M bounds ∥ϕ j ( f ( X ))∥2 , it scales with dimension d sometimes polynomially or even exponentially. Therefore, the greedy selection achieves a performance near the best transformation with minimal convergence cost, whereas stacking transformations may yield even better asymptotic variance but suffers from parameter estimation variance with limited labeled data. In practice, PPI++ typically works with black-box predictors f (e.g., LLMs, pre-trained models) that provide reasonable signals but require minor corrections for biases or behavioral changes due to distribution shifts, systematic over/underestimation, threshold adjustments, or capping effects. In these settings, corrections often require only simple transformations (linear, polynomial, threshold functions,...). Making a broad initial guess of possible transformations and applying greedy selection provides stable improvements without the variance penalty of high-dimensional parameter estimation from finite labeled data. If the model needs significant and complex modification, then further training might be a better approach. 3.3

Dimension scaling

Following the discussion above, we dive deeper into the impact of high dimensional transformations In practice, GPPI provides the greatest benefits when predictions exhibit systematic nonlinear errors. However, the transformation dimension must remain moderate to avoid large variance from estimating λ. The following result characterizes the permissible growth rate of dimension d with sample size n. Theorem 3.3. Let ϕd : R → Rd denote a d-dimensional transformation where d = dn may grow with n. Assume: (i) E[Y 2 ] < ∞ and limn→∞ n/N = r > 0 (ii) ∥ϕd ( f ( X ))∥2 ≤ Md almost surely (iii) The eigenvalues of Σd = Cov(ϕd ( f (p X ))) satisfy λmin (Σd ) ≥ cd > 0 (iv) The asymptotic √ variance Vd∗ converges to a limit V∞ If Md2 log d/c2d = o ( n), then

√

n(θ̂GPPI − θ ∗ ) − → N (0, V∞ ),

V̂ − → V∞ ,

and the confidence interval CI1−α achieves asymptotic coverage probability 1 − α.

Experiments

We validate our findings through experiments on population-level directional predictions (Section 2) and individual-level predictions (Section 3) using persona-based LLM simulations on four datasets: (i) MovieLens-25M (Harper & Konstan, 2015): 25M ratings from 162K users on 59K movies, (ii) Jester (Goldberg et al., 2001): 1.4M ratings from 26K users on 140 jokes, (iii) Last.fm (Bertin-Mahieux et al., 2011): 423K listening events from 8.6K users on 20K tracks, and (iv) Book-Crossing (Ziegler et al., 2005): 296K ratings from 7.3K users on 146K books. For each dataset, we select 15 test items based on recency and popularity (detailed in Appendix C)), with remaining items forming the training set for persona construction. Given a user’s historical interactions, we construct textual personas summarizing their preferences and provide these to LLMs for test item predictions. We use three models to make predictions: mistral-7b, mistral-24b, and mistral-123b. Comprehensive experimental details and additional results are provided in Appendix C. We compare GPPI with greedy transform selection against PPI++. For numerical stability, 1 we use ridge regularization in estimating the weights: λ̂ = 1+n/N (Σ̂ϕ + γI )−1 σ̂Y,ϕ with γ = 7

10−3 . Furthermore, to avoid biasing the greedy selection algorithm towards selecting high dimensional transformations, we implement two penalized selection variants. Denoting by d j the dimension of the transformation ϕ j : • Greedy AIC: selects the family minimizing Var(θ̂ GPPI (ϕ j )) + 2d j /n. This penalization proportional to d is inspired by the Akaike-type correction (Akaike, 1974) for model selection. • Greedy BIC: selects the family minimizing Var(θ̂ GPPI (ϕ j )) + d j log(n)/n, which is inspired by the Bayesian information criterion penalty (Schwarz, 1978). To correct nonlinear prediction biases without prior knowledge of the error structure, we use m = 10 complementary transformations ϕ j : R → Rd j with dimensions d ≤ 6, selected to span diverse function classes while maintaining small dimension: • Polynomials: standard and Bernstein polynomials (d ∈ {3, 5}) • Logarithmic: p 7→ (log(| p| + 1), . . . , log(| p| + 1)3 ), handle magnitude-dependent errors • Piecewise: cubic B-splines and tent functions (d = 4), provide corrections for local biases • Sigmoidal: logistic and softplus functions (d ∈ {3, 5}), capture saturation effects and bounded prediction ranges To all families, we add the identity transformation, ensuring PPI++ remains recoverable. As explained in Section 3, GPPI is designed to guarantee better mean estimation than PPI++, and this naturally extends to better guarantees also for hypothesis testing and A/B tests. We experimentally test it on both tasks. 4.1

Mean estimation with individual-level predictions

We measureimprovement using effective sample size (ESS) gain (Krsteski et al., 2025): Var(θ̂classical ) −1 Var(θ̂method )

× 100, which measures how much the sample size can be reduced with the method we’re testing (PPI, GPPI, greedy selection, ...) compared to the baseline estimator that uses only labeled data, while maintaining the same level of precision in the estimation. The higher this metric the better is the improvement. ESSgain% =

We conduct a first experiment with N = 5n over 5000 Monte Carlo repetitions, sampling n labeled users and N unlabeled users to compute estimators and 95% confidence intervals. The figure below show ESS gains for Jester datasets, with columns representing items and rows representing different models. The subplot corresponding to each item and model shows ESS gain across different sample sizes n and different estimators.

Jester Jokes 7B params ESS Gain

108

117 10.0%

5.0%

0.0%

24B params ESS Gain

10.0%

1000 2000

4.0%

250 500 750

0.0%

500 1000 1500

40.0%

10.0% 5.0%

0.0%

500 1000 1500

30.0%

1000 2000

0.0%

500 10001500

10.0%

10.0% 0.0%

0.0%

20.0%

250 500 750

1000 2000

0.0%

PPI++

2000

0.0%

36 7.5%

4.0%

5.0%

0.0%

143

6.0%

2.0%

10.0%

134

10.0%

6.0%

20.0%

123B params ESS Gain

127

1000

2000

2.5%

0.0%

20.0%

10.0%

0.0%

500 1000

5.0%

2.0% 0

1000

2000

10.0%

0.0%

500 1000

20.0%

10.0%

Greedy

Greedy AIC

1000

2000

0.0%

1000 2000

5.0% 0

1000 2000

30.0%

10.0%

0.0%

15.0%

20.0%

1000 2000

500

10.0%

30.0%

250

15.0%

30.0%

0.0%

5.0%

0.0%

10.0%

0.0%

500 1000 1500

30.0% 20.0% 10.0% 0

1000 2000

0.0%

1000 2000

Greedy BIC

Figure 1: ESS gain comparison on Jester dataset across different jokes and Mistral models 8

In Appendix C.5, we present results for the three remaining datasets and compare the coverage of PPI++ and GPPI across all four datasets. The experiments demonstrate that GPPI with greedy selection consistently outperforms PPI++ in ESS gain while maintaining comparable coverage. For the Book Crossing and Last.FM datasets, coverage falls below the nominal level for certain items; however, these failure modes are common to both PPI++ and GPPI. They occur for items whose individual values contain outliers or exhibit heavy-tailed distributions, which violate the finite-sample CLT approximation. 4.2

A/B testing with population- and individual-level predictions

We evaluate our framework’s effectiveness in hypothesis testing by comparing item pairs to determine which has higher average ratings. We compare individual- and population-level predictions against classical estimators using ESS gain as our quality measure, with N = 5n. Here, ESS gain is defined as the reduction of the sample size that we can achieve compared to a classical A/B test using only the labeled data, while maintaining the same power of the test. In this experiment, we set the type I error to α = 5%. For population-level predictions, we aggregate all available model predictions and compare mean predicted ratings. We test asymmetry values λ ∈ {0.25, 0.5, 0.75}. Figures 13 and 2 show that GPPI with greedy transform selection consistently outperforms PPI++, with both methods achieving significant improvements over the classical baselines which only uses labeled data, up to 30% sample size reduction. Population-level predictions yield substantial improvements when predictions are accurate (MovieLens), but can negatively impact performance with noisy predictions (Jester).

Mean ESS Gain (pairs with | | > )

7B parameters

Jester Jokes

Mean ESS Gain for pairs with | | >

20.0%

30.0%

0.0%

10.0%

0.0%

20.0%

10.0% 0.0

0.2

0.4

0.6

0.8

Minimum effect size

PPI++

Greedy

1.0

Greedy AIC

24B parameters

123B parameters

20.0% 10.0% 0.0% 10.0% 0.0

0.2

0.4

0.6

0.8

Minimum effect size Greedy BIC

1.0

Asym =0.25

0.0

0.2

0.4

0.6

0.8

Minimum effect size

Asym =0.5

1.0

Asym =0.75

Figure 2: Comparison of the learning-augmented A/B testing approaches for Jester dataset While the experiments above were not conducted in a live production A/B test environment, the methodology and conclusions directly translate to such settings, as A/B tests fundamentally rely on hypothesis testing to decide the best treatment.

Conclusion

We presented a framework for learning-augmented hypothesis testing that leverages predictions to reduce experimental costs while maintaining validity. For population-level directional predictions, we established a consistency-robustness trade-off through asymmetric tests with linear guarantees. For individual-level predictions, we introduced Generalized PPI (GPPI) with greedy transformation selection, extending Prediction-Powered Inference to handle nonlinear prediction errors. In particular, we applied our methods to personabased LLM simulations on four real-world datasets, demonstrating the usefulness of our framework and how it improves over prior methods.

References Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023. 9

Gati V Aher, Rosa I Arriaga, and Adam Tauman Kalai. Using large language models to simulate multiple humans and replicate human subject studies. International Conference on Machine Learning, pp. 337–371, 2023. Hirotugu Akaike. A new look at the statistical model identification. IEEE Transactions on Automatic Control, 19(6):716–723, 1974. Anastasios N Angelopoulos, Stephen Bates, Clara Fannjiang, Michael I Jordan, and Tijana Zrnic. Prediction-powered inference. Science, 382(6671):669–674, 2023a. Anastasios N Angelopoulos, John C Duchi, and Tijana Zrnic. Ppi++: Efficient predictionpowered inference. arXiv preprint arXiv:2311.01453, 2023b. Antonios Antoniadis, Themis Gouleakis, Pieter Kleer, and Pavel Kolev. Secretary and online matching problems with machine learned advice. In Advances in Neural Information Processing Systems, volume 33, pp. 7933–7944, 2020. Lisa P Argyle, Ethan C Busby, Nancy Fulda, Joshua R Gubler, Christopher Rytting, and David Wingate. Out of one, many: Using language models to simulate human samples. Political Analysis, 31(3):337–351, 2023. Etienne Bamas, Andreas Maggiori, and Ola Svensson. The primal-dual method for learning augmented algorithms. Advances in Neural Information Processing Systems, 33:20083–20094, 2020. Ziyad Benomar and Christian Coester. Learning-augmented priority queues. Advances in Neural Information Processing Systems, 37:124163–124197, 2024. Ziyad Benomar and Vianney Perchet. Advice querying under budget constraint for online algorithms. Advances in Neural Information Processing Systems, 36:75026–75047, 2023. Ziyad Benomar and Vianney Perchet. Non-clairvoyant scheduling with partial predictions. arXiv preprint arXiv:2405.01013, 2024. Ziyad Benomar and Vianney Perchet. On tradeoffs in learning-augmented algorithms. arXiv preprint arXiv:2501.12770, 2025. Ziyad Benomar, Romain Cosson, Alexander Lindermayr, and Jens Schlöter. Non-clairvoyant scheduling with progress bars. Advances in Neural Information Processing Systems, 38: 93520–93561, 2026. Thierry Bertin-Mahieux, Daniel P.W. Ellis, Brian Whitman, and Paul Lamere. The million song dataset. In Proceedings of the 12th International Conference on Music Information Retrieval (ISMIR 2011), 2011. Aditya Bhaskara, Ashok Cutkosky, Ravi Kumar, and Manish Purohit. Logarithmic regret from sublinear hints. Advances in Neural Information Processing Systems, 34:28222–28232, 2021. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. Nitay Calderon, Roi Reichart, and Rotem Dror. The alternative annotator test for llm-as-ajudge: How to statistically justify replacing human annotators with llms. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 16051–16081, 2025. Alberto Castelo, Zahra Zanjani Foumani, Ailin Fan, Keat Yang Koay, Vibhor Malik, Yuanzheng Zhu, Han Li, Meysam Feghhi, Ronie Uliana, Shuang Xie, et al. Simgym: Traffic-grounded browser agents for offline a/b testing in e-commerce. arXiv preprint arXiv:2602.01443, 2026. 10

Senthil Chandrasegaran, Peter Lloyd, et al. Synthetic users: insights from designers’ interactions with persona-based chatbots. AI EDAM, 39:e2, 2025. Ivi Chatzi, Eleni Straitouri, Suhas Thejaswi, and Manuel G Rodriguez. Prediction-powered ranking of large language models. Advances in Neural Information Processing Systems, 37: 113096–113133, 2024. Charlie Cowen-Breen, Alekh Agarwal, Stephen Bates, William W. Cohen, Jacob Eisenstein, Amir Globerson, and Adam Fisch. Multiple-prediction-powered inference. In The Fourteenth International Conference on Learning Representations, 2026. URL https: //openreview.net/forum?id=gJZ5rf2bS4. Alex Deng and Xiaolin Shi. Data-driven metric development for online controlled experiments: Seven lessons learned. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 77–86, 2017. Danica Dillion, Niket Tandon, Yuling Gu, and Kurt Gray. Can ai language models replace human participants? Trends in Cognitive Sciences, 27(7):597–600, 2023. Paul Dütting, Silvio Lattanzi, Renato Paes Leme, and Sergei Vassilvitskii. Secretaries with advice. In Proceedings of the 22nd ACM Conference on Economics and Computation, pp. 409–429, 2021. Ken Goldberg, Theresa Roeder, Dhruv Gupta, and Chris Perkins. Eigentaste: A constant time collaborative filtering algorithm. Information Retrieval, 4(2):133–151, 2001. Noah Golowich and Ankur Moitra. Can q-learning be improved with advice? In Conference on Learning Theory, pp. 4548–4619. PMLR, 2022. Igor Grossmann, Matthew Feinberg, Dawn C Parker, Nicholas A Christakis, Philip E Tetlock, and William A Cunningham. Ai and the transformation of social science research. Science, 380(6650):1108–1109, 2023. F Maxwell Harper and Joseph A Konstan. The movielens datasets: History and context. ACM Transactions on Interactive Intelligent Systems (TiiS), 5(4):1–19, 2015. John J Horton, Apostolos Filippas, and Benjamin S Manning. Large language models as simulated economic agents: What can we learn from homo silicus? Technical report, National Bureau of Economic Research, 2023. Tiancheng Hu and Nigel Collier. Quantifying the persona effect in llm simulations. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 10289–10307, 2024. Wenlong Ji, Yihan Pan, Ruihao Zhu, and Lihua Lei. Multi-armed bandits with machine learning-generated surrogate rewards. arXiv preprint arXiv:2506.16658, 2025. Wenlong Ji, Weizhe Yuan, Emily Getzen, Kyunghyun Cho, Michael I Jordan, Song Mei, Jason Weston, Weijie J Su, Jing Xu, and Linjun Zhang. An overview of large language models for statisticians. The American Statistician, (just-accepted):1–106, 2026. Ron Kohavi, Roger Longbotham, Dan Sommerfield, and Randal M Henne. Controlled experiments on the web: survey and practical guide. Data mining and knowledge discovery, 18(1):140–181, 2009. Stefan Krsteski, Giuseppe Russo, Serina Chang, Robert West, and Kristina Gligorić. Valid survey simulations with limited human data: The roles of prompting, fine-tuning, and rectification. arXiv preprint arXiv:2510.11408, 2025. Silvio Lattanzi, Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Online scheduling via learned weights. In Proceedings of the Fourteenth Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 1859–1877, 2020. 11

Ang Li, Haozhe Chen, Hongseok Namkoong, and Tianyi Peng. Llm generated persona is a promise with a catch. arXiv preprint arXiv:2503.16527, 2025. Lihong Li, Jin Young Kim, and Imed Zitouni. Toward predicting the outcome of an a/b experiment for search relevance. In Proceedings of the Eighth ACM International Conference on Web Search and Data Mining, pp. 37–46, 2015. Honghao Lin, Tian Luo, and David Woodruff. Learning augmented binary search trees. In International Conference on Machine Learning, pp. 13431–13440. PMLR, 2022. Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 8086–8098, 2022. Thodoris Lykouris and Sergei Vassilvitskii. Competitive caching with machine learned advice. International Conference on Machine Learning, pp. 3296–3305, 2018. Benjamin S Manning, Kehang Zhu, and John J Horton. Automated social science: Language models as scientist and subjects. Technical report, National Bureau of Economic Research, 2024. Saab Mansour, Leonardo Perelli, Lorenzo Mainetti, George Davidson, and Stefano D’Amato. Paars: Persona aligned agentic retail shoppers. In Proceedings of the 1st Workshop for Research on Agent Language Models (REALM 2025), pp. 143–159, 2025. Samuel McCauley, Ben Moseley, Aidin Niaparast, and Shikha Singh. Online list labeling with predictions. Advances in Neural Information Processing Systems, 36:60278–60290, 2023. Michael Mitzenmacher. Scheduling with predictions and the price of misprediction. In Proceedings of the 11th Innovations in Theoretical Computer Science Conference, pp. 14:1–14:18, 2020. Michael Mitzenmacher and Sergei Vassilvitskii. Algorithms with predictions. Communications of the ACM, 65(7):33–35, 2022. 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. Manish Purohit, Zoya Svitkina, and Ravi Kumar. Improving online algorithms via ml predictions. In Advances in Neural Information Processing Systems, volume 31, 2018. Vinod Raman and Ambuj Tewari. Online classification with predictions. Advances in Neural Information Processing Systems, 37:55884–55914, 2024. Tim Rieder, Marian Schneider, Mario Truss, Vitaly Tsaplin, Alina Rublea, Sinem Dere, Francisco Chicharro Sanz, Tobias Reiss, and Mustafa Doga Dogan. Simab: Simulating a/b tests with persona-conditioned ai agents for rapid design evaluation. arXiv preprint arXiv:2603.01024, 2026. Dhruv Rohatgi. Near-optimal bounds for online caching with machine learned advice. In Proceedings of the Fourteenth Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 1834–1845, 2020. Adam Sales, AF Botelho, Thanaporn Patikorn, and Neil T Heffernan. Using big data to sharpen design-based inference in a/b tests. In Proceedings of the eleventh international conference on educational data mining, 2018. Gideon Schwarz. Estimating the dimension of a model. The Annals of Statistics, 6(2):461–464, 1978. 12

Noa Shoham, Ron Dorfman, Shalev Shaer, Kfir Y Levy, and Yaniv Romano. Predictionpowered semi-supervised learning with online power tuning. arXiv preprint arXiv:2510.22586, 2025. Burrhus Frederic Skinner. Science and human behavior. Number 92904. Simon and Schuster, 1965. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. Joel A Tropp. An introduction to matrix concentration inequalities. Foundations and Trends in Machine Learning, 8(1-2):1–230, 2015. Aad W Van der Vaart. Asymptotic statistics, volume 3. Cambridge university press, 2000. Ya Xu, Nanyu Chen, Addrian Fernandez, Omar Sinno, and Anmol Bhasin. From infrastructure to culture: A/b testing challenges in large scale social networks. In Proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining, pp. 2227–2236, 2015. Ali Zeynali, Shahin Kamali, and Mohammad Hajiesmaili. Robust learning-augmented dictionaries. In Forty-first International Conference on Machine Learning. Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International conference on machine learning, pp. 12697–12706. Pmlr, 2021. Cai-Nicolas Ziegler, Sean M McNee, Joseph A Konstan, and Georg Lausen. Improving recommendation lists through topic diversification. In Proceedings of the 14th international conference on World Wide Web, pp. 22–32, 2005.

Population level sign prediction

In this section we provide the proof of Theorem 2.1. The first Lemma below gives the expression of the power of the asymmetric Z-test. Then we prove the consistency and robustness bounds in separate lemmas. Lemma A.1. For the asymmetric Z-test with critical values at quantiles (1 − λ)α/2 and 1 − (1 + λ)α/2, the power function is: Power(θ, α, λ) = Φ(z(1−λ)α/2 − θ ) + Φ(θ − z1−(1+λ)α/2 ) Proof. We test H0 : θ = 0 vs H1 : θ ̸= 0 using test statistic Z ∼ N (θ, 1) under the alternative with effect size θ. The rejection region is: Z < z(1−λ)α/2

Z > z1−(1+λ)α/2 .

Therefore, since Z − θ ∼ N (0, 1) the power is: Power(θ, α, λ) = P( Z < z(1−λ)α/2 | θ ) + P( Z > z1−(1+λ)α/2 | θ ) = P Z − θ < z(1−λ)α/2 − θ + P Z − θ > z1−(1+λ)α/2 − θ = Φ(z(1−λ)α/2 − θ ) + 1 − Φ(z1−(1+λ)α/2 − θ )

= Φ(z(1−λ)α/2 − θ ) + Φ(θ − z1−(1+λ)α/2 ) , where we used for the last inequality the symmetry Φ(−u) = 1 − Φ(u). A.1

Consistency

√ 2 Lemma A.2. Denote Rby ϕ : x 7→ e− x /2 / 2π the PDF of the standard normal distribution x N (0, 1), and Φ : x 7→ −∞ ϕ(t)dt its CDF. Then f = ϕ/Φ is decreasing. Proof. To demonstrate the result, we will show that the derivative of f is negative. It holds for all x ∈ R that ϕ ′ ( x ) Φ ( x ) − ϕ ( x )2 f ′ (x) = . Φ ( x )2 √ 2 Recalling that ϕ( x ) = e− x /2 / 2π, we have that ϕ′ ( x ) = − xϕ( x ), and since ϕ and Φ are both positives, we have for all x ∈ R that f ′ ( x ) < 0 ⇐⇒ ϕ′ ( x )Φ( x ) < ϕ( x )2

⇐⇒ − xϕ( x )Φ( x ) < ϕ( x )2 ⇐⇒ − xΦ( x ) < ϕ( x ) . Therefore, it suffices to prove that − xΦ( x ) < ϕ( x ) for all x. Let x ∈ R, we have 2

e−t /2 √ dt = −∞ 2π

e−t /2 (− x ) √ dt −∞ 2π 2 Z x e−t /2 1 h −t2 /2 i x = ϕ( x ) . < (−t) √ dt = √ e −∞ −∞ 2π 2π

− xΦ( x ) = − x

Z x

which concludes the proof. Lemma A.3 (Consistency). For any fixed λ ∈ (0, 1), we have for all θ > 0 and α ∈ (0, 1) that Power(θ, α, λ) 1+λ ≥ Power(θ, α, 1) 2 14

Proof. Let λ ∈ [0, 1], θ ≥ 0, and α ∈ (0, 1). We have by Lemma A.1 that Power(θ, α, λ) = Φ(z(1−λ)α/2 − θ ) + Φ(θ − z1−(1+λ)α/2 ) Power(θ, α, 1) = Φ(θ − z1−α ) We need to prove that Φ(z(1−λ)α/2 − θ ) + Φ(θ − z1−(1+λ)α/2 ) ≥

1+λ Φ ( θ − z 1− α ) . 2

Since Φ > 0, it suffices to prove that Φ(θ − z1−(1+λ)α/2 ) ≥

1+λ Φ ( θ − z 1− α ) . 2

(6)

Let us denote by t = z1−α and s = z1−(1+λ)α/2 . We have that (1 + λ)α/2 ≤ α and u 7→ z1−u is decreasing, hence s ≥ t. Now observe that for all x ∈ R we have d Φ( x − s) ϕ( x − s)Φ( x − t) − Φ( x − s)ϕ( x − t) = dx Φ( x − t) Φ ( x − t )2 Φ( x − s) ϕ( x − s) ϕ( x − t) = − Φ( x − t) Φ( x − s) Φ( x − t) ≥0, where we used for the last inequality that ϕ/Φ is decreasing (by Lemma A.2) and that Φ( x −s) x − s ≤ x − t. We deduce that x 7→ Φ( x−t) is non-decreasing, hence for θ ≥ 0 we have Φ(θ − s) Φ(−s) 1 − Φ(s) ≥ = Φ(θ − t) Φ(−t) 1 − Φ(t) 1 − Φ(z1−(1+λ)α/2 ) (1 + λ)α/2) 1+λ = = = 1 − Φ ( z 1− α ) α 2 where we used the symmetry property of the normal CDF Φ(−u) = 1 − Φ(u) and the definition of the quantiles zu = Φ−1 (u). Therefore, we deduce that Equation (6) is true, which concludes the proof.

A.1.1

Robustness

Lemma A.4 (Robustness). For any fixed λ ∈ (0, 1), we have for all θ > 0 and α ∈ (0, 1) that Power(θ, α, λ) ≥ 1−λ Power(θ, α, 0) Proof. The power functions are Power(θ, α, 0) = Φ(zα/2 − θ ) + Φ(θ − z1−α/2 ) Power(θ, α, λ) = Φ(z(1−λ)α/2 − θ ) + Φ(θ − z1−(1+λ)α/2 ) We need to prove that

∀θ ≤ 0 :

Power(θ, α, λ) − (1 − λ)Power(θ, α, 0) ≥ 0 ,

Defining for all θ ≤ 0 the functions g(θ ) = Φ(z(1−λ)α/2 − θ ) − (1 − λ)Φ(zα/2 − θ ) h(θ ) = Φ(θ − z1−(1+λ)α/2 ) − (1 − λ)Φ(θ − z1−α/2 ) , The objective (7) is equivalent to

∀θ ≤ 0 :

g(θ ) + h(θ ) ≥ 0 .

In the following, we will prove separately that both g and h and non-negative for θ ≤ 0. 15

(7)

Analysis of g(θ ) Let a = zα/2 , b = z(1−λ)α/2 . We can write g(θ ) = Φ(b − θ ) − (1 − λ)Φ( a − mu), and we have for all θ ≤ 0 that g ′ ( θ ) = − ϕ ( b − θ ) + (1 − λ ) ϕ ( a − θ ) √ 2 Since α/2 ≥ (1 − λ)α/2, then a ≥ b. Moreover, ϕ : x 7→ e− x /2 / 2π is decreasing on [0, ∞) and we have for θ ≤ 0 that a − θ ≥ b − θ ≥ 0, hence ϕ( a − θ ) ≥ ϕ(b − θ ), which yields g ′ ( θ ) = − ϕ ( b − θ ) + (1 − λ ) ϕ ( a − θ )

≤ − ϕ ( b − θ ) + (1 − λ ) ϕ ( b − θ ) = −λϕ(b − θ ) ≤ 0 . We deduce that g is non-increasing on (−∞, 0], and consequently we have for all θ ≤ 0 that g ( θ ) ≥ g (0)

= Φ ( b ) − (1 − λ ) Φ ( a ) = Φ(z(1−λ)α/2 ) − (1 − λ)Φ(zα/2 ) = (1 − λ)α/2 − (1 − λ)α/2 =0, which proves that g(θ ) ≥ 0 for all θ ≤ 0. Analysis of h(θ ): Let c = z1−α/2 and d = z1−(1+λ)α/2 , then h(θ ) = Φ(θ − d) − (1 − λ)Φ(θ − c), and we have for all θ ≤ 0 that h ′ ( θ ) = ϕ ( θ − d ) − (1 − λ ) ϕ ( θ − c ) . Using similar arguments as for the analysis of g, observe that since (1 + λ)α/2 > α/2, then √ 2 d < c. Moreover, ϕ : x 7→ e− x /2 / 2π is increasing on (−∞, 0] and we have for θ ≤ 0 that θ − c ≤ θ − d ≤ 0, hence ϕ(θ − c) ≤ ϕ(θ − d). The derivative of h can be lower bounded as h ′ ( θ ) = ϕ ( θ − d ) − (1 − λ ) ϕ ( θ − c )

≥ ϕ ( θ − c ) − (1 − λ ) ϕ ( θ − c ) = λϕ(θ − c) ≥ 0 . Therefore, h is non-decreasing on (−∞, 0] and it holds for all θ ≤ 0 that h(θ ) ≥ lim h(u) u→−∞

= lim [Φ(u − d) − (1 − λ)Φ(u − c)] u→−∞

=0 where the last inequality is true because lim Φ(u − d) = lim Φ(u − c) = 0. u→−∞

u→−∞

All in all, we proved that g and h are both non-negative on (−∞, 0], and consequently Power(θ, α, λ) − (1 − λ)Power(θ, α, 0) ≥ 0 for all θ ≤ 0, which concludes the result.

Individual-level predictions - GPPI

B.1

Proof of Theorem 3.1

We establish three key lemmas, then prove the main result. Lemma B.1 (Fixed Weight Asymptotics). For any fixed λ ∈ Rd , define θ̂ (λ) = Ȳn + λ T (ϕ̄U N− ϕ̄nL ). Then: √ d n(θ̂ (λ) − θ ∗ ) − → N (0, V (λ)) where V (λ) = σR2 (λ) + rλ T Σϕ λ and σR2 (λ) = Var(Y − λ T ϕ( f ( X ))). 16

Proof. By linearity of expectation, E[θ̂ (λ)] = θ ∗ (unbiasedness), and defining µϕ = E[ϕ( f ( X ))] we have that θ (λ) − θ ∗ = (Ȳn − λ T ϕ̄nL ) − (θ ∗ − λ T µϕ ) + λ T (ϕ̄U N − µϕ )

1 n λT N Ri ( λ ) + (ϕ( f ( X̃ j )) − µϕ ) , ∑ n i =1 N i∑ =1

where Ri (λ) = (Yi − θ ∗ ) − λ T (ϕ( f ( Xi )) − µϕ ), i = 1, . . . , n . √ where µϕ = E[ϕ( f ( X ))]. Muliplying by n gives r √ 1 n n λT N U n(θ̂ (λ) − θ ∗ ) = √ ∑ Ri (λ) + · √ ∑ ϕj . N n i =1 N j =1 Term 1:

(8)

Ri (λ) are i.i.d. with E[ Ri ] = 0 and Var( Ri ) = Var(Y − λ T ϕ( f ( X )))

= σY2 + λ T Σϕ λ − 2λ T σY,ϕ =: σR2 (λ) < ∞ since E[Y 2 ] < ∞ and E[∥ϕ( f ( X ))∥2 ] < ∞. By the Central Limit Theorem 1 n d √ ∑ Ri ( λ ) − → N (0, σR2 (λ)) . n i =1 U U Term 2: ϕU j are i.i.d. with E[ ϕ j ] = 0 and Cov( ϕ j ) = Σϕ . The Multivariate Central Limit Theorem yields 1 N d √ ∑ ϕU → Nd (0, Σϕ ) , j − N j =1

and we deduce that

λT N d √ ∑ ϕU → N (0, λ T Σϕ λ) . j − N j =1 √ √ Since rn = n/N → r, we have n/N → r. By Slutsky’s theorem: r n λT N U d · √ ∑ ϕj − → N (0, rλ T Σϕ λ) N N j =1

Since labeled and unlabeled data are independent, the two sums in (8) are independent. Therefore √ d n(θ̂ (λ) − θ ∗ ) − → N (0, σR2 (λ) + rλ T Σϕ λ) = N (0, V (λ)) . −1 1 Lemma B.2 (Optimal Weight). V (λ) is minimized at λ∗ = 1+ r Σϕ σY,ϕ with minimum value:

V ∗ = σY2 − Proof. Expanding σR2 (λ): Therefore

1 σ T Σ −1 σ 1 + r Y,ϕ ϕ Y,ϕ

σR2 (λ) = σY2 + λ T Σϕ λ − 2λ T σY,ϕ V (λ) = σY2 − 2λ T σY,ϕ + (1 + r )λ T Σϕ λ , 17

and

∇λ V = −2σY,ϕ + 2(1 + r )Σϕ λ .

The Hessian is ∇2λ V = 2(1 + r )Σϕ ≻ 0 since Σϕ ≻ 0, hence V (λ) is convex, and it is minimal at the point where the gradient is zero, i.e. at λ∗ =

1 Σ −1 σ . 1 + r ϕ Y,ϕ

Substituting λ∗ gives V ∗ = σY2 − 2(λ∗ ) T σY,ϕ + (1 + r )(λ∗ ) T Σϕ λ∗ 2 1 T 1 = σY2 − σY,ϕ Σ− σ T Σ −1 σ ϕ σY,ϕ + 1+r 1 + r Y,ϕ ϕ Y,ϕ 1 = σY2 − σ T Σ −1 σ . 1 + r Y,ϕ ϕ Y,ϕ Lemma B.3 (Consistency of Plug-in Estimators). Under the conditions of Theorem 3.1: p

Σ̂ϕ − → Σϕ ,

σ̂Y,ϕ − → σY,ϕ ,

λ̂ − → λ∗ .

Proof. By the Strong Law of Large Numbers, since E[∥ϕ( f ( X ))∥2 ] < ∞ and E[Y2 ] < ∞, we have p p Σ̂ϕ − → Σϕ , σ̂Y,ϕ − → σY,ϕ . Since Σϕ ≻ 0, matrix inversion is continuous at Σϕ . By the continuous mapping theorem p

1 1 Σ̂− → Σ− ϕ − ϕ .

Since rn → r and 1/(1 + rn ) → 1/(1 + r ), we deduce by Slutsky’s theorem that λ̂ =

B.1.1

p 1 1 1 → Σ −1 σ = λ ∗ . Σ̂− ϕ σ̂Y,ϕ − 1 + rn 1 + r ϕ Y,ϕ

Proof of Main Theorem

Proof of Theorem 3.1. We prove below the the three statements of the theorem. (i): Asymptotic normality. Decompose √ √ √ n(θ̂GPPI − θ ∗ ) = n(θ̂ (λ∗ ) − θ ∗ ) + n(θ̂ (λ̂) − θ̂ (λ∗ )) . √ d We have by Lemma B.1 that n(θ̂ (λ∗ ) − θ ∗ ) − → N (0, V ∗ ). For the second term, we can write √ √ L n(θ̂ (λ̂) − θ̂ (λ∗ )) = n(λ̂ − λ∗ ) T (ϕ̄U N − ϕ̄n ) , −1/2 ) and ϕ̄ L − µ = O ( n−1/2 ). Since and we have by the CLT that ϕ̄U ϕ p n N − µϕ = O p ( N n/N → r ∈ (0, ∞), we have N = Θ(n), hence L −1/2 ϕ̄U ). N − ϕ̄n = O p ( n

Furthermore, Lemma B.3 gives λ̂ − λ∗ = o p (1). Therefore: √ √ n(θ̂ (λ̂) − θ̂ (λ∗ )) = n · o p (1) · O p (n−1/2 ) = o p (1) . Finally, we deduce by Slutsky’s theorem that √ d n(θ̂GPPI − θ ∗ ) − → N (0, V ∗ ) 18

(ii): Variance estimation.

Define

R̃i = (Yi − θ ∗ ) − (λ∗ ) T (ϕ( f ( Xi )) − µϕ ) . we have that Var( R̃i ) = σR2 (λ∗ ). The sample residuals satisfy R̂i = (Yi − Ȳn ) − λ̂ T (ϕ( f ( Xi )) − ϕ̄nL )

= R̃i + (λ∗ − λ̂)T (ϕ( f ( Xi )) − µϕ ) + O p (n−1/2 ) . By standard results on sample variance (see for example (Van der Vaart, 2000)) and Lemma B.3: n p 1 R̂2i − → σR2 (λ∗ ) . σ̂R2 = ∑ n − 1 i =1 Using again Lemma B.3 and the continuity of (λ, Σ) 7→ (λ T Σλ, we obtain p

λ̂ T Σ̂ϕ λ̂ − → (λ∗ ) T Σϕ λ∗ Since rn = n/N → r: V̂ = σ̂R2 + rn λ̂ T Σ̂ϕ λ̂ p

− → σR2 (λ∗ ) + r (λ∗ )T Σϕ λ∗ = V (λ∗ ) = V ∗ (iii): Confidence interval validity. From claims (i)–(ii) and Slutsky’s theorem, we deduce that √ n(θ̂GPPI − θ ∗ ) d √ − → N (0, 1) . V̂ Therefore ! √ n(θ̂GPPI − θ ∗ ) √ ≤ zα/2 = 1 − α lim P n,N →∞ V̂ which is equivalent to limn,N →∞ P(θ ∗ ∈ CI1−α ) = 1 − α. B.2

Efficiency gains

Here we quantify how much the confidence interval width can be reduced using GPPI compared to PPI++. Lemma B.4. Under the conditions of Theorem 3.1, the following hold: (i) The asymptotic variance admits the decomposition V ∗ = σY2

1−

R2ϕ 1+r

! ,

(9)

T Σ −1 σ 2 where R2ϕ = σY,ϕ ϕ Y,ϕ /σY ∈ [0, 1] is the population coefficient of determination from regressing Y on ϕ( f ( X )). 2 , where ρ (ii) For PPI++, with ϕ(z) = z, we have R2ϕ = ρY, Y, f = Cor(Y, f ( X )). f

(iii) The asymptotic confidence interval half-widths satisfy s s 2 ρY, R2ϕ f WGPPI = Wclassical 1 − ≤ WPPI = Wclassical 1 − ≤ Wclassical , 1+r 1+r q where Wclassical = zα/2 σY2 /n. 19

(10)

(iv) The width ratio is given by v u u 1 − R2ϕ /(1 + r ) WGPPI =t . 2 / (1 + r ) WPPI 1 − ρY, f

(11)

2 , which occurs when nonlinear Strict inequality WGPPI < WPPI holds if and only if R2ϕ > ρY, f transformations capture additional structure beyond linear correlation.

Proof. We prove below the three caims of the lemma. Claim (i):

From Theorem 3.1(i), the asymptotic variance is V

∗

= σY2 −

1 σ T Σ−1 σ = σY2 1 + r Y,ϕ ϕ Y,ϕ

T Σ −1 σ σY,ϕ 1 ϕ Y,ϕ 1− · 1+r σY2

! .

(12)

T Σ −1 σ 2 2 Define R2ϕ = σY,ϕ ϕ Y,ϕ /σY . This is the population R from the best linear predictor of Y

given ϕ( f ( X )). By the Cauchy-Schwarz inequality, R2ϕ ∈ [0, 1]. Claim (ii): When ϕ(z) = z, we have Σϕ = Var( f ( X )) = σ2f and σY,ϕ = Cov(Y, f ( X )). Thus Cov(Y, f ( X ))2 2 = Cor(Y, f ( X ))2 = ρY, (13) R2ϕ = f. σY2 · σ2f

√ Claim (iii): The asymptotic (1 − α) confidence interval has half-width W = zα/2 V ∗ /n. q

For the classical estimator, Vclassical = σY2 , giving Wclassical = zα/2 σY2 /n. Substituting the variance decomposition from Claim (i): s q R2ϕ 2 2 WGPPI = zα/2 σY (1 − Rϕ /(1 + r ))/n = Wclassical 1 − . (14) 1+r q 2 / (1 + r ) by Claim (ii). Similarly, WPPI = Wclassical 1 − ρY, f

Since ϕ(z) = z is a special case of the general transformation (taking d = 1 and including only the linear term), and R2ϕ measures variance explained by a d-dimensional transfor2 mation while ρY, f measures variance explained by a 1-dimensional predictor, we have 2 . Therefore W R2ϕ ≥ ρY, GPPI ≤ WPPI ≤ Wclassical . f

Claim (iv): Taking the ratio of the expressions in Claim (iii) yields the stated formula. B.3

Matrix concentration inequalitites

The two Lemmas we prove below are immediate corollaries of the Bernstein inequalities for symmetrical then rectangular matrices (Theorems 1.6.2 and 7.3.1 in (Tropp, 2015)). Lemma B.5. Let Z1 , . . . , Zn be independent random vectors in Rd , satisfying for all i ∈ [n] that ∥ Zi ∥2 ≤ M and having the same covariance matrix Σ. Let Σ̂ = n1 ∑in=1 Zi ZiT , then for all t ≥ 0 −nt2 P ∥Σ̂ − Σ∥op ≥ t ≤ 2d exp . 8M (∥Σ∥op + 2t/3) In particular, if d = ω (1), M = Ω(1) and M2 log d = o (n) then ! r M2 log d ∥Σ̂ − Σ∥op = O p . n 20

Proof. By the matrix Bernstein inequality (Tropp, 2015), for independent centered symmetric matrices Q1 , . . . , Qn with ∥ Qi ∥op ≤ L almost surely:    n  t2 /2 P Q ≥ t ≤ 2d exp − , i  i∑  σ2 + Lt/3 =1 op

where σ2 =

∑in=1 E[ Q2i ] op .

Define for all i ∈ [n] the matrix Qi = n1 ( Z̃i Z̃iT − Σ). The matrices Qi are centered, and symmetric. To use Bernstein inequality, we will prove first the bound L on the operator norm almost surely, then the bound on σ2 . Using that ∥ Z̃i ∥2 ≤ M a.s., by definition of the operator norm and Cauchy Schwarz inequality we obtain

∥ Z̃i Z̃iT ∥op = sup u T Z̃i Z̃iT u = sup ( Z̃iT u)2 ≤ sup (∥ Z̃i ∥∥u∥)2 ≤ M a.s. , ∥u∥=1

∥u∥=1

and by Jensen’s inequality ∥Σ∥op = ∥E[ Z̃i Z̃iT ]∥op ≤ E∥ Z̃i Z̃iT ∥op ≤ M. Therefore ∥ Qi ∥op ≤ 1 1 T n ∥ Z̃i Z̃i ∥op + n ∥ Σ ∥op ≤ 2M/n a.s. On the other hand, ∑in=1 E[ Q2i ] op ≤ ∑in=1 E[ Q2i ] op = n12 ∑in=1 E[( Z̃i Z̃iT − Σ)2 ] op , and for all i ∈ [n] we have E[( Z̃i Z̃iT − Σ)2 ] = E[( Z̃i Z̃iT )2 ] − Σ2

⪯ E[( Z̃i Z̃iT )2 ] = E[ Z̃i ( Z̃iT Z̃i ) Z̃iT ] = E[∥ Z̃i ∥2 Z̃i Z̃iT ] ⪯ ME[ Z̃i Z̃iT ] = MΣ . Therefore σ2 = ∑in=1 E[ Q2i ] op ≤ M ∥Σ∥op /n Therefore, using Bernstein’s inequality with L = 2M/n and σ2 ≤ M ∥Σ∥op /n gives −nt2 P ∥Σ̂ − Σ∥op ≥ t ≤ 2d exp . 2M (∥Σ∥op + 2t/3) which proves the concentration inequality for centered random variables. Now in the general case, if we do not assume that ( Zi )i are centered, then we define for all i the centered random √ variable Z̃i = Zi − E[ Zi ]. With the triangle then Jensen’s inequality we have that ∥ Z̃i ∥ ≤ 2 M, and the previous result holds by replacing M with 4M: −nt2 P ∥Σ̂ − Σ∥op ≥ t ≤ 2d exp , 8M (∥Σ∥op + 2t/3) which proves the claimed concentration bound. Using that ∥Σ∥op ≤ M as we proved earlier, we also have that −nt2 . P ∥Σ̂ − Σ∥op ≥ t ≤ 2d exp 2M ( M + 2t/3) √ 4M log d √ Now assuming that M = Ω(1) and M2 log d = o (n), taking t = gives n ( ) p 4M log d −16M2 log d √ P ∥Σ̂ − Σ∥op ≥ ≤ 2d exp 8M( M + o (1)) n

= 2d exp (−(2 − o (1)) log d) = 2/d1−o(1) = o (1) , hence

∥Σ̂ − Σ∥op = O p

M2 log d n

! .

Lemma B.6. Let Z1 , . . . , Zn be independent random vectors in Rd such that ∥ Zi ∥2q≤ M a.s. and

Y1 , . . . , Yn be real-valued random variables such that Yi2 ≤ KY a.s. Then for t ≥ 11 KYnM , −nt2 d √ P{∥Cov(Y, Z ) − Cov(Y, Z )∥ ≥ t} ≤ 4(d + 1) exp . 16(8KY M + KY Mt/3) Proof. Define the centered variables Ỹi = Yi − E[Y1 ] and Z̃i = Zi − E[ Z1 ], then: n

d (Y, Z ) = 1 ∑ Ỹi Z̃i , Cov n i =1

Cov(Y, Z ) = E[Ỹ1 Z̃1 ]

Define Si = n1 (Ỹi Z̃i − E[Ỹ1 Z̃1 ]) ∈ Rd×1 for i = 1, . . . , n. Then: n

d (Y, Z ) − Cov(Y, Z ) = ∑ Si Cov i =1

We will use for our proof the Intrinsic Matrix Bernstein inequality (Theorem 7.3.1 in (Tropp, 2015)). Let us first verify that (Si )i satisfy the assumptions of the theorem. 1. Zero mean E[Si ] = n1 (E[Ỹi Z̃i ] − E[Ỹ1 Z̃1 ]) = 0. 2. Uniform bound We need to bound ∥Ỹi Z̃i ∥2 = |Ỹi |∥ Z̃i ∥2 . By triangle inequality then Jensen inequality we have: |Ỹi | = |Yi − E[Y1 ]| ≤ |Yi | + E[|Y1 |] ≤ √ √ 2 KY a.s., and similarly ∥ Z̃i ∥2 ≤ 2 M a.s. Therefore: √ p p ∥Ỹi Z̃i ∥ ≤ 2 KY · 2 M = 4 KY M a.s. √ And by Jensen’s inequality ∥E[Ỹ1 Z̃1 ]∥ ≤ 4 KY M. Thus:

∥Si ∥op ≤

√ p 1 p 8 KY M ( 4 KY M + 4 KY M ) = := L n n

3. Matrix-valued variances We need to prove semidefinite upper bounds on E[Si SiT ] and E[SiT Si ]. We have immediately that n

i =1

∑ E[Si SiT ] ⪯ ∑ E[∥Si ∥22 ] Id ⪯

64KY M Id n

and SiT Si ∈ R and satisfies n

i =1

∑ E[SiT Si ] = ∑ E[∥Si ∥22 ] ≤

64KY M n

4. Apply Intrinsic Matrix Bernstein Therefore, using the intrinsic Matrix Bernstein inequal√ ity (Theorem 7.3.1 in (Tropp, 2015)) with v = 64KnY M and L = 8 KnY M , we obtain for all √

√

KY M t ≥ 8 √KnY M + 8 3n :

d (Y, Z ) − Cov(Y, Z )∥2 ≥ t} ≤ 4(d + 1) exp P{∥Cov It follows that for all t ≥ 11

−t2 /2 v + Lt/3

KY M n :

d (Y, Z ) − Cov(Y, Z )∥2 ≥ t} ≤ 4(d + 1) exp P{∥Cov

−nt2 √ 16(8KY M + KY Mt/3)

B.4

Adaptive Transformation Selection

This section provides the proof of Theorem 3.2, which establishes that empirical variance minimization asymptotically selects the optimal transformation. B.4.1

Notation and Setup

We remind here the notation and assumptions used in the proof. For each transformation j ϕ j : R → Rd with j ∈ [m], define: • Population quantities: Σ j = Cov(ϕ j ( f ( X ))), σY,j = Cov(Y, ϕ j ( f ( X ))) −1 1 • Optimal weight: λ∗j = 1+ r Σ j σY,j −1 1 T • Asymptotic variance: Vj = σY2 − 1+ r σY,j Σ j σY,j

• Sample estimators: Σ̂ j , σ̂Y,j , λ̂ j , V̂j as defined in Section 3.2 Let j∗ = arg min j∈[m] Vj denote the oracle-optimal transformation and ĵ = arg min j∈[m] V̂j the empirically selected transformation. Assumption B.7. For all j ∈ [m]: (i) ∥ϕ j ( f ( X ))∥2 ≤ M almost surely, where M is a uniform bound (ii) Y 2 ≤ KY almost surely (iii) Σ j = Cov(ϕ j ( f ( X ))) ≻ 0 (positive definite) (iv) d j ≤ d for all j, where d is a fixed constant (v) d log(m) = o (n) B.4.2

Preliminary lemmas

Lemma B.8 (Uniform Covariance Concentration). Under Assumption B.7, for any δ ∈ (0, 1), we have for N sufficiently large that ( ) r M2 log(2dm/δ) P max ∥Σ̂ j − Σ j ∥op ≥ 3 ≤δ N j∈[m] Proof. Let j ∈ [ M ], the random vectors (ϕ j ( f ( X̃i )))i are independent and have their norm bounded by M almost surely. Therefore, by union-bound and Lemma B.5 we have for all t ≥ that ( ) m − Nt2 P max ∥Σ̂ j − Σ j ∥op ≥ t ≤ ∑ P ∥Σ̂ j − Σ j ∥op ≥ t ≤ 2dm exp . 8M ( M + 2t/3) j∈[m] j =1 q q M2 log(2dm/δ) 4 log(2dm/δ) Let δ > 0, and t = 3 . For N sufficiently large we have ≤ 1 and N N ( ) r M2 log(2dm/δ) − Nt2 P max ∥Σ̂ j − Σ j ∥op ≥ 3 ≤ 2dm exp N 8M ( M + 2t/3) j∈[m]     −9 log(2dm/δ)  = 2dm exp  q   4 log(2dm/δ) 8 1+ N

≤ 2dm exp (− log(2dm/δ)) = δ

Lemma B.9 (Uniform Cross-Covariance Concentration). Under Assumption B.7, for any δ ∈ (0, 1), we have for N sufficiently large that ( ) r 2KY M log(2md/δ) P max ∥σ̂Y,j − σY,j ∥2 ≥ 3 ≤δ. n j∈[m] Proof. For each j ∈ [m], we apply Lemma B.6 with Zi = ϕ j ( f ( Xi )) and Yi . By Assumption B.7(i) and (ii), we have ∥ϕ j ( f ( Xi ))∥2 ≤ M a.s. and Yi2 ≤ KY a.s. q By Lemma B.6, for t ≥ 11 KYnM : −nt2 j √ P{∥σ̂Y,j − σY,j ∥2 ≥ t} ≤ 4(d + 1) exp 16(8KY M + KY Mt/3) q

KY M log(2md/δ) with C ≥ 11. Assuming that d ≥ 2 or m ≥ 2 or δ ≤ 1/2 we have n q √ log(2md/δ) ≥ 1, hence t ≥ 11 KYnM . For N sufficiently large we have K3Y M t ≤ 2KY M, we

Set t = C have:

nt2 C2 log(2md/δ) √ , ≥ 160 16(8KY M + KY Mt/3)

and taking C2 ≥ 320 (i.e. C ≥ 18) ensures: P{∥σ̂Y,j − σY,j ∥2 ≥ t} ≤ 4(d + 1) exp(−2 log(2md/δ))

4( d + 1) . (2md/δ)2

By the union bound over all j ∈ [m]: ( P max ∥σ̂Y,j − σY,j ∥2 ≥ t j∈[m]

)

≤

4m(d + 1) 8md ≤ (2md/δ)2 (2md/δ)2

2δ2 =δ. md where the last inequality is true when md ≥ 2δ, for example if d ≥ 2 or m ≥ 2. This completes the proof.

Lemma B.10 (Uniform Weight Estimation Error). Under Assumption B.7, for any δ ∈ (0, 1) and n sufficiently large: ( ) r KY M3 log(md/δ) ∗ ≤δ P max ∥λ̂ j − λ j ∥ ≥ C n j∈[m] for some universal constant C > 0. −1 −1 1 1 Proof. Recall that λ Tj = 1+ r Σ j σY,j and λ̂ j = 1+r Σ̂ j σ̂Y,j . We decompose the error as:

1 −1 1 Σ̂ j σ̂Y,j − Σ− σ Y,j j 1+r 1 1 −1 −1 −1 = Σ̂− σ̂ − Σ̂ σ + Σ̂ σ − Σ σ Y,j Y,j Y,j Y,j j j j j 1+r 1 1 −1 −1 = Σ̂− j ( σ̂Y,j − σY,j ) + ( Σ̂ j − Σ j ) σY,j 1+r

λ̂ j − λ Tj =

By the triangle then Cauchy Schwarz inequalities: 1 −1 1 1 ∥λ̂ j − λ Tj ∥ ≤ ∥Σ̂ j ∥ · ∥σ̂Y,j − σY,j ∥ + ∥Σ̂− − Σ− ∥ · ∥σY,j ∥ j j 1+r 24

Bound on ∥σY,j ∥. Jensen’s inequality gives

∥σY,j ∥ = ∥E[Yϕ j ( f ( X ))]∥ ≤ E[|Y | · ∥ϕ j ( f ( X ))∥] ≤

KY M

1 Bound on ∥Σ̂− j ∥op . Let ρmin ( Σ j ) denote the minimum eigenvalue of Σ j . By Assumption B.7(iii), ρmin (Σ j ) > 0. Let κ = min j∈[m] ρmin (Σ j ) > 0.

By Lemma B.8, with probability at least 1 − δ/2: r 3M2 log(4md/δ) κ max ∥Σ̂ j − Σ j ∥op ≤ ≤ N 2 j∈[m]

(15)

for N sufficiently large. On this event, by Weyl’s inequality: ρmin (Σ̂ j ) ≥ ρmin (Σ j ) − ∥Σ̂ j − Σ j ∥op ≥ κ − Therefore: 1 ∥Σ̂− j ∥op =

1 −1 Bound on ∥Σ̂− j − Σ j ∥.

κ κ = . 2 2

2 1 ≤ . κ ρmin (Σ̂ j )

1 −1 −1 −1 Using the resolvent identity Σ̂− j − Σ j = Σ̂ j ( Σ j − Σ̂ j ) Σ j :

1 −1 −1 −1 ∥Σ̂− j − Σ j ∥op ≤ ∥ Σ̂ j ∥op · ∥ Σ̂ j − Σ j ∥op · ∥ Σ j ∥op

On the high-probability event (15): r r 2 3M2 log(4dm/δ) 2 3M2 log(4dm/δ) 1 −1 −1 · = 2 max ∥Σ̂ j − Σ j ∥op ≤ · κ N κ N κ j∈[m] By Lemma B.9, with probability at least 1 − δ/2: r C1 KY M log(4md/δ) max ∥σ̂Y,j − σY,j ∥ ≤ n j∈[m]

Combining the bounds.

Combining all bounds, and assuming M ≥ 1, we have with probability at least 1 − 2δ/3 that: ! r r √ 1 2 C1 KY M log(4md/δ) 2 KY M 3M2 log(4md/δ) T max ∥λ̂ j − λ j ∥ ≤ + 1+r κ n N κ2 j∈[m] r KY M3 log(md/δ) ≤C n for some universal constant C depending on κ and r. This concludes the proof. Lemma B.11 (Uniform Inverse Covariance Bounds). Under Assumption B.7, let κ = min j∈[m] ρmin (Σ j ) > 0. For any δ ∈ (0, 1) and N sufficiently large, with probability at least 1 − δ: r 2 6 M2 log(4dm/δ) −1 −1 −1 and max ∥Σ̂ j − Σ j ∥op ≤ 2 max ∥Σ̂ j ∥op ≤ κ N κ j∈[m] j∈[m] Proof. By Lemma B.8, with probability at least 1 − δ: r 3M2 log(4md/δ) κ max ∥Σ̂ j − Σ j ∥op ≤ ≤ N 2 j∈[m] for N sufficiently large. 25

(16)

1 Bound on ∥Σ̂− j ∥op . On the event (16), by Weyl’s inequality:

ρmin (Σ̂ j ) ≥ ρmin (Σ j ) − ∥Σ̂ j − Σ j ∥op ≥ κ − Therefore: 1 ∥Σ̂− j ∥op =

κ κ = 2 2

1 2 ≤ κ ρmin (Σ̂ j )

1 −1 −1 −1 −1 −1 Bound on ∥Σ̂− j − Σ j ∥op . Using the resolvent identity Σ̂ j − Σ j = Σ̂ j ( Σ j − Σ̂ j ) Σ j :

−1 −1 −1 1 ∥Σ̂− j − Σ j ∥op ≤ ∥ Σ̂ j ∥op · ∥ Σ̂ j − Σ j ∥op · ∥ Σ j ∥op

On the event (16): r

3M2 log(4dm/δ) 1 · N κ √ r 2 3 M2 log(4dm/δ) = 2 N κr 2 M log(4dm/δ) 6 ≤ 2 N κ

2 −1 1 max ∥Σ̂− j − Σ j ∥op ≤ κ · j∈[m]

This completes the proof. Lemma B.12 (Uniform Residual Variance Concentration). Under Assumption B.7, for any δ ∈ (0, 1) and n sufficiently large: ( ) r log(md/δ) 2 P max |V̂j − Vj | ≥ CKY M ≤δ n j∈[m] −1 1 T for some universal constant C > 0, where Vj = σY2 − 1+ r σY,j Σ j σY,j and V̂j is defined analogously with sample quantities.

Proof. Recall that: Vj = σY2 −

1 σ T Σ −1 σ 1 + r Y,j j Y,j

V̂j = σ̂Y2 −

1 σ̂ T Σ̂−1 σ̂ 1 + r Y,j j Y,j

and

We decompose the error as: V̂j − Vj = σ̂Y2 − σY2 −

1 T −1 T 1 σ̂Y,j Σ̂ j σ̂Y,j − σY,j Σ− σ Y,j . j 1+r

In the rest of proof, we will demonstrate high probability bounds on these two terms. 1. Bound on |σ̂Y2 − σY2 |.

We have σY2 = Var(Y ) and σ̂Y2 is the sample variance:

σ̂Y2 =

1 n (Yi − Ȳn )2 , n i∑ =1

σY2 = E[(Y − E[Y ])2 ] .

Define the centered variables Ỹi = Yi − E[Y ]. Then E[Ỹi ] = 0, E[Ỹi2 ] = σY2 , and σ̂Y2 =

n 1 n ¯ )2 = 1 ( Ỹ − Ỹ n i ∑ Ỹi2 − Ỹ¯n2 n i∑ n =1 i =1

(17)

where Ỹ¯n = n1 ∑in=1 Ỹi . We can write

|σ̂Y2 − σY2 | =

1 n 2 Ỹi − Ỹ¯ n2 − σY2 n i∑ =1

≤

1 n 2 Ỹi − σY2 + Ỹ¯ n2 . n i∑ =1

We will bound separately these two terms. We have p |Ỹi | = |Yi − E[Y ]| ≤ |Yi | + |E[Y ]| ≤ 2 KY

a.s.

Hence Ỹi2 ≤ 4KY a.s., and by Jensen’s inequality σY2 = E[Ỹi2 ] ≤ 4KY . Define Wi = Ỹi2 − σY2 . Then: • E[Wi ] = E[Ỹi2 ] − σY2 = 0 • Since 0 ≤ Ỹi2 ≤ 4KY and 0 ≤ σY2 ≤ 4KY , we have |Wi | = |Ỹi2 − σY2 | ≤ 4KY a.s. By Hoeffding’s inequality: ( P q

1 n Wi ≥ t n i∑ =1

)

nt2 ≤ 2 exp − 32KY2

! .

K2 log(4/δ)

Y Setting t = 6 , we get: n   s  1 n 2 log(4/δ )  K 36 log(4/δ) 2 2 Y ≥ 6 P Ỹ − σ ≤ 2 exp − < δ/4 . Y i  n i∑  n 32 =1

Similarly, Ỹ¯ n = n1 ∑in=1 Ỹi is the sum of n iid random variables all bounded almost surely in √ √ [−2 KY , 2 KY ], and Chernoff bound gives for all t ≥ 0 that ! nt2 ¯ P(|Ỹn | ≥ t) ≤ 2 exp − , 32KY2 q KY2 log(4/δ) . hence with probability at least 1 − δ/4 we have that Ỹ¯n ≤ 6 n By the union bound, we deduce that with probability at least 1 − δ/2: 1 n 2 Ỹi − σY2 + Ỹ¯ n2 n i∑ =1 s KY2 log(4/δ) 36KY2 log(4/δ) + ≤6 n n s KY2 log(1/δ) ≤ C1 n

|σ̂Y2 − σY2 | ≤

for some universal constant C1 > 0 and n sufficiently large. −1 T Σ̂−1 σ̂ T 2. Bound on |σ̂Y,j Y,j − σY,j Σ j σY,j |. j

We use the identity:

T 1 T −1 σ̂Y,j Σ̂− j σ̂Y,j − σY,j Σ j σY,j 1 T −1 = (σ̂Y,j − σY,j )T Σ̂− j σ̂Y,j + σY,j Σ̂ j ( σ̂Y,j − σY,j ) T 1 −1 + σY,j (Σ̂− j − Σ j ) σY,j

By the triangle inequality: T 1 T −1 |σ̂Y,j Σ̂− j σ̂Y,j − σY,j Σ j σY,j | 1 ≤ ∥σ̂Y,j − σY,j ∥ · ∥Σ̂− j ∥op · ∥ σ̂Y,j ∥ 1 + ∥σY,j ∥ · ∥Σ̂− j ∥op · ∥ σ̂Y,j − σY,j ∥

−1 1 + ∥σY,j ∥2 · ∥Σ̂− j − Σ j ∥op

In the proof of Lemma B.11, we showed that with probability at least 1 − δ/4, it holds for all q 1 −1 −1 6 j ∈ [m] that ∥Σ̂− j ∥op ≤ 2/κ and ∥ Σ̂ j − Σ j ∥op ≤ κ 2

M2 log(4dm/δ) N

Also, by Lemma B.9, with probability at least 1 − δ/4 we have r 2KY M log(4md/δ) max ∥σ̂Y,j − σY,j ∥ ≤ 3 n j∈[m] and by Jensen’s inequality we have

∥σY,j ∥ = ∥E[Yϕ j ( f ( X ))]∥ ≤ E[|Y | · ∥ϕ j ( f ( X ))∥] ≤

KY M .

We deduce that for n sufficiently large, we have under the previous event that p ∥σ̂Y,j ∥ ≤ ∥σY,j ∥ + ∥σ̂Y,j − σY,j ∥ ≤ 2 KY M . Combining these two bounds, we obtain with probability at least 1 − δ: T 1 T −1 |σ̂Y,j Σ̂− j σ̂Y,j − σY,j Σ j σY,j | r 2KY M log(4md/δ) 2 p · · 2 KY M ≤3 n κ r p 2KY M log(4md/δ) 2 + KY M · · 3 κ n r 2 6 M log(4dm/δ) + KY M · 2 N κ r log(md/δ) ≤ C2 KY M2 n for some constant C2 depending on κ and r.

Substituting in (17), by the union bound over all events, with probability at least 1 − δ: s r KY2 log(1/δ) log(md/δ) 2 max |V̂j − Vj | ≤ C1 + C2 KY M n n j∈[m] r log(md/δ) ≤ CKY M2 n for some universal constant C depending on κ and r. p

Lemma B.13 (Consistency of Selection). Under Assumption B.7, ĵ − → j∗ . Proof. Recall that j∗ = arg min j∈[m] Vj and ĵ = arg min j∈[m] V̂j . We need to show that for any ϵ > 0: P{ ĵ ̸= j∗ } → 0 as n → ∞ Let ∆ = min j̸= j∗ (Vj − Vj∗ ) be the minimum gap between the optimal transformation and all other bases. We assume that ∆ > 0, i.e. j∗ is the unique minimizer. 28

The event { ĵ ̸= j∗ } implies that there exists some j ̸= j∗ such that V̂j ≤ V̂j∗ . This can be written as: V̂j ≤ V̂j∗ ⇐⇒ V̂j − Vj + Vj ≤ V̂j∗ − Vj∗ + Vj∗

⇐⇒ (V̂j − Vj ) − (V̂j∗ − Vj∗ ) ≤ −(Vj − Vj∗ ) Since Vj − Vj∗ ≥ ∆ for all j ̸= j∗ , we have: [ { ĵ ̸= j∗ } ⊆ (V̂j − Vj ) − (V̂j∗ − Vj∗ ) ≤ −∆ j̸= j∗

By the triangle inequality:

|(V̂j − Vj ) − (V̂j∗ − Vj∗ )| ≤ |V̂j − Vj | + |V̂j∗ − Vj∗ | ≤ 2 max |V̂k − Vk | k ∈[m]

Therefore:

∗

{ ĵ ̸= j } ⊆

∆ max |V̂k − Vk | ≥ 2 k ∈[m]

Step 4: Apply uniform convergence. By Lemma B.12, for any δ ∈ (0, 1): ( ) r log(md/δ) 2 P max |V̂j − Vj | ≥ CKY M ≤δ n j∈[m] q log(md/δ) For n sufficiently large such that CKY M2 < ∆2 , we have: n ∆ ≤δ P{ ĵ ̸= j∗ } ≤ P max |V̂k − Vk | ≥ 2 k ∈[m] Therefore, we proved that for any δ ∈ (0, 1), we have for n sufficiently large that P{ ĵ ̸= j∗ } ≤ δ, which means that P{ ĵ ̸= j∗ } → 0 as n → ∞ , p

hence ĵ − → j∗ . B.4.3

Proof of Theorem 3.2

We now prove the three claims of Theorem 3.2 using the lemmas established above. Proof of Theorem 3.2. Part (i): Consistent selection.

This follows immediately from

Lemma B.13, which establishes that ĵ − → j∗ under Assumption B.7. Part (ii): Asymptotic normality. We need to show that:

√

ĵ

n(θ̂GPPI − θ ∗ ) − → N (0, Vj∗ )

By Theorem 3.1, for each fixed transformation j, the GPPI estimator with estimated weights satisfies: √ j d n(θ̂GPPI − θ ∗ ) − → N (0, Vj ) We decompose the adaptive estimator:

√

ĵ

n(θ̂GPPI − θ ∗ ) =

√

j∗

n(θ̂GPPI − θ ∗ ) + 29

√

ĵ

j∗

n(θ̂GPPI − θ̂GPPI )

On the event { ĵ = j∗ }, the second term is zero, and we have for any ϵ > 0: n√ o j∗ ĵ P n(θ̂GPPI − θ̂GPPI ) > ϵ ≤ P{ ĵ ̸= j∗ } → 0 by Lemma B.13. Tis means that

√

By Slutsky’s theorem, we obtain √

ĵ

j∗

n(θ̂GPPI − θ̂GPPI ) = o p (1). d

ĵ

n(θ̂GPPI − θ ∗ ) − → N (0, Vj∗ )

Part (iii): Variance bound. We need to show that: r Vĵ ≤ min Vj + O p j∈[m]

log(m) n

By the definition of ĵ and j∗ , we have that ĵ minimizes the empirical variance, i.e. V̂ĵ ≤ V̂j∗ , hence V̂ĵ = Vj∗ + (V̂j∗ − Vj∗ ) + (V̂ĵ − V̂j∗ )

≤ Vj∗ + (V̂j∗ − Vj∗ ) ≤ min Vj + max |V̂j − Vj | j∈[m]

j∈[m]

By Lemma B.12, for any δ ∈ (0, 1): (

P max |V̂j − Vj | ≥ CKY M

j∈[m]

log(md/δ) n

)

≤δ

Therefore, with probability 1 − δ r V̂ĵ ≤ min Vj + CKY M j∈[m]

log(md/δ) n

This completes the proof of all three parts of Theorem 3.2. Corollary B.14. Under the assumptions of Theorem 3.2, the estimator of GPPI with greedy selection satisfies r C ′ zα/2 KY M2 log(mdn) ∗ Pr(θ ∈ Iĵ ) ≥ 1 − α − . Vj∗ n which proves asymptotic coverage. p Proof. Fix δ ∈ (0, 1) and let ε = CKY M2 log(md/δ)/n. By Lemma B.12, with probability at least 1 − δ, V̂ĵ ≥ Vĵ − ε ≥ Vj∗ − ε. q ĵ Denoting by E this event, Iĵ = θ̂GPPI ± zα/2 V̂ĵ /n and I ∗ the analogous CI with Vj∗ − ε in √ ĵ p d place of V̂ĵ , we have I ∗ ⊆ Iĵ . By Theorem 3.2(ii), Zn := n(θ̂GPPI − θ ∗ )/ Vj∗ − → N (0, 1), hence q Pr(θ ∗ ∈ Iĵ | E ) ≥ Pr(θ ∗ ∈ I ∗ | E ) = Pr | Zn | ≤ zα/2 1 − ε/Vj∗ ≥ Pr | Zn | ≤ zα/2 (1 − ε/Vj∗ ) , Denoting by Φ the standard √ normal CDF and using Berry–Esseen, assuming finite third moments, together with 1/ 2π-Lipschitz continuity of Φ, 2z Pr(θ ∗ ∈ Iĵ | E ) ≥ 2Φ zα/2 (1 − ε/Vj∗ ) − 1 − O(n−1/2 ) ≥ 1 − α − √ α/2 ε − O(n−1/2 ). 2π Vj∗ 30

Accounting for the failure event of probability at most δ, 2z Pr(θ ∗ ∈ Iĵ ) ≥ 1 − α − √ α/2 ε − O(n−1/2 ) − δ. 2π Vj∗ p √ Choosing δ = 1/ n gives ε = O( log(mdn)/n) and r C ′ zα/2 KY M2 log(mdn) ∗ Pr(θ ∈ Iĵ ) ≥ 1 − α − . Vj∗ n

B.5

Proof of Theorem 3.3

We prove that GPPI++ remains asymptotically valid when the transformation dimension d = dn grows with sample size n at an appropriate rate. In all the proof, we assume without loss of generality that cd ≤ 1 ≤ Md for all d. We demonstrate below several auxiliary lemmas before proving the theorem. B.5.1

Concentration of Sample Covariance

Lemma B.15 (Covariance Concentration). Under the assumptions of Theorem 3.3 s  Md2 log d . ∥Σ̂d − Σd ∥op = O p  N Proof. Without loss of generality, we assume that κ0 ≥ 1. For transformation ϕd : R → Rd , we have ϕd ( f ( XiU )) are i.i.d. random vectors. Define the centered random matrix: Zi = ϕd ( f ( XiU ))ϕd ( f ( XiU )) T − Σd ∈ Rd×d . Then Σ̂d − Σd = N1 ∑iN=1 Zi with E[ Zi ] = 0. By the matrix Bernstein inequality (Theorem 1.6.2 in Tropp (2015)), for independent centered self-adjoint matrices Z̃1 , . . . , Z̃N with ∥ Z̃i ∥op ≤ L almost surely:    N  t2 /2 Z̃i , P ≥ t ≤ 2d exp − 2  i∑  σ + Lt/3 =1 op

where σ2 = ∑iN=1 E[ Zi2 ]

. To use this concentration bound on ( Zi )i , we need to prove

bounds σ2 and L. Bound on σ2 : By Assumption, ∥ϕd ( f ( X ))∥2 ≤ Md . Therefore: E[ Zi2 ] = E[(ϕd ( f ( X ))ϕd ( f ( X )) T − Σd )2 ]

⪯ 2E[(ϕd ( f ( X ))ϕd ( f ( X ))T )2 ] + 2Σ2d = 2E[∥ϕd ( f ( X ))∥2 · ϕd ( f ( X ))ϕd ( f ( X ))T ] + 2Σ2d ⪯ 2Md · E[ϕd ( f ( X ))ϕd ( f ( X ))T ] + 2Σ2d = 2Md Σd + 2Σ2d ⪯ 2Md ∥Σd ∥op Id + 2∥Σd ∥2op Id . Note that ∥Σd ∥op = sup∥v∥=1 v T Σd v = sup∥v∥=1 E[(v T ϕd ( f ( X )))2 ] ≤ E[∥ϕd ( f ( X ))∥2 ] ≤ Md by assumption. Thus: E[ Zi2 ] ⪯ 2Md2 Id + 2Md2 Id = 4Md2 Id . 31

Therefore: σ2 =

∑ E[Zi2 ]

i =1

= O( N Md2 ). op

Bound on L: We need to bound ∥ Zi ∥op . By the triangle inequality:

∥ Zi ∥op = ∥ϕd ( f ( Xi ))ϕd ( f ( Xi ))T − Σd ∥op ≤ ∥ϕd ( f ( Xi ))ϕd ( f ( Xi ))T ∥op + ∥Σd ∥op . For any vector v ∈ Rd , the rank-one matrix vv T has operator norm ∥vv T ∥op = ∥v∥2 (since vv T has a single non-zero eigenvalue equal to ∥v∥2 ). Therefore:

∥ϕd ( f ( Xi ))ϕd ( f ( Xi ))T ∥op = ∥ϕd ( f ( Xi ))∥2 ≤ Md

a.s. .

Since ∥Σd ∥op ≤ Md (as shown above), we have:

∥ Zi ∥op ≤ Md + Md = 2Md . So we can take L = 2Md . Applying Bernstein: We have shown that: • E[ Zi2 ] ⪯ 4Md2 Id , so σ2 = ∑iN=1 E[ Zi2 ]

≤ 4N Md2

• ∥ Zi ∥op ≤ 2Md almost surely By the matrix Bernstein inequality, for any t > 0:   !  N  t2 /2 t2 /2 Zi ≥ t ≤ 2d exp − 2 P ≤ 2d exp −  i∑  σ + Lt/3 4N Md2 + 2Md t/3 =1 op ! t2 ≤ 2d exp − 8Md2 N + 2Md t q Set t = 24N Md2 log d. Then: 24N Md2 log d t2 q = 8Md2 N + 2Md t 8M2 N + 2M 24N M2 log d d

6 log(d) p 2 + 6 log(d)/N ≥ 2 log d ,

where the last inequality holds for d sufficiently large, as log d = o (n). Thus   q   N 2 P Zi ≥ 24N Md2 log d ≤ 2d exp (−2 log d) = → 0. ∑   i =1 d op

This gives: N

∑ Zi

i =1

= Op

Therefore:

∥Σ̂d − Σd ∥op =

N Md2 log d . s

1 Zi N i∑ =1

= Op  op

 Md2 log d . N

B.5.2

Concentration of Cross-Covariance

Lemma B.16 (Cross-Covariance Concentration). Under the theorem assumptions: ! r Md ∥σ̂Y,d − σY,d ∥ = O p . n Proof. Define Wi = Yi ϕd ( f ( XiL )) − σY,d . Since (Wi )i are iid and have null expectations, it holds that 1 E[∥σ̂Y,d − σY,d ∥2 ] = E[∥W1 ∥2 ] . (18) n Furthermore, using the variance decomposition formula then gives E[∥W1 ∥2 ] = E[∥Yϕd ( f ( X )) − σY,d ∥2 ]

= E[∥Yϕd ( f ( X ))∥2 ] − ∥σY,d ∥2 ≤ E[∥Yϕd ( f ( X ))∥2 ] = E[Y2 ∥ϕd ( f ( X ))∥2 ] ≤ E [Y 2 ] M d . The last inequality holds because ∥ϕd ( f ( X ))∥2 ≤ Md almost surely. Substituting in (18) and using that Y has bounded second moment yields E[∥σ̂Y,d − σY,d ∥2 ] ≤

E [Y 2 ] M d = O( Md /n) , n

√ and we conclude from Chebyshev inequality that ∥σ̂Y,d − σY,d ∥ = O p ( Md /n). B.5.3

Weight Estimation Error

Lemma B.17 (Weight Error). Under the assumptions of Theorem 3.3 ! p Md3/2 log(d) ∗ √ , ∥λ̂d − λd ∥ ≤ O p c2d n Proof. Using Weyl’s inequality for eigenvalues then Lemma B.15 yields s  Md2 log d  . |ρmin (Σ̂d ) − ρmin (Σd )| ≤ ∥Σ̂d − Σd ∥op = O p  N Moreover, using the assumptions ρmin (Σd ) ≥ cd and

Md2 log d/N = o (cd ) we deduce that

ρmin (Σ̂d ) ≥ cd − |ρmin (Σ̂d ) − ρmin (Σd )|

≥ cd − o p (cd ) = Ω p (cd ) where the last inequality holds for n sufficiently large. It follows that 1 ∥Σ̂− d ∥op =

1 = O p (1/cd ) . ρmin (Σ̂d )

Decompose: 1 1 Σ̂−1 σ̂ − Σ −1 σ 1 + r d Y,d 1 + r d Y,d 1 1 1 1 = Σ̂− (σ̂Y,d − σY,d ) + (Σ̂−1 − Σ− d d ) σY,d . 1+r 1+r d

λ̂d − λ∗d =

(19)

1 First term: By Lemma B.16 and ∥Σ̂− d ∥op = O p (1/cd ):

1 Σ̂−1 (σ̂Y,d − σY,d ) ≤ O p (1/cd ) · O p 1+r d

Md n

√

= Op

M √d cd n

(20)

1 −1 −1 −1 Second term: Using the identity Σ̂− d − Σd = − Σ̂d ( Σ̂d − Σd ) Σd , we obtain

1 1 1 −1 −1 (Σ̂−1 − Σ− d ) σY,d ≤ 1 + r · ∥ Σ̂d ∥op · ∥ Σd ∥op · ∥ Σ̂d − Σd ∥op · ∥ σY,d ∥ 1+r d s  Md2 log d  · ∥σY,d ∥. ≤ O p (1/cd ) · 1/cd · O p  N The term ∥σY,d ∥ can be bounded using Jensen’s inequality and assumption that ϕd ( f ( X )) is bounded a.s.:

∥σY,d ∥2 = ∥E[Yϕd ( f ( X ))]∥2 ≤ E[∥Yϕd ( f ( X ))∥2 ] = E[Y2 ∥ϕd ( f ( X ))∥2 ] ≤ E[Y2 ] Md , √ thus ∥σY,d ∥ = O( Md ). It follows that  q Md3 log d 1 1  .  √ (21) (Σ̂−1 − Σ− d ) σY,d = O p 1+r d c2d n Finally, combining (19), (20) and (21) we deduce that  q √ Md3 log d M  = Op √ d + Op  √ ∥λ̂d − λ∗d ∥ ≤ O p cd n c2d n

√

! p Md log d Md /n , 1+ cd cd

and with the assumption that cd ≤ 1 ≤ Md we deduce that ! p 3/2 M log ( d ) d √ ∥λ̂d − λ∗d ∥ ≤ O p , c2d n

B.5.4

Feature Mean Concentration

Lemma B.18 (Feature Mean Concentration). Under the theorem assumptions, let µd = E[ϕd ( f ( X ))]. Then: ! r M d U ∥ϕ̄d,N − µd ∥ = O p , N ! r Md L ∥ϕ̄d,n − µd ∥ = O p . n N 1 n U = 1 U L L Proof. Since ϕ̄d,N N ∑i =1 ϕd ( f ( Xi )) and ϕ̄d,n = n ∑i =1 ϕd ( f ( Xi )) are sample means of i.i.d. random vectors, the proof follows the same concentration argument as Lemma B.16. √ Using the assumption ∥ϕd ( f ( X ))∥ ≤ Md and applying again matrix Bernstein inequality yields the stated bounds.

Lemma B.19 (Second-Order Bias Vanishes). Under the assumptions of Theorem 3.3 √ n(θ̂ (λ̂d ) − θ̂ (λ∗d )) = o p (1) . 34

U − ϕ̄ L . By definition: Proof. Define ∆λ = λ̂d − λ∗d and ∆ϕ = ϕ̄d,N d,n √ √ n(θ̂ (λ̂d ) − θ̂ (λ∗d )) = n∆λT ∆ϕ .

Let µd = E[ϕd ( f ( X ))]. Decompose: U L ∆ϕ = (ϕ̄d,N − µd ) + (µd − ϕ̄d,n ).

Using Lemmas B.18 and B.17, and that N = Ω(n) we obtain √ √ n|∆λT ∆ϕ | ≤ n∥∆λ ∥ · ∥∆ϕ ∥ √ U L ≤ n∥∆λ ∥ · ∥ϕ̄d,N − µd ∥ + ∥ϕ̄d,n − µd ∥ ! p r √ Md3/2 log(d) Md √ n· · = Op n c2d n ! p Md2 log d √ = o p (1) . = Op c2d n −1 1 T Lemma B.20 (Variance Estimator Consistency). Let V̂ = σ̂Y2 − 1+ r σ̂Y,d Σ̂d σ̂Y,d be the plug-in variance estimator. Under the assumptions of Theorem 3.3, it holds that p

V̂ − → V∞ . Proof. Recall that V̂ = σ̂Y2 −

1 σ̂ T Σ̂−1 σ̂ , 1 + r Y,d d Y,d

hence we have that V̂ − V∞ = V̂ − Vd∗ + Vd∗ − V∞ 1 T −1 T 1 ∗ = σ̂Y2 − σY2 − σ̂Y,d Σ̂d σ̂Y,d − σY,d Σ− σ Y,d + (Vd − V∞ ) . d 1+r

(22)

−1 1 1 n 2 T 2 We have by assumption that Vd∗ = σY2 − 1+ r σY,d Σd σY,d → V∞ , and since σ̂Y = n ∑i =1 Yi and E[Y 2 ] < ∞, by the strong law of large numbers: p

σ̂Y2 − σY2 − →0. p −1 T Σ̂−1 σ̂ T Therefore, it remains to prove that σ̂Y,d → 0. Let us decompose Y,d − σY,d Σd σY,d − d this as two terms: T 1 T −1 σ̂Y,d Σ̂− d σ̂Y,d − σY,d Σd σY,d T 1 T −1 T −1 T −1 ≤ σ̂Y,d Σ̂− d σ̂Y,d − σY,d Σ̂d σY,d + σY,d Σ̂d σY,d − σY,d Σd σY,d .

Subterm 1.

(23)

The first term can be upper bounded as T −1 T 1 T −1 σ̂Y,d Σ̂− d σ̂Y,d − σY,d Σ̂d σY,d = ( σ̂Y,d − σY,d ) Σ̂d ( σ̂Y,d + σY,d )

1 ≤ ∥σ̂Y,d − σY,d ∥∥Σ̂− (24) d ∥op ∥ σ̂Y,d + σY,d ∥. √ By Lemma B.16: ∥σ̂Y,d − σY,d ∥ = O p ( Md /n). Also, from the proof of Lemma B.17: 1 ∥Σ̂− d ∥op = O p (1/cd ). And the norm of the sum satisfies p p p ∥σ̂Y,d + σY,d ∥ ≤ ∥σ̂Y,d ∥ + ∥σY,d ∥ = O p ( Md ) + O( Md ) = O p ( Md ) .

Therefore:

]

T 1 T −1 σ̂Y,d Σ̂− d σ̂Y,d − σY,d Σ̂d σY,d

= Op

Md 1 p · · Md n cd

= Op

Md √ cd n

(25)

Subterm 2: For the second term in (23), we have T −1 −1 T 1 T −1 σY,d Σ̂− d σY,d − σY,d Σd σY,d = σY,d ( Σ̂d − Σd ) σY,d

−1 1 ≤ ∥σY,d ∥2 ∥Σ̂− d − Σd ∥op .

(26)

1 −1 −1 −1 Using the identity Σ̂− d − Σd = − Σ̂d ( Σ̂d − Σd ) Σd then Lemma B.15 yields −1 −1 −1 1 ∥Σ̂− d − Σd ∥op ≤ ∥ Σ̂d ∥op ∥ Σ̂d − Σd ∥op ∥ Σd ∥op

s

= O p (1/cd ) · O p  = Op

log d √ 2 cd N



Md2 log d  N

· O(1/cd )

! ,

Since ∥σY,d ∥2 = O( Md ) (as shown in Lemma B.17), we have T 1 T −1 σY,d Σ̂− d σY,d − σY,d Σd σY,d

= Op

Md2

log d √ 2 cd N

! .

(27)

Substituting (25) and (27) into (23), we obtain: T 1 T −1 σ̂Y,d Σ̂− d σ̂Y,d − σY,d Σd σY,d

= Op = Op

Md √ + Op cd n ! p Md2 log d √ c2d N

where we used that cd ≤ 1 ≤ Md hence c M√dn ≤ c M√dn · Md /cd · d

Md2

log d √ 2 cd N

log d =

Md2

√

c2d

log d √ . MoreN

over, the assumption of Theorem 3.3 guarantees that this term vanishes, hence 1 T 1 σ̂ T Σ̂−1 σ̂ − σY,d Σ− d σY,d = o p (1) , 1 + r Y,d d Y,d which concludes the proof. B.5.5

Main Proof

Proof of Theorem 3.3. The proof of the theorem follows immediately from the previous lemmas. Asymptotic normality By assumption (v), Vd∗ → V∞ . As in the proof of Theorem 3.1, for the oracle estimator with known optimal weights: √ d n(θ̂ (λ∗d ) − θ ∗ ) − → N (0, V∞ ) . On the other hand, we can decompose √ √ √ n(θ̂GPPI − θ ∗ ) = n(θ̂ (λ∗d ) − θ ∗ ) + n(θ̂ (λ̂d ) − θ̂ (λ∗d )). By Lemma B.19, the second term is o p (1). Therefore, we deduce By Slutsky’s theorem that √ d n(θ̂GPPI − θ ∗ ) − → N (0, V∞ ). p

Variance consistency: By Lemma B.20, V̂ − → V∞ .

Experimental Details

This appendix provides comprehensive details on the experimental setup, including dataset statistics, persona construction procedures, prompt templates, and implementation specifics. C.1

Persona Construction

Jester Jokes personas. For each user, we construct a persona from their ratings on the training jokes. The persona includes: • Rating statistics: Total number of ratings, mean rating, standard deviation, min, max, and median. • Rating distribution: Counts and percentages in four buckets: very negative ([−10, −5)), negative ([−5, 0)), positive ([0, 5)), very positive ([5, 10]). • Favorite jokes: Up to 15 jokes rated above 7, sorted by rating, with full joke text. • Least favorite jokes: Up to 10 jokes rated below −5, sorted by rating, with full joke text. MovieLens personas. For each user, we construct a persona from their complete rating history on training movies. The persona includes: • Rating statistics: Total ratings, mean, standard deviation, and rating distribution across the 10 possible values. • Temporal evolution: If the user’s rating history spans more than 6 months, we divide it into 3 equal time periods and compute the mean rating in each period to capture temporal trends. • Genre preferences: For each of the 20 genres, we compute the average rating and number of movies rated. We report the top 5 genres by average rating. • Favorite movies: Up to 10 movies rated 5.0, with title and genres. • Recent history: The last 100 movies rated (chronologically), with date, title, genres, and rating. C.2

Prompt Templates

We provide below a representative prompt template for MovieLens. All prompts for he other datasets follow a similar structure, including the available data from each dataset. You are impersonating a movie viewer. Based on YOUR complete viewing history below, predict how YOU would rate the following movie on a scale of 0.5-5.0 (in 0.5 increments). YOUR VIEWING STATISTICS: - Total movies rated: {total_ratings} - Overall average rating: {mean_rating:.2f} - Evolution of average rating with time: {temporal_evolution} - Top genres: {top_genres_with_avg_ratings} - Rating distribution: {rating_distribution} - Active period: {first_date} to {last_date} SAMPLE OF YOUR FAVORITE MOVIES (rated 5.0): {favorite_movies_with_genres} YOUR RECENT RATING HISTORY (last 100 movies, chronological): {recent_history_with_dates_titles_genres_ratings} MOVIE TO RATE: Title: {movie_title}

Genres: {movie_genres} INSTRUCTIONS: 1. Consider YOUR genre preferences and how this movie's genres align with YOUR tastes. 2. Reflect on YOUR rating patterns over time and YOUR typical rating level. 3. Think about similar movies YOU have rated and how this compares. Provide ONLY a numerical rating in {0.5, 1.0, 1.5, ..., 5.0}. No explanations. C.3

Model Configuration

LLM models.

We use three Mistral models:

• Mistral-7B-Instruct-v0.2: 7 billion parameters • Mistral-24B-Instruct: 24 billion parameters • Mistral-123B-Instruct: 123 billion parameters These models were selected for their open availability and performance on simple tasks such as the ones we test. All models use temperature 0.0 and the same prompt format with no model-specific adjustments. For inference, we used Amazon bedrock. The cost of all the inferences is less than 2000 USD. The rest of the experiments consisted in sampling and doing parameter estimation / hypothesis testing. They did not incur any further costs. C.4

Population-level predictions: Synthetic experiment.

To better illustrate strenghts and weaknesses of the population level prediction framework, we do a synthetic experiments visualizing the different modes. We generate n = 100 samples from N (µ0 + δ, 1) for effect sizes δ ∈ {0.05, 0.1, 0.2, 0.3, 0.4}, testing H0 : µ = µ0 at α = 0.05 with prediction accuracy p ∈ [0.5, 1] and asymmetry λ ∈ [0, 0.99]. Figure 3 shows moderate λ ∈ [0.25, 0.75] achieve substantial gains with accurate predictions while maintaining robustness when predictions fail, validating Theorem 2.1. = 0.1

0.25

0.08

0.20

60%

80%

100%

Classical ( =0)

0.5 0.4

0.15

Prediction Accuracy

Power

0.10

60%

80%

100%

Prediction Accuracy

=0.25

=0.50

=0.75

60%

80%

100%

Prediction Accuracy =0.99

= 0.3

0.9

0.6

Power

0.12

Learning-Augmented Z-test

= 0.2

0.8 0.7 0.6

= 0.4 Power

Power vs Prediction Accuracy

= 0.05

60%

80%

0.90

0.85 100%

Prediction Accuracy

Classical baseline

0.95

Oracle (correct)

60%

80%

100%

Prediction Accuracy Oracle (wrong)

Figure 3: Synthetic directional predictions. Power vs prediction accuracy for various δ and λ. Moderate λ balances gains with robustness. C.5

Item-Level GPPI Performance Analysis

We provide in this section a comprehensive item-level analysis examining how GPPI performs relative to PPI++ on the task of mean estimation for each test individual movie, joke, book and song from the datasets we used. In all the experiments, we construct confidence intervals with type I error α = 5%. For each item and each model (Mistral-7B, 24B, 123B), we sample n labeled samples and N unlabeled samples, where N = 5n and n is chosen maximally, i.e. every user is used in 38

the experiment, either as labeled or unlabled sample. All the values in these experiments are computed over 1000 trials. Note that the predictions with the LLM models are only computed once, the source of randomness in the experiment is in splitting the users into labeled and unlabeled. For each dataset and for each item we compute • PPI++ ESS gain: The baseline improvement over the classical estimator using linear correction. • GPPI gap: The difference between the ESS gain of GPPI with greedy transform selection (using the same transformations mentioned in the experiments section) and PPI++’s ESS gain. A positive gap indicates GPPI improves over PPI++; negative indicates degradation. We visualize these metrics as heatmaps where rows represent models and columns represent items. Items are labeled with their ID and available sample size n. We also visualize heatmaps for the coverage of each method. For each dataset we show two figures next to each other. In the right figure, the top panel shows PPI++ ESS gain (%) for each of the 15 movies (columns) and three models (rows). The lower three panels: GPPI with greedy selection gap (GPPI ESS gain minus PPI++ ESS gain) without penalization, with AIC penalization, then with BIC penalization. The right figure has one additional top panel for the classical estimator. Each panel panels corresponds to a mean estimator (classical/PPI++/GPPI greedy selection with different penalizations), and shows the coverage obtained with that estimator for different items and models, i.e. percentage of times when the true mean is in the CI constructed by that algorithm. CI Coverage

Movie Lens

10%

-5%

-10%

-15%

75 9

15%

77 6 17

76 1 17

53 0 17

40 5 17

17 6 17

82 5 16

66 3 16

65 3 16

65 2 16

29 2

29 0 12

18 75

17 77

17 76

17 53

17 40

17 17

16 82

16 66

16 65

16 64

15 82

15 20

12 29

80%

82%

64 6

85%

88%

90%

82 3

92%

95%

98%

20 8

100%

PPI++ ESS Gain vs Classical 7B params 8.4% 8.5% 9.9% 5.9% 10.4% 6.6% 3.6% 9.3% 4.1% 13.3% 3.8% 8.4% 7.7% 5.6% 10.3% 24B params 7.4% 12.1% 17.7% 12.0% 16.8% 11.3% 14.7% 26.9% 6.0% 5.4% 9.2% 13.8% 12.2% 10.5% 8.7% 123B params 19.5% 31.3% 40.4% 21.2% 34.5% 28.4% 21.4% 31.0% 15.2% 20.3% 15.4% 21.4% 25.3% 31.3% 20.7% Greedy PPI++ 7B params 1.5% 2.1% 1.6% 1.5% 1.2% 2.1% 1.2% 1.5% 1.3% 2.0% 0.9% 1.7% 1.9% 1.8% 1.8% 24B params 1.5% 1.8% 0.9% 0.9% 1.6% 1.9% 1.2% 1.6% 1.6% 1.0% 1.2% 0.8% 0.8% 0.6% 1.2% 123B params 2.1% 2.1% 1.2% 1.6% 1.2% 1.6% 1.5% 1.4% 1.4% 1.7% 2.4% 3.0% 1.8% 0.7% 1.1% Greedy AIC PPI++ 7B params 1.6% 1.0% 1.7% 1.8% 1.0% 2.0% 1.4% 1.6% 1.4% 1.3% 1.9% 1.6% 1.5% 1.7% 1.8% 24B params 2.1% 1.4% 1.7% 1.7% 0.8% 1.5% 1.5% 1.4% 1.6% 1.1% 1.1% 1.4% 1.2% 1.3% 1.0% 123B params 1.4% 1.9% 1.5% 2.0% 1.7% 1.4% 1.6% 1.9% 2.0% 2.6% 2.0% 3.2% 1.8% 2.7% 1.1% Greedy BIC PPI++ 7B params 1.4% 1.3% 1.7% 2.0% 2.0% 1.6% 1.4% 2.4% 1.2% 1.1% 1.5% 1.8% 1.4% 1.8% 1.5% 1.4% 1.5% 1.1% 2.1% 1.6% 0.9% 1.5% 2.5% 1.5% 2.1% 1.9% 1.1% 1.6% 0.9% 1.5% 24B params 123B params 1.6% 1.6% 1.6% 2.1% 1.8% 1.2% 1.6% 1.2% 2.4% 1.9% 1.4% 2.4% 1.8% 1.9% 2.2% 15

Movie Lens

('classical', 0) Coverage 7B params 96.8% 96.9% 96.8% 96.8% 96.9% 96.8% 96.7% 96.8% 96.6% 96.9% 96.8% 97.0% 96.5% 96.8% 96.6% 24B params 96.8% 96.8% 96.8% 97.1% 96.9% 96.8% 96.9% 96.6% 96.9% 96.7% 96.6% 96.7% 96.8% 97.0% 96.9% 123B params 96.5% 96.6% 96.7% 96.8% 96.8% 96.7% 96.9% 96.6% 96.9% 96.7% 96.9% 96.8% 96.6% 96.5% 96.6% PPI++ Coverage 7B params 96.8% 97.0% 97.0% 96.9% 97.0% 96.8% 96.8% 97.0% 96.6% 97.0% 96.9% 97.1% 96.5% 96.8% 96.8% 24B params 97.0% 96.8% 97.2% 97.1% 97.0% 97.1% 97.2% 96.9% 96.9% 96.8% 96.7% 96.8% 96.7% 97.1% 96.9% 123B params 96.8% 97.0% 97.2% 97.0% 97.1% 97.2% 97.2% 96.8% 97.0% 96.9% 96.9% 97.1% 97.0% 96.9% 97.0% Greedy Coverage 7B params 96.8% 96.9% 96.9% 96.8% 96.9% 96.7% 96.6% 96.9% 96.6% 97.0% 96.8% 96.9% 96.5% 96.7% 96.6% 24B params 97.0% 96.7% 97.1% 97.0% 96.8% 97.0% 97.1% 96.9% 96.8% 96.7% 96.6% 96.6% 96.6% 96.9% 96.9% 123B params 96.7% 96.9% 97.1% 96.9% 97.0% 97.2% 97.1% 96.8% 96.9% 96.9% 96.8% 96.9% 96.7% 96.7% 96.9% Greedy AIC Coverage 7B params 96.8% 97.0% 96.9% 96.9% 97.0% 96.8% 96.8% 97.0% 96.6% 97.1% 96.8% 97.0% 96.6% 96.8% 96.8% 24B params 97.0% 96.7% 97.2% 97.1% 96.9% 97.0% 97.2% 97.0% 96.9% 96.8% 96.7% 96.8% 96.7% 97.0% 96.9% 123B params 96.7% 97.0% 97.2% 97.0% 97.0% 97.2% 97.2% 96.9% 96.9% 96.9% 96.9% 97.0% 97.0% 96.8% 97.0% Greedy BIC Coverage 7B params 96.9% 97.1% 97.0% 97.0% 97.1% 96.8% 96.9% 97.1% 96.7% 97.1% 96.9% 97.1% 96.7% 96.9% 96.8% 24B params 97.0% 96.8% 97.3% 97.1% 97.0% 97.1% 97.3% 97.1% 96.9% 96.8% 96.8% 96.9% 96.8% 97.1% 97.0% 123B params 96.8% 97.1% 97.3% 97.1% 97.1% 97.2% 97.2% 97.0% 97.0% 97.0% 97.0% 97.2% 97.1% 96.9% 97.1%

Figure 4: Item-level performance for MovieLens. CI Coverage

2.3%

6.5%

0.9%

5.5%

3.6%

0.1%

4.5%

7.3%

5.7%

24B params 20.1% 17.6% 10.2%

9.0%

17.1%

17.2%

14.5%

11.3%

12.7%

17.3%

5.4%

14.4%

14.1%

14.2%

11.0%

15%

123B params 31.1% 28.1% 18.4% 23.2% 29.3% 26.2% 23.5% 19.2% 22.8% 28.8% 17.9% 23.1% 26.3% 29.0% 27.6% Greedy PPI++ 7B params 2.4% 5.5% 3.6% 2.4% 3.8% 2.9% 4.3% 2.1% 2.5% 2.6% 1.3% 1.1% 3.6% 2.6% 2.3%

10%

24B params 5.3%

5.9%

4.8%

3.7%

5.5%

5.4%

4.1%

3.9%

4.0%

2.8%

2.6%

4.5%

3.8%

2.7%

123B params 6.7%

4.4%

3.6%

3.3%

4.2%

4.9%

1.8%

2.6%

4.4%

4.2%

3.8%

3.5%

3.9%

4.0%

2.4%

7B params 2.0%

4.8%

3.8%

2.4%

5.0%

3.9%

3.3%

1.7%

1.6%

1.7%

1.6%

3.8%

2.8%

2.5%

24B params

5.2%

5.7%

4.8%

3.2%

6.1%

4.4%

3.8%

4.2%

3.9%

2.7%

2.0%

2.5%

4.6%

3.5%

2.6%

123B params 6.2%

4.2%

3.2%

3.8%

4.3%

4.6%

2.6%

4.5%

3.7%

2.7%

2.9%

3.5%

3.8%

2.6%

Greedy AIC PPI++

90%

0.9%

24B params 4.9%

5.5%

4.6%

2.6%

4.9%

3.4%

4.4%

3.8%

3.1%

2.3%

2.7%

3.0%

3.1%

3.6%

2.2%

123B params 6.2%

3.9%

2.1%

3.2%

4.1%

4.0%

1.2%

2.8%

3.6%

3.4%

3.6%

1.6%

2.5%

2.9%

1.6%

-15%

1.4%

3.6%

1.6%

1.0%

1.3%

2.2%

1.8%

3.8%

2.3%

4.6%

1.7%

3.1%

2.7%

80%

-5%

-10%

7B params 1.8%

82%

Greedy BIC PPI++

85%

88%

92%

4.6%

95%

3.7%

98%

7.4%

100%

PPI++ ESS Gain vs Classical 4.6%

8 14

3 14

4 13

7 12

9 11

7 11

8 10

Jester Jokes

7B params 9.9%

Jester Jokes

('classical', 0) Coverage 7B params 96.9% 96.8% 96.8% 96.7% 96.8% 96.7% 97.0% 96.9% 96.7% 96.8% 96.9% 96.9% 96.7% 96.8% 97.0% 24B params 96.5% 96.8% 96.8% 96.7% 96.9% 97.0% 96.8% 96.8% 96.9% 96.8% 96.6% 96.9% 96.7% 96.8% 96.8% 123B params 96.6% 96.8% 96.7% 96.8% 96.8% 96.6% 97.0% 96.9% 96.7% 96.9% 97.0% 96.9% 96.9% 97.0% 96.9% PPI++ Coverage 7B params 97.0% 96.9% 97.0% 96.9% 96.8% 96.7% 97.0% 96.9% 96.9% 96.8% 96.8% 97.0% 96.8% 96.8% 96.9% 24B params 96.9% 97.2% 97.0% 96.8% 97.1% 97.0% 97.0% 97.0% 97.1% 97.0% 96.7% 97.0% 97.0% 96.9% 97.0% 123B params 97.3% 97.1% 97.1% 97.1% 97.3% 96.9% 97.2% 97.2% 97.0% 97.3% 97.2% 97.0% 97.2% 97.3% 97.3% Greedy Coverage 7B params 97.0% 96.9% 96.9% 96.8% 96.9% 96.8% 97.0% 96.9% 96.8% 96.7% 96.7% 96.9% 96.7% 96.9% 96.7% 24B params 97.0% 97.2% 97.0% 96.8% 97.1% 97.0% 97.0% 97.0% 97.1% 97.0% 96.7% 97.0% 97.0% 97.0% 97.0% 123B params 97.3% 97.1% 97.0% 97.1% 97.4% 96.9% 97.2% 97.2% 97.1% 97.2% 97.1% 97.0% 97.2% 97.3% 97.3% Greedy AIC Coverage 7B params 97.0% 97.0% 97.0% 96.9% 96.9% 96.8% 97.0% 96.9% 96.9% 96.8% 96.8% 97.0% 96.7% 96.9% 96.8% 24B params 97.1% 97.2% 97.1% 96.9% 97.2% 97.0% 97.1% 97.0% 97.1% 97.0% 96.8% 97.0% 97.1% 97.0% 97.0% 123B params 97.3% 97.1% 97.0% 97.1% 97.4% 96.9% 97.2% 97.2% 97.1% 97.3% 97.1% 97.0% 97.2% 97.3% 97.4% Greedy BIC Coverage 7B params 97.1% 97.0% 97.1% 96.9% 97.0% 96.8% 97.0% 97.0% 96.9% 96.9% 96.9% 97.1% 96.8% 96.9% 96.9% 24B params 97.1% 97.2% 97.2% 96.9% 97.2% 97.0% 97.1% 97.1% 97.2% 97.0% 96.8% 97.1% 97.2% 97.0% 97.1% 123B params 97.3% 97.2% 97.1% 97.2% 97.4% 97.0% 97.3% 97.3% 97.2% 97.3% 97.2% 97.1% 97.3% 97.3% 97.3%

Figure 5: Item-level performance for Jester jokes. Across the four datasets, GPPI with greedy transformation selection outperforms PPI++ in 98% of the tested items. Improvements are mostly modest on the Movielens dataset, which suggests that linear correction of the predictions is already sifficient. On the Jester and Book crossing datasets the improvements are again moderate for some items but more important for others. 39

CI Coverage

Book Crossing

3.3%

-0.4%

-8.7%

0.6%

9.0%

-5.4%

16.5%

8.0%

-25.5% 14.1%

-1.3%

-0.2%

2.5%

5.9%

-2.7%

-12.4%

2.8%

-2.2%

13.4%

-4.6%

6.3%

7.2%

6.4%

13.0%

2.4%

-5.7%

1.2%

1.1%

2.0%

-6.5%

0.2%

-5.5%

-6.3%

0.3%

4.2%

14.3%

17.7%

4.7%

3.0%

7B params 4.3% 24B params 1.1% 123B params 2.2%

4.3%

1.0%

-0.3%

2.5%

1.3%

3.2%

0.7%

-0.2%

5.3%

1.3%

0.1%

1.3%

5.2%

7B params 4.3% 24B params 1.3% 123B params 1.2%

2.2%

0.9%

-0.5%

1.7%

2.0%

4.3%

1.6%

1.8%

1.7%

-0.2%

5.0%

-5.9%

2.5%

0.6%

0.5%

1.8%

4.5%

0.6%

1.8%

7B params 2.2% 24B params 1.3% 123B params 1.6%

1.1%

1.7%

-0.8%

-0.2%

1.3%

4.1%

1.4%

1.2%

0.5%

4.2%

-1.3%

2.5%

1.7%

1.2%

2.0%

3.0%

5.0%

1.1%

2.1%

2.7%

8.2%

-0.2%

0.8%

8.3%

3.2%

0.8%

-7.2%

0.4%

0.2%

1.5%

1.0%

1.1%

1.5%

2.4%

1.3%

0.2%

-0.3%

0.1%

8.3%

0.7%

3.1%

4.7%

1.0%

2.9%

0.9%

1.8%

7.4%

2.0%

1.7%

8.8%

2.4%

1.0%

1.9%

2.5%

1.1%

1.8%

1.6%

1.2%

1.8%

1.6%

7.9%

0.9%

2.4%

2.5%

0.8%

3.3%

1.1%

1.7%

8.5%

2.6%

1.6%

8.6%

2.8%

1.6%

1.9%

1.1%

1.8%

1.9%

1.5%

2.1%

7.7%

1.6%

2.5%

1.7%

0.2%

2.7%

1.3%

-5%

11 04 52 28 21 52 06 79 76 40 2X 07 43 41 81 74 09 71 88 01 07

-15%

04 46 67 22

-10%

04 46 31 07

03 75 72 73

03 16 76 94

01 56 02 73

01 42 00 17

80%

5.1%

Greedy BIC PPI++

03 16 66 63

82%

03 16 60 19

85%

81 41

71 09

1.4%

Greedy AIC PPI++

88%

15%

10%

6.0%

90%

43 07

76 06

22 46

07 31 46

2X 72 23

42 50

85 03

75 03

16 03

60 16

12 03

02 56 01

00 42 01

Greedy PPI++

92%

9.2%

95%

5.9%

04 40 23 72

98%

2.8%

03 85 50 42

100%

PPI++ ESS Gain vs Classical

7B params 2.1% 24B params 2.9% 123B params 2.3%

03 12 19 55

Book Crossing

('classical', 0) Coverage 7B params 94.8% 94.5% 94.7% 94.8% 95.0% 94.7% 94.6% 94.6% 94.9% 94.3% 94.7% 94.8% 94.4% 94.8% 94.8% 24B params 94.6% 94.5% 94.6% 94.7% 94.9% 94.7% 94.6% 94.7% 94.5% 94.1% 94.9% 94.7% 94.3% 94.9% 94.9% 123B params 94.7% 94.5% 94.6% 95.0% 95.3% 94.3% 94.6% 94.7% 94.7% 94.3% 94.9% 94.6% 94.4% 94.9% 94.9% PPI++ Coverage 7B params 94.7% 94.3% 94.0% 94.6% 94.9% 94.4% 90.2% 93.7% 95.0% 93.9% 92.9% 93.8% 84.2% 94.5% 93.9% 24B params 93.6% 94.5% 94.0% 90.2% 93.5% 90.1% 94.1% 93.5% 94.3% 92.9% 92.7% 94.6% 94.3% 92.9% 94.9% 123B params 94.3% 93.0% 94.4% 91.6% 95.4% 90.9% 94.3% 92.0% 84.1% 94.5% 94.2% 94.4% 93.8% 88.0% 94.8% Greedy Coverage 7B params 94.0% 93.3% 93.6% 93.4% 94.7% 94.2% 91.0% 93.7% 94.4% 94.5% 92.8% 92.7% 86.1% 94.5% 93.8% 24B params 93.3% 94.4% 93.6% 89.0% 93.4% 86.2% 93.9% 92.9% 94.1% 92.4% 92.7% 94.3% 94.0% 91.7% 94.8% 123B params 94.3% 92.5% 94.1% 92.0% 95.2% 90.6% 93.9% 91.3% 88.3% 94.4% 93.9% 94.3% 93.8% 87.4% 94.8% Greedy AIC Coverage 7B params 94.7% 94.5% 94.1% 94.0% 94.8% 94.7% 91.3% 94.0% 94.7% 94.8% 93.5% 93.9% 86.9% 94.8% 94.1% 24B params 93.9% 94.6% 94.2% 90.3% 94.0% 88.4% 94.5% 93.8% 94.4% 93.1% 93.2% 94.8% 94.5% 93.1% 95.0% 123B params 94.5% 93.1% 94.5% 92.8% 95.3% 91.2% 94.5% 92.2% 88.8% 94.7% 94.3% 94.6% 94.1% 89.2% 94.9% Greedy BIC Coverage 7B params 95.0% 94.7% 94.3% 94.6% 95.1% 95.0% 91.6% 94.3% 95.3% 95.0% 93.8% 94.3% 87.2% 95.1% 94.4% 24B params 94.2% 94.9% 94.5% 91.2% 94.4% 90.2% 94.7% 94.1% 94.7% 93.5% 93.7% 95.2% 94.8% 93.5% 95.2% 123B params 94.7% 93.5% 94.8% 93.2% 95.6% 91.6% 94.8% 92.9% 89.5% 94.9% 94.7% 94.9% 94.3% 89.9% 95.1%

Figure 6: Item-level performance for Book Crossing. CI Coverage

Last.fm

6.7%

-2.1%

2.3%

1.1%

-2.8%

8.8%

1.0%

1.3%

-1.6%

7.8%

-9.0%

-79.7% 14.5%

-0.0%

12.4%

1.0%

0.1%

-48.0%

-4.1%

-3.0%

-5.4%

20.2%

2.1%

5.1%

0.1%

-80.8% 10.4% -30.3%

2.8%

5.0%

-0.5%

4.7%

-0.5%

-5.6%

-10.6%

5.0%

-8.0%

Greedy PPI++ -1.5%

12.6%

2.6%

-0.1%

2.2%

12.2%

1.3%

-1.7%

1.4%

5.9%

12.9%

5.7%

90.4%

1.4%

3.6%

4.6%

4.9%

36.9%

-3.0%

-2.4%

-4.4%

-4.3%

6.1%

1.7%

97.6%

0.5%

13.0%

3.2%

2.6%

3.6%

8.5%

3.7%

2.4%

31.1%

18.0%

12.9%

7B params 4.9% 0.7% 24B params 4.8% -10.4% 123B params -10.4% -3.4%

4.4%

-2.0%

13.2%

1.9%

0.9%

1.6%

12.0%

0.3%

-1.1%

1.2%

1.5%

3.5%

11.3%

2.5%

90.9%

1.6%

4.8%

2.8%

4.3%

3.4%

36.8%

-1.7%

0.9%

-1.6%

2.3%

6.9%

1.3%

95.6%

0.8%

16.0%

1.6%

2.6%

4.0%

6.0%

3.5%

1.5%

35.8%

15.2%

12.9%

7B params 2.8% 0.8% 24B params 5.3% 0.4% 123B params -11.0% -1.8%

3.8%

-2.6%

16.2%

1.1%

0.9%

8.8%

0.6%

0.4%

1.2%

2.5%

5.6%

0.5%

90.4%

1.5%

4.8%

2.1%

5.0%

1.6%

36.7%

0.9%

3.7%

-1.0%

4.3%

7.9%

1.2%

91.7%

1.4%

19.0%

1.5%

2.0%

2.2%

10.2%

2.5%

1.8%

25.2%

12.8%

14.6%

Greedy AIC PPI++

-5%

F 3B

28 F9

UF TB Y1 TR

W1 SW

-15%

28 F9 TR 34 WA 50 QO B8 C1 29 03 TR CB XW 84 AZ CA C1 28 F9 31 4B 3E

F9 2

03 C

29 K1 PD

34 6

-10%

F7 1

D1 C9

YY L1 PF TR

Y1 2

8F 92

93 20

9 C0

4 4B 92 DE

KC G OM

28 F

03 CB

8B 0C E0 78

H1 29 HL

GC TR

F4 2

28 OJ1

FW G TR

TR CR CB

T1 28

1 C6 2E 12

Greedy BIC PPI++

8F 92

F4 2 X1 28

H1 LA

BN TR

TR AA

28 E0 78

80%

XID

82%

85%

15%

10%

4.6%

90%

88%

-3.2%

7B params 4.6% 0.9% 24B params 4.0% -10.0% 123B params -9.2% -1.8%

34 F9 28

K1 PD

UF 63 TB BF Y1 28 F9 TR 34 WA 50 QO B8 C1 29 03 TR CB XW 84 AZ CA C1 28 F9 31 4B 3E

1 65

71 2F

CC C

03 29

28 YY L1

PF TR

9D 2C

F9 28 Y1 YH

09 0C

B4 28

KC G1

2D E4 F9

03 29

B C8

8E J12

GO FW TR

D5 92

8F 12

TR CR CB T

4A 23

22 F4

XID TR CP

BN TR

TR AA

21.4%

92%

-9.9%

95%

1.2%

98%

PPI++ ESS Gain vs Classical

7B params 4.5% 24B params 22.0% 123B params 27.2%

12 8F

100%

TR CP

Last.fm

('classical', 0) Coverage 7B params 94.2% 92.9% 94.6% 94.8% 94.5% 94.8% 94.0% 94.7% 94.7% 94.5% 93.9% 94.3% 93.2% 94.5% 94.1% 24B params 94.4% 93.0% 94.7% 94.6% 94.3% 93.8% 94.3% 95.1% 94.9% 94.7% 93.1% 94.5% 93.7% 94.5% 94.5% 123B params 94.7% 93.3% 94.9% 94.7% 94.5% 93.5% 94.4% 94.9% 94.7% 94.5% 92.9% 94.7% 93.7% 95.0% 94.2% PPI++ Coverage 7B params 93.9% 92.7% 93.3% 87.5% 81.5% 94.8% 94.6% 94.4% 92.5% 88.5% 93.7% 91.0% 90.8% 93.2% 93.9% 24B params 84.6% 80.6% 93.3% 69.9% 90.7% 85.0% 87.0% 95.0% 94.7% 57.4% 91.1% 83.2% 73.7% 86.6% 84.2% 123B params 89.4% 92.5% 95.0% 77.1% 89.6% 71.3% 93.6% 94.4% 94.6% 75.6% 91.7% 87.6% 71.8% 88.6% 80.7% Greedy Coverage 7B params 93.6% 92.4% 93.9% 89.1% 83.6% 94.8% 94.3% 94.3% 90.4% 88.0% 93.7% 91.5% 90.6% 93.0% 94.0% 24B params 85.9% 89.9% 94.1% 76.3% 91.4% 83.2% 86.5% 94.4% 92.8% 68.5% 91.9% 84.8% 75.6% 84.8% 85.5% 123B params 90.3% 91.8% 94.9% 79.1% 89.7% 74.5% 92.5% 94.4% 93.9% 79.5% 91.7% 86.9% 72.8% 87.0% 82.1% Greedy AIC Coverage 7B params 93.8% 92.7% 93.8% 89.1% 84.0% 94.8% 94.5% 94.4% 91.0% 88.5% 93.8% 91.3% 90.9% 93.0% 94.1% 24B params 86.5% 88.9% 93.6% 77.1% 91.5% 83.6% 87.5% 94.5% 93.0% 69.2% 91.9% 85.0% 75.3% 84.8% 86.5% 123B params 90.2% 92.0% 95.0% 79.6% 89.9% 75.2% 92.7% 94.5% 94.0% 80.0% 91.8% 86.6% 74.2% 88.3% 82.9% Greedy BIC Coverage 7B params 94.0% 92.7% 93.8% 88.9% 84.4% 94.9% 94.6% 94.5% 91.3% 88.9% 93.8% 91.3% 91.0% 93.1% 94.0% 24B params 87.1% 82.0% 93.4% 77.6% 91.6% 84.2% 88.0% 94.6% 93.9% 69.7% 91.4% 85.2% 75.7% 85.2% 87.0% 123B params 89.9% 92.5% 95.1% 79.9% 90.1% 75.7% 93.4% 94.5% 94.3% 80.2% 92.0% 86.7% 74.5% 88.7% 83.2%

Figure 7: Item-level performance for Song Rating.

On the LFM dataset for music engagement, GPPI gives huge improvement over PPI++. The particularity of this dataset is that the engagement is measured by the number of times a user listened to a song, which is an integer that can be very high. The distribution of play counts for each song is a highly skewed distribution with some rare very high values. PPI++ fails to correctly debias the predictions on such dataset, and its performance, both in terms of coverage and ESS gains, drops significantly. Interestingly, GPPI with greedy selection has globally a better coverage than PPI++, although still does not meet the nominal coverage for difficult items, and it has a significantly better ESS gain. While we would expect that on difficult instances PPI++ would be more stable than GPPI with greedy selection, this example shows that, on the opposite, GPPI our approach can be more stable and robust than PPI++. ESS gain vs Sample size Below, we give visualizations of how the ESS gains varies with the sample size for Book-Crossing and Last.fm datasets. The plots for Movielens and Jester datasets are in the main body of the paper. C.6

A/B testing using GPPI with greedy selection

We first present the plots of the experiment 4.2 with Last.fm and Book-Crossing datasets: Consistently with the results for Movielens and Jester datasets, we see that GPPI with greedy transformation selection has better performance than the other approaches. Observe that in some plots (i.e. specific dataset and model choice), PPI++ has negative ESS-gain, which means that the estimation of λ is noisy on the labeled data which leads to degrading the quality of the predictions. On the other hand, GPPI with greedy transform selection always give positve ESS gain, and always better than the vanilla PPI++. 40

Movie Lens 122904

122926

152081 10.0%

5.0%

2.0%

0.0%

7B params ESS Gain

10.0%

1000 2000

24B params ESS Gain 123B params ESS Gain

1000 2000

20.0%

0.0%

250

1000 2000

500 1000 1500

0.0%

250

5.0%

0.0%

PPI++

100 200 300

0.0%

250

0.0%

500

0.0%

250

5.0% 100 200 300

Greedy

0.0%

250

500

250

500

10.0% 5.0%

5.0% 100 200 300

0.0%

20.0%

10.0%

0.0%

500

187593

10.0%

5.0%

5.0% 500 1000 1500

177615 10.0%

10.0%

0.0%

500

10.0%

15.0%

20.0%

10.0%

500

5.0%

40.0%

20.0%

250

10.0%

0.0%

500

30.0%

10.0%

500 1000 1500

10.0%

5.0%

175303 10.0%

4.0%

20.0%

10.0%

5.0%

0.0%

500

15.0%

10.0%

0.0%

250

174055

6.0%

10.0%

100 200 300

Greedy AIC

0.0%

100 200 300

0.0% 20.0% 10.0%

100 200 300

0.0%

250

500

Greedy BIC

Figure 8: ESS gain comparison on MovieLens dataset across different movies and Mistral models

Book Crossing 7B params ESS Gain

0142001740 6.0% 4.0%

10.0%

0316601950

5.0%

2.0%

24B params ESS Gain

0.0%

100

150

2.0%

123B params ESS Gain

120

140

100

150

10.0%

2.0%

5.0%

0.0%

100

200

300

0.0%

2.00%

10.0%

2.5%

1.00%

5.0%

0.0%

2.0% 150

4.0%

15.0%

2.0%

15.0%

3.00%

4.0%

100

0316666343

5.0%

4.0%

0.0%

100

7.5%

4.0%

0.0%

6.0%

0.0%

100

120

140

0.00%

100

200

300

0.0%

044023722X

20.0% 10.0% 0.0%

100

120

0.0%

5.0%

2.5%

1.00%

2.5%

PPI++

0.00%

100

125

150

0.0%

100

120

140

100

Greedy

120

0.0%

100

125

150

Greedy AIC

0.0%

120

100

120

100

120

5.0% 100

120

140

0.0%

20.0%

10.0%

5.0%

0.0%

100

10.0%

5.0% 100

0.0% 15.0%

2.5%

7.5%

300

10.0%

5.0%

2.00%

0743418174 15.0%

10.0%

3.00%

200

0452282152

7.5%

5.0%

100

10.0%

5.0% 100 110 120

7.5%

0.0%

0446672211

100

120

140

0.0%

Greedy BIC

Figure 9: ESS gain comparison on Book-crossing dataset across different jokes and Mistral models To better illustrate the improvement in ESS gains achieved by GPPI with greedy transformation selection compared to PPI++, we conduct A/B tests on all the pairs of items for every dataset, in the same setting described in the experiments section, then we plot the ESS gain of GPPI with greedy transform versus that of PPI++. We also plot the identity line as a reference. Points above this line are those where GPPI with greedy selection is better than PPI++, and below below it are those where PPI++ has a better performance. These figures below all illustrate again the power of our method, showing that it consistently and sometimes significantly outperforms PPI++. C.7

Replication of (Angelopoulos et al., 2023b) experiments

For completeness, we also replicate experiments from (Angelopoulos et al., 2023b) on deforestation and galaxy classification datasets, where they illustrate how PPI++ reduces the confidence interval width compared to the classical estimator. Since this is a classification 41

Last.fm TROMKCG128F9320C09

7B params ESS Gain

4.00%

TRMYECS128F92DE4B4

2.00%

24B params ESS Gain

0.00% 15.0%

10.0%

2.00%

2.0%

5.0%

1.00%

2.0%

0.0%

123B params ESS Gain

5.0%

4.0%

500 1000 1500

0.0%

1000

1.00%

0.100%

0.00%

0.000%

1000 2000

2.00% 0.00%

500 1000 1500

1000

2000

30.0%

6.0%

20.0%

4.0%

10.0% 0

1000 2000

PPI++

0.0%

250

1000

Greedy

0.0%

2000

0.0%

1000

10.0%

5.0%

0.0%

500

0.0%

1000

20.0%

10.0%

5.0%

0.0%

500 1000

Greedy AIC

500

20.0%

500 1000

2.0% 0

TRXWAZC128F9314B3E

5.0%

0.0%

500

0.200%

0.0%

10.0%

2.0%

0.300%

2.00%

TRWAQOC12903CB84CA

4.0%

0.00%

2000

2.0%

4.00%

0.0%

500 1000

1000 2000

4.0%

5.0%

2.0%

3.00%

10.0%

6.0%

0.0%

500 1000

4.0% 500 1000

TRPGPDK12903CCC651

10.0%

6.0%

10.0%

TRPFYYL128F92F7144

4.0%

0.0%

500 1000

TRONYHY128F92C9D11

500

0.0%

1000

500 1000 1500

Greedy BIC

Figure 10: ESS gain comparison on Last.fm across different jokes and Mistral models Book Crossing

Mean ESS Gain (pairs with | | > )

7B parameters

Mean ESS Gain for pairs with | | > 24B parameters

12.5%

4.0%

10.0%

2.0%

7.5%

0.0%

5.0%

4.0% 2.0%

2.5%

2.0% 0

Minimum effect size PPI++

Greedy

0.0%

123B parameters

6.0%

Minimum effect size

Greedy AIC

Greedy BIC

0.0%

Asym =0.25

Minimum effect size

Asym =0.5

Asym =0.75

Figure 11: Comparison of the learning-augmented A/B testing approaches for BookCrossing task, we use transformations different than before, of the form ϕ( p) = [1( p > t1 ), . . . , 1( p > td )], where t1 , . . . , td are evenly spaced in (0, 1). Greedy selection operates on a family of 8 transformations: identity and threshold transformations with d ∈ {5, 10, 15, 20, 25, 30, 35}.

Mean ESS Gain (pairs with | | > )

Figure 18 demonstrates that our method achieves narrower confidence intervals. The coverage is slightly weaker than that of PP++, but still meets the nominal value of 90% for this experiment for most values of n.

7B parameters

12.5% 10.0% 7.5% 5.0% 2.5% 0.0%

Music Rating

Mean ESS Gain for pairs with | | >

10.0%

24B parameters

0.0%

10.0%

20.0%

30.0% 0.0

0.2

0.4

0.6

0.8

1.0

Minimum effect size

PPI++

Greedy

1.2

1.4

Greedy AIC

123B parameters

0.0

0.2

0.4

0.6

0.8

1.0

Minimum effect size Greedy BIC

1.2

1.4

Asym =0.25

0.0

0.2

0.4

0.6

0.8

1.0

Minimum effect size

Asym =0.5

1.2

1.4

Asym =0.75

Figure 12: Comparison of the learning-augmented A/B testing approaches for Last.fm 42

Mean ESS Gain (pairs with | | > )

7B parameters

30.0%

Movie Lens

Mean ESS Gain for pairs with | | >

30.0%

20.0%

10.0%

20.0% 10.0% 0.0%

0.0

0.1

0.2

0.3

Minimum effect size

PPI++

0.4

0.0%

0.5

Greedy

24B parameters

0.0

0.1

0.2

0.3

Minimum effect size

Greedy AIC

Greedy BIC

0.4

0.0%

0.5

Asym =0.25

123B parameters

0.0

0.1

0.2

0.3

Minimum effect size

Asym =0.5

0.4

0.5

Asym =0.75

Figure 13: Comparison of the learning-augmented A/B testing approaches for MovieLens 7B parameters Greedy Greedy AIC Greedy BIC

14.0%

Greedy ESS Gain

Movie Lens

12.0%

Greedy vs PPI++ ESS Gain 24B parameters

35.0%

20.0%

30.0%

8.0%

15.0%

25.0%

6.0%

10.0%

20.0%

4.0%

5.0%

10.0%

4.0% 6.0% 8.0% 10.0% 12.0% 14.0%

15.0% 5.0%

PPI++ ESS Gain

123B parameters

40.0%

25.0%

10.0%

15.0%

20.0%

PPI++ ESS Gain

25.0%

15.0% 20.0% 25.0% 30.0% 35.0% 40.0%

PPI++ ESS Gain

Greedy ESS Gain

Figure 14: Movielens: comparison of GPPI with greedy selection versus PPI++ 14.0% 12.0% 10.0% 8.0% 6.0% 4.0% 2.0% 0.0%

7B parameters

Jester Jokes

Greedy Greedy AIC Greedy BIC

Greedy vs PPI++ ESS Gain 24B parameters

25.0%

123B parameters

35.0%

20.0%

30.0%

15.0%

25.0%

10.0%

20.0%

0.0% 2.0% 4.0% 6.0% 8.0% 10.0% 12.0% 14.0%

10.0%

PPI++ ESS Gain

15.0%

20.0%

PPI++ ESS Gain

25.0%

20.0%

25.0%

30.0%

PPI++ ESS Gain

35.0%

Figure 15: Jester: comparison of GPPI with greedy selection versus PPI++

Greedy ESS Gain

20.0% 10.0%

7B parameters

Book Crossing

Greedy Greedy AIC Greedy BIC

0.0% 10.0% 20.0% 20.0%

10.0%

0.0%

10.0%

PPI++ ESS Gain

20.0%

Greedy vs PPI++ ESS Gain 24B parameters

15.0% 10.0% 5.0% 0.0% 5.0% 10.0% 15.0% 20.0% 20.0%15.0%10.0% 5.0%0.0% 5.0% 10.0%15.0%

123B parameters 20.0% 15.0% 10.0% 5.0% 0.0% 5.0% 5.0% 0.0% 5.0% 10.0% 15.0% 20.0%

PPI++ ESS Gain

Figure 16: Book crossing: comparison of GPPI with greedy selection versus PPI++ 7B parameters Greedy ESS Gain

20.0% 15.0%

Music Rating

Greedy Greedy AIC Greedy BIC

10.0% 5.0% 0.0% 5.0% 5.0% 0.0% 5.0% 10.0% 15.0% 20.0%

PPI++ ESS Gain

Greedy vs PPI++ ESS Gain 24B parameters

123B parameters 20%

20% 0%

20%

40%

60%

80%

60%

40%

20%

PPI++ ESS Gain

20%

80%

60%

40%

20%

PPI++ ESS Gain

20%

Figure 17: LFM music: comparison of GPPI with greedy selection versus PPI++ 43

CI Width: Forest

200

400

600

800

1000

n (labeled sample size)

Coverage

0.6

CI Width

Coverage

Classical PPI++ Greedy Greedy AIC Greedy BIC Nominal 90%

0.8

Coverage: Galaxies

1.0

0.08

Classical PPI++ Greedy Greedy AIC Greedy BIC Nominal 90%

0.8

0.06

0.6

0.04 200

400

600

800

1000

200

n (labeled sample size)

400

600

800

CI Width: Galaxies

0.10

1000

n (labeled sample size)

0.08

CI Width

Coverage: Forest

1.0

0.06 0.04 200

400

600

800

1000

n (labeled sample size)

Figure 18: Mean estimation results on deforestation (left) and galaxies (right) datasets showing coverage and confidence interval width.

Record · ID 1028707 · SHA-256 b0e4e6efe33425e6
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.