Learning U-Statistics with Active Inference

arXiv:2605.11638 · stat.ML, cs.LG · Submitted 2026-05-12 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Learning U-Statistics with Active Inference".

Jane: The gist Learning U-Statistics with Active Inference develops an active inference framework for U-statistics that selectively queries informative labels to improve estimation efficiency under a fixed labeling budget while preserving…

Tom: First, who's behind it and why it matters.

Paper summary: Tom: We started by looking at how we can improve estimation efficiency when acquiring labels is costly, and this paper proposes an active inference framework for U-statistics.

Jane: Essentially, they claim you can selectively query informative labels to boost estimation efficiency while sticking to a fixed labeling budget without ruining the statistical validity.

Lu: The mechanism they use is the augmented inverse probability weighting U-statistic, which mixes the sampling rule and machine learning predictions into one estimator.

Meng: So it’s not just a standard U-statistic plus some tacked-on prediction term; it’s integrated into the core structure of how you calculate things.

Tom: Right, they define this augmented inverse probability weighting U-statistic as U pi AIPW = one over n sum i=one n r X C n,r h Y one Y r + one over n sum i=one n r X C n,r h(i one i r) xi i one xi i r pi i one pi i r <ref:2605.11638#pg1,augmented inverse probability weighting U-statistic>

Jane: That looks dense. The key is the IPW term that corrects the bias from plugging in machine learning predictions Y = i = mu(X i) into the U-statistic.

Lu: They show that this AIPW U-statistic is unbiased, satisfying EU pi AIPW = theta*, which is the target parameter for the U-statistic >

Meng: Unbiasedness is crucial, but then they immediately ask how to choose the sampling rule pi(times) to minimize its variance under that budget constraint.

Tom: That’s where they characterize the optimal sampling rule in Theorem two point two, stating it should be proportional to q E

h one(Y) - mu one(Y) squared X: > <ref:2605.11638#pg1,characterize the optimal sampling rule>

Jane: So the strategy is driven by minimizing the residual uncertainty from the first-order Hoeffding projection, that uncertainty in h one(Y) - mu one(Y), rather than just predicting how far Y is from > <ref:2605.11638#pg1>

Lu: That’s a big conceptual point because it shifts the focus to the model's own uncertainty about the functional being estimated, not just the raw prediction error >

Meng: If that holds up, it suggests we should be sampling where our model is most uncertain about its own function outputs, which is a much more targeted approach for label acquisition >

Tom: And they take this further by applying it to U-statistics-based empirical risk minimization, leading to an optimal policy pi*(X) = (one nb/n times p S(X)/E

p S(X): ) > <ref:2605.11638#pg1>

Jane: That policy aims to minimize the trace of the covariance matrix of the resulting estimator pi act, which is a standard way to look for optimal sampling designs in these problems >

Lu: It connects this active inference idea directly to minimizing the trace of that covariance matrix, which leads to a concrete, actionable policy >

Meng: I wonder how much computational work it takes to calculate that policy pi*(X) on real-time data streams; is it feasible for deployment?

Tom: The paper suggests the strategy is robust enough for higher-order U-statistic tasks and remains reasonably stable even if there's misspecification in the predictive model, which is a plus >

Jane: So, we see this as a way to make label acquisition an active part of the learning process rather than just a passive data gathering step.

Conclusion: Tom: To wrap up, "Learning U-Statistics with Active Inference" shows that you can actively use machine learning predictions to guide how you acquire labels for complex statistics like U-statistics.

Jane: The authors successfully developed the AIPW U-statistic and derived a principled sampling rule that targets variance minimization under a labeling budget constraint.

Lu: It's an extension of active learning, moving it from just prediction uncertainty to inference uncertainty about the functional being estimated >

Meng: So for those deploying models, this implies you should think about which data points you need labeled based on what your current AI model is uncertain about its own function outputs >

Tom: The implication is that we can get much better estimation efficiency on real datasets while still maintaining the necessary statistical rigor for inference >

Jane: The paper concludes that this active inference method remains efficient for higher-order U-statistic tasks and is reasonably robust to predictive model misspecification, which is a key takeaway >

Nankai University

stat.ML, cs.LG

Submitted: 2026-05-12

Updated: 2026-05-12

Journal ref: Proceedings of the 43rd International Conference on Machine Learning (ICML), PMLR 306:127065-127088, 2026

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 77/100

The gist: The gist Learning U-Statistics with Active Inference develops an active inference framework for U-statistics that selectively queries informative labels to improve estimation efficiency under a fixed

Key concepts

Augmented Inverse Probability Weighting U-statistic (AIPW): This is the core statistical tool used, designed to combine the standard U-statistic with machine learning predictions and an inverse probability weighting term. It helps correct bias by incorporating both sampling rules and model outputs into the estimation process, leading to more accurate results.
Optimal Sampling Rule: This rule dictates which data points should be selected for labeling based on a specific criterion. The paper shows that the optimal selection probability is proportional to the uncertainty in the first-order Hoeffding projection, meaning it prioritizes samples where the model's prediction error is highest.
U-statistic-based Empirical Risk Minimization (ERM): This extension applies the active inference framework to finding parameters by minimizing an empirical risk defined by a U-risk. Instead of just estimating a single statistic, it uses active learning to select labels that best minimize this risk function under the labeling budget.
Labeling Budget Constraint: This refers to the fixed, limited number of labels available for acquisition. The framework is specifically designed to maximize estimation efficiency—getting the most information out of these limited labels—while ensuring that the final statistical inference remains valid.

Terminology

Active Inference — 3× in this paper · explained in 1 episode

Summary

The gist Learning U-Statistics with Active Inference develops an active inference framework for U-statistics that selectively queries informative labels to improve estimation efficiency under a fixed labeling budget while preserving valid statistical inference

Learning U-Statistics with Active Inference

The paper introduces an active inference framework for U-statistics designed to selectively query informative labels to improve estimation efficiency under a fixed labeling budget while preserving valid statistical inference The approach is built on the augmented inverse probability weighting U-statistic, which incorporates the sampling rule and machine learning predictions The goal is to design label-acquisition strategies that strategically select samples for labeling to minimize the estimation error of a target U-statistic, while enabling valid follow-up inference under a fixed labeling budget

Active Inference for U-statistics

The framework utilizes the augmented inverse probability weighting (AIPW) U-statistic to incorporate machine learning predictions and correct bias using an IPW term The AIPW U-statistic is defined as U pi AIPW = 1/n r X Cn,r h Yˆi1,..., Yˆir + 1/n r X Cn,r ∆h(i1,..., ir) ξi1 · · · ξir πi1 · · · πir The optimal sampling rule is characterized by Theorem 2.2, which states that the optimal sampling probability should satisfy π(X) ∝ q E[h1(Y) − hˆµ 1(Yˆ) 2 X] to minimize the asymptotic variance of U pi AIPW under the budget constraint This optimal rule is driven by the residual uncertainty in the first-order Hoeffding projection, h1(Y)− hˆµ 1(Yˆ), rather than the prediction residual Yˆ − Y that is natural for mean estimation

Asymptotic Normality and Variance Estimation

The paper establishes the asymptotic normality of the proposed active U-statistics in Theorem 2.4, showing that √n (U pi hat AIPW − θ∗) d→ N (0, σ2∗), where σ2∗ is given by a specific formula involving Var(hµ1(Yˆ)) and [h1(Y) − hˆµ 1(Yˆ)] ξ∗τ π∗τ (X) Furthermore, the variance estimator σb2 is shown to be consistent for the CLT variance parameter σ2∗ by estimating both components using Horvitz–Thompson (HT) corrections

Active Inference for U-estimation

The framework is extended to U-statistics-based empirical risk minimization (ERM), where the target parameter is the minimizer of a population U-risk The active U-estimator is defined as the minimizer of the empirical U-risk L act n,π(θ), which combines predicted labels with an IPW correction based on actively acquired labels The optimal sampling rule for active U-estimation is derived by minimizing the trace of the covariance matrix tr(Cov(ˆθ pi act)), leading to a policy π∗(X) = min (1, nb/n · p S(X)/E[p S(X)])

Experiments and Results

Numerical experiments on real datasets demonstrate substantial gains in estimation efficiency over baseline methods, while maintaining target coverage On the Income dataset (ACS), the active sampling method achieves a markedly larger effective sample size, corresponding to roughly a 60% labeling-budget reduction relative to classical one The results show that incorporating predictions improves efficiency: uniform sampling achieves a positive saving ratio relative to the classical estimator

Discussion on Degenerate U-Statistics

When the target U-statistic is degenerate, the first-order projection vanishes and the leading contribution comes from higher-order Hoeffding components The resulting limiting distribution of the AIPW U-statistic is generally nonGaussian, which makes the construction of valid confidence intervals substantially more delicate For degenerate U-statistics, a possible practical approach is to estimate the pairwise score T2(x, x′) from a pilot labeled sample and then solve the resulting policy by a fixed-point or alternating optimization procedure

Conclusion

Overall, this strategy highlights a key advantage of our approach: by exploiting the U-risks structure, labels are actively queried for efficient U-statistics-based ERM The paper concludes that the proposed active inference method remains efficient for higher-order U-statistic tasks and is reasonably robust to predictive-model misspecification

--- Page 1 ---

How it works

The approach is built on the augmented inverse probability weighting U-statistic, which incorporates the sampling rule and machine learning predictions The AIPW U-statistic is defined as U pi AIPW = 1/n r X Cn,r h Yˆi1,..., Yˆir + 1/n r X Cn,r ∆h(i1,..., ir) ξi1 · · · ξir πi1 · · · πir The optimal sampling rule is characterized by Theorem 2.2, which states that the optimal sampling probability should satisfy π(X) ∝ q E[h1(Y) − hˆµ 1(Yˆ) 2 X] to minimize the asymptotic variance of U pi AIPW under the budget constraint

Asymptotic Normality and Variance Estimation

The paper establishes the asymptotic normality of the proposed active U-statistics in Theorem 2.

Improvements for AI systems

textbfActive U-statistic Label Acquisition for Parameter Estimation (ERM): It allows AI systems to efficiently estimate complex parameters like population risk or ranking rules in label-limited settings by adaptively querying informative labels that minimize estimation error under a fixed budget, as shown in Algorithm 1 and Section 3.1.
textbfVariance-Optimal Label Querying: The system can select samples based on residual uncertainty in the first-order Hoeffding projection, specifically by setting the sampling probability to min (1, nb/n · p s(X) E[p s(X)]) as derived in Theorem 2.2, leading to substantially improved estimation efficiency compared to existing baselines.
textbfRobust Inference under Model Misspecification: The framework ensures statistical validity even when the predictive model is inaccurate, because the AIPW U-statistic is corrected by the IPW label term, allowing for valid statistical inference despite using a potentially misspecified model µ(X).
textbf Higher-Order Structure Estimation: The system can be extended to estimate higher-order statistics, such as the third central moment (skewness), by utilizing the U-estimation framework in Section 3, demonstrating effectiveness beyond pairwise targets like the Gini index or Kendall’s tau.
textbf Adaptive Trimming for Stability: To prevent variance inflation from rare sampling events, the system employs an adaptive trimming strategy defined as π trim i = τπ i + (1 − τ)nb/n, ensuring that the minimal sampling probability is bounded below by (1−τ)nb/n throughout the process.