Heterogeneity-Aware Distributional Framework (HDF)

• HDF is a methodological framework that models heterogeneity as an explicit distributional object rather than a mere nuisance, using tools like latent variables, transport maps, and robust objectives.
• It spans diverse domains including federated learning, dynamic facial expression recognition, causal inference, meta-analysis, and latent-variable modeling, each tailoring specific distributional mechanisms.
• Empirical results demonstrate that HDF-based methods improve performance metrics and reveal deeper insights into treatment heterogeneity and distributional shifts across multiple applications.

The papers considered here use the label Heterogeneity-aware Distributional Framework (HDF) for several distinct methodological frameworks that treat heterogeneity as a distributional problem rather than only a nuisance in conditional means. In these works, the relevant distribution may be a client-specific data law in federated learning, a video-feature distribution under sample heterogeneity in dynamic facial expression recognition, one-sided outcome distributions at a regression discontinuity threshold, effect-size distributions in meta-analysis, subject-specific density responses, latent person-level distributional features, or treated and counterfactual outcome laws in distributional Difference-in-Differences (Yuan et al., 2022, Cui et al., 21 Jul 2025, Schindl et al., 22 Feb 2026, Yang et al., 1 Apr 2026, Bhattacharjee et al., 20 Jun 2026). This suggests that HDF is best understood not as a single canonical architecture, but as a recurring methodological stance in which heterogeneous structure is modeled explicitly through transformations, robust objectives, latent variables, transport maps, or distributional regression.

1. Scope and domain-specific uses

The summarized literature spans machine learning, computer vision, causal inference, meta-analysis, latent-variable modeling, and nonparametric distributional regression. In the federated-learning setting, the corresponding framework is called DisTrans in the paper itself; in dynamic facial expression recognition, the term Heterogeneity-aware Distributional Framework is used directly; in several statistical settings, the summary applies the same label to frameworks centered on Wasserstein distances, distributional regression, or latent distributional features (Yuan et al., 2022, Cui et al., 21 Jul 2025, Schindl et al., 22 Feb 2026, Yang et al., 1 Apr 2026, Fazio et al., 14 Jun 2026).

Area	Representative method	Key distributional mechanism
Federated learning	DisTrans (Yuan et al., 2022)	client-specific offsets $\Delta_i$ and double-input transformation $T(x;\Delta_i)$
Dynamic facial expression recognition	HDF (Cui et al., 21 Jul 2025)	DRO objective with DAM and DSM
Regression discontinuity and kink designs	Distributional discontinuity designs (Schindl et al., 22 Feb 2026)	$W_1(F^+,F^-)$ and Wasserstein derivative
Meta-analysis	Distributional regression models (Yang et al., 1 Apr 2026)	location, scale, and shape as functions of covariates
Density-valued response modeling	additive model with latent groups (Han et al., 2021)	LQD map $\Psi$ and grouped subject-specific curves
Latent-variable modeling	DFLVM (Fazio et al., 14 Jun 2026)	random intercepts for distributional features
Difference-in-Differences	distributional DiD test (Bhattacharjee et al., 20 Jun 2026)	optimal transport counterfactual and MMD
Distributional regression forests	DRF (Ćevid et al., 2020)	MMD-based splits and weighted empirical conditional law
OOD generalization	HRM (Liu et al., 2021)	latent environment discovery and invariant prediction

A common feature across these settings is that heterogeneity is not treated as a residual irregularity. It is instead encoded into an explicit mathematical object: a transport map, an offset vector, a conditional law, a latent random intercept, a group-specific function, or an environment partition.

2. Distributional representations and transformations

In federated learning, DisTrans begins from the standard cross-device setting with local datasets $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$, local distributions $P_i(x,y)$, and a global model $f_\theta:\mathbb R^m\to\mathbb R^N$. Its central representation of heterogeneity is the client-specific offset $\Delta_i\in\mathbb R^m$, which shifts client data toward a more canonical distribution. The transformation is

$$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$

so that two transformed views are fed in parallel to a shared backbone. The resulting double-input-channel architecture computes $h_{i,1}=f_{\text{backbone}}(x'_{i,1})$ and $T(x;\Delta_i)$0, concatenates $T(x;\Delta_i)$1, and then predicts logits. The offset magnitude is controlled by limiting $T(x;\Delta_i)$2, using small learning rates for $T(x;\Delta_i)$3, and optionally by an $T(x;\Delta_i)$4 penalty (Yuan et al., 2022).

In dynamic facial expression recognition, HDF casts the problem as Distributionally Robust Optimization. The objective replaces empirical-risk minimization

$T(x;\Delta_i)$5

with the worst-case formulation

$T(x;\Delta_i)$6

Its Time–Frequency Distributional Attention Module (DAM) contains a frequency branch and a temporal branch. The frequency branch applies a $T(x;\Delta_i)$7 DCT-like convolution, adds adversarial-style perturbations, normalizes through a dynamic activation fitting step, and then applies self-attention and residual fusion. The temporal branch combines global context, local saliency, Wasserstein regularization, and local compensation into a gated temporal attention map. The two branches are fused adaptively through

$T(x;\Delta_i)$8

with $T(x;\Delta_i)$9 increasing when temporal inconsistency dominates and $W_1(F^+,F^-)$0 increasing when style variation dominates (Cui et al., 21 Jul 2025).

In density-valued response modeling, the key representation is the Log-Quantile-Density map

$W_1(F^+,F^-)$1

which embeds densities into the linear space $W_1(F^+,F^-)$2. This supports the additive-heterogeneity model

$W_1(F^+,F^-)$3

where subject-specific curves are homogeneous within latent groups but heterogeneous across groups (Han et al., 2021).

In the Distributional Feature Latent Variable Model, heterogeneity is encoded through person-specific parameters

$W_1(F^+,F^-)$4

where $W_1(F^+,F^-)$5 is an individual-specific random intercept for a distributional feature such as location, scale, or shape. The same latent intercepts then enter a downstream outcome model for $W_1(F^+,F^-)$6, so the associations between distributional features and outcomes are estimated in a single step rather than via plug-in summaries (Fazio et al., 14 Jun 2026).

3. Optimization, estimation, and algorithmic structure

DisTrans uses per-client joint optimization. In round $W_1(F^+,F^-)$7 at client $W_1(F^+,F^-)$8, it initializes $W_1(F^+,F^-)$9 and $\Psi$0 and minimizes

$\Psi$1

Each minibatch alternates two SGD-style steps: one update for $\Psi$2 and one update for $\Psi$3. After local training, the server aggregates model weights in FedAvg style,

$\Psi$4

and updates offsets by one of three modes: no-agg, avg-agg, or NN-agg, where the last uses a small server-side network keyed by the client’s class-ratio embedding $\Psi$5 (Yuan et al., 2022).

HRM uses a different algorithmic decomposition: a heterogeneity identification module $\Psi$6 and an invariant-prediction module $\Psi$7. The clustering stage fits a $\Psi$8-component mixture model for the joint $\Psi$9 distribution by minimizing

$D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$0

optimized by EM. The invariant-prediction stage learns model parameters $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$1 and a soft feature-selection mask $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$2 through a regularized objective containing both an $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$3 penalty and an environment-wise gradient-variance penalty,

$D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$4

thereby enforcing invariance across the discovered environments (Liu et al., 2021).

DRF uses a tree-based optimization principle centered on the Maximum Mean Discrepancy. At each node, for candidate split sets $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$5 and $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$6, it maximizes

$D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$7

where $D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$8 is a random-feature approximation to a characteristic kernel. The induced forest weights define the empirical conditional law

$D_i=\{z_i=(x_i,y_i)\}_{i=1}^{n_i}$9

Because the estimator is target-free, arbitrary downstream quantities can then be obtained by plug-in evaluation on $P_i(x,y)$0 (Ćevid et al., 2020).

Distributional regression for meta-analysis instead specifies a parametric family

$P_i(x,y)$1

with each parameter linked to covariates through

$P_i(x,y)$2

Estimation can proceed by Maximum penalized likelihood, Restricted maximum likelihood, or Fully Bayesian inference, and heterogeneity enters both through within-study sampling variance and through random effects in the distributional parameters themselves (Yang et al., 1 Apr 2026).

4. Distributional estimands, decomposition, and causal heterogeneity

In distributional discontinuity designs, the core estimand is the 1-Wasserstein distance between the two one-sided limit distributions at the cutoff:

$P_i(x,y)$3

The usual mean-based regression discontinuity estimand is

$P_i(x,y)$4

and the framework shows

$P_i(x,y)$5

with equality if and only if the treatment effect is purely additive. The decomposition

$P_i(x,y)$6

implies that the gap between the mean jump and the Wasserstein distance measures heterogeneity across quantiles. This is formalized by the heterogeneity index

$P_i(x,y)$7

The same framework gives an orthogonal $P_i(x,y)$8-moment decomposition,

$P_i(x,y)$9

which attributes the total distributional shift to location, scale, skewness, and higher-order shape components (Schindl et al., 22 Feb 2026).

The framework extends to distributional kink designs by defining the Wasserstein derivative at a kink. In the sharp case,

$f_\theta:\mathbb R^m\to\mathbb R^N$0

which describes the instantaneous rate of mass-transport through the kink. In the fuzzy case, identification proceeds through the local Wald ratio for the derivative of the CDF at the kink and the corresponding fuzzy Wasserstein derivative (Schindl et al., 22 Feb 2026).

In distributional Difference-in-Differences, the counterfactual treated post-treatment law is built by transporting the treated baseline distribution using the control-group drift map

$f_\theta:\mathbb R^m\to\mathbb R^N$1

The null hypothesis is

$f_\theta:\mathbb R^m\to\mathbb R^N$2

The test statistic is based on the RKHS distance

$f_\theta:\mathbb R^m\to\mathbb R^N$3

with empirical statistic

$f_\theta:\mathbb R^m\to\mathbb R^N$4

Under the null, $f_\theta:\mathbb R^m\to\mathbb R^N$5 converges to a Gaussian quadratic form; under Pitman-contiguous alternatives it yields a noncentral chi-square mixture; under moderate deviations it is consistent. Because characteristic kernels metrize weak convergence, the test is sensitive to changes in location, scale, shape, and tail behavior (Bhattacharjee et al., 20 Jun 2026).

These causal-inference variants make explicit that a distributional estimand is not merely a descriptive supplement to a mean effect. In the cited formulations, it is the primary object through which treatment heterogeneity is identified and decomposed.

5. Empirical results and application domains

In federated learning, DisTrans is evaluated on CH-MNIST, CIFAR-10, CIFAR-100, Bird-200, BioID, and CelebA, under distributional heterogeneity levels

$f_\theta:\mathbb R^m\to\mathbb R^N$6

Across all datasets and all DH values, it outperforms FedAvg, pFedMe, pFedHN, MOON, and FedAwS. The reported typical absolute gains versus FedAvg range from 1%–10% in top-1 accuracy. Ablations show +3%–9% gain for the double-channel model over a single-channel version; NN-agg is best for DH<50% and no-agg is best for DH>50%; the best setting is $f_\theta:\mathbb R^m\to\mathbb R^N$7; the best number of local epochs is $f_\theta:\mathbb R^m\to\mathbb R^N$8; convergence occurs in $f_\theta:\mathbb R^m\to\mathbb R^N$9 rounds; and model-size overhead is $\Delta_i\in\mathbb R^m$0 extra bytes for typical backbones (Yuan et al., 2022).

In dynamic facial expression recognition, HDF is evaluated on DFEW and FERV39k. On DFEW, it achieves 71.60% WAR and 60.40% UAR, exceeding the reported 70.84% WAR of CLIPER and 58.89% UAR of LG-DSTF. On FERV39k, it reaches 50.30% WAR and 40.49% UAR. The DFEW ablations report: Baseline (X3D+CE): 68.62% WAR, 58.21% UAR; +DAM only: 69.73% / 59.02%; +DSM only: 71.79% / 59.93%; DAM + DSM: 73.24% / 61.31%. The stated improvement over baseline is +4.62% WAR and +3.10% UAR (Cui et al., 21 Jul 2025).

In causal inference, the distributional discontinuity framework is illustrated by two re-analyses. For U.S. House incumbency (Lee 2008 RDD), the estimate is $\Delta_i\in\mathbb R^m$1 versus $\Delta_i\in\mathbb R^m$2, with a small heterogeneity index $\Delta_i\in\mathbb R^m$3 and an $\Delta_i\in\mathbb R^m$4-moment split showing that $\Delta_i\in\mathbb R^m$5 of shift comes from the mean. For the Swedish grant kink (Lundqvist 2014 RKD), $\Delta_i\in\mathbb R^m$6 but $\Delta_i\in\mathbb R^m$7, and most distance $\Delta_i\in\mathbb R^m$8 comes from higher-order moments, indicating tail or skew effects even when mean drift is null (Schindl et al., 22 Feb 2026).

In meta-analysis, the illustrative example uses 67,393 meta-analyses from Cochrane and a location-scale Normal model. The reported findings are that about 13.0% of meta-analyses showed $\Delta_i\in\mathbb R^m$9 at $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$0, an Egger’s-test analogue for the location parameter, and about 2.3% showed $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$1 at $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$2, indicating heterogeneity increased with the mean standard error (Yang et al., 1 Apr 2026).

In latent-variable modeling, simulation results compare Oracle, DFLVM, LC two-step, and FP two-step. The main findings are: DFLVM $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$3 Oracle: almost zero bias, nominal Type I error, well-calibrated posteriors, best $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$4; LC two-step: small downward bias on standardized $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$5, slight inflation in Type I error when $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$6 is small or $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$7 large; FP two-step: large bias, poor coverage, inflated error rates, worst $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$8; and DFLVM takes $$T(x;\Delta_i)\equiv((1-\alpha)x+\alpha\Delta_i,\;(1+\alpha)x-\alpha\Delta_i),$$9–$h_{i,1}=f_{\text{backbone}}(x'_{i,1})$0 the time of the LC first stage, but remains feasible (Fazio et al., 14 Jun 2026).

In distributional Difference-in-Differences, the application to the Card–Krueger minimum-wage data reports that classical DiD finds no mean effect on employment, whereas the HDF test rejects at $h_{i,1}=f_{\text{backbone}}(x'_{i,1})$1, indicating changes in dispersion and tail behavior of employment (Bhattacharjee et al., 20 Jun 2026).

In distributional regression forests, the documented examples include multivariate distribution benchmarks, Air Quality, conditional copulas and independence, heterogeneous regression and causal effects, birthweight, and fairness in wages. One reported empirical contrast is a fairness-adjusted pay gap (11%) instead of observed 17%. In HRM, synthetic and real-world tasks show the best mean accuracy (or RMSE), the smallest variance across test environments, and lowest worst-case loss, even outperforming IRM given true env labels (Ćevid et al., 2020, Liu et al., 2021).

6. Interpretation, related methods, and recurring misconceptions

A common misconception is that heterogeneity-aware distributional methods are simply mean-based models with an auxiliary variance correction. The summarized frameworks instead make the full distribution, or a distribution-sensitive surrogate, the primary modeling target. Distributional regression for meta-analysis allows all parameters of the effect size distribution, such as location, scale, and shape, to be modelled as functions of explanatory variables, and explicitly subsumes random-effects, multilevel, multivariate, location-scale, and outlier-robust meta-analyses as special cases. DRF is described as a nonparametric, target-free method for estimating the full multivariate conditional distribution. DFLVM is motivated by the claim that two-step estimation of distributional features ignores estimation error and can therefore lead to biased estimates and increased error rates (Yang et al., 1 Apr 2026, Ćevid et al., 2020, Fazio et al., 14 Jun 2026).

A second misconception is that HDF denotes a single framework with a fixed set of modules. The papers summarized here instead include client-specific offsets and double-input channels in federated learning, DAM and DSM in DFER, Wasserstein distances and Wasserstein derivatives in discontinuity and kink designs, optimal-transport counterfactuals and MMD in DiD, latent-environment discovery in HRM, and LQD-space additive modeling with hierarchical agglomerative clustering in density-valued responses (Yuan et al., 2022, Cui et al., 21 Jul 2025, Schindl et al., 22 Feb 2026, Bhattacharjee et al., 20 Jun 2026, Liu et al., 2021, Han et al., 2021).

What unifies these otherwise different constructions is methodological rather than architectural. Each framework begins by positing that the relevant heterogeneity resides in a distributional object: a shifted input distribution, a worst-case uncertainty set, a quantile function, a latent feature distribution, a multivariate conditional law, or a counterfactual transport. Estimation or learning then proceeds by a mechanism matched to that object: alternating SGD with server aggregation, adversarial-style perturbation and adaptive loss weighting, orthogonal decomposition in $h_{i,1}=f_{\text{backbone}}(x'_{i,1})$2-moments, EM plus invariance regularization, one-step Bayesian inference, or MMD-based forest splitting. This suggests a broad research program in which heterogeneity is neither marginalized away nor reduced to a scalar nuisance parameter, but treated as structural information about how distributions vary across clients, subjects, studies, environments, or policy regimes.