Expected Gain-based Escalation in Vertical Federated Learning

Technical Summary: Expected Gain-based Escalation in Vertical Federated Learning

Problem Statement
In Vertical Federated Learning (VFL), multiple clients hold different views of the same sample, and a server aggregates this information to improve predictive performance. While collaborative inference (fusing intermediate representations from all clients) generally yields higher accuracy than local inference, it incurs significant communication and computational overhead. The central challenge addressed in this work is determining when to invoke this expensive collaborative route. Applying fusion to every sample is inefficient, as many inputs are easily resolved by lightweight local predictions. The goal is to develop a per-sample routing policy that decides whether the expected improvement in correctness justifies the additional communication cost, specifically within a two-round inference protocol.

Methodology
The authors propose an analytical routing rule based on expected gain estimation, avoiding the need for a separately trained routing network. The system operates in two rounds:

1. Round 1 (Local): Clients transmit low-dimensional local posteriors to the server. The server forms a pooled prediction via average pooling.
2. Round 2 (Collaborative): If a sample is escalated, clients transmit high-dimensional embeddings, which are fused by a stronger server-side VFL model.

The routing decision is modeled as a threshold-based policy: escalate if the expected gain $E[G | \bar{p}]$ exceeds a threshold $\tau$, where $G$ is the difference in correctness between the Round 2 and Round 1 predictions.

To compute this score analytically without training a router, the authors derive a plug-in approximation using two components estimated from held-out calibration data:

1. Calibrated Pooled Posterior: The raw pooled posterior from Round 1 is often uncalibrated. The authors apply a multiclass Dirichlet calibration map to estimate the conditional class posterior $P(y=k | \bar{p})$.
2. Classwise Reliability: The authors approximate the conditional correctness of the Round 2 model, $P(\hat{y}_{VFL}=k | y=k, \bar{p})$, using a class-level reliability coefficient (the diagonal entry of the Round 2 confusion matrix, $CVFL[k,k]$). This assumes that once the true class is fixed, the Round 2 model's performance is relatively stable regardless of the Round 1 posterior.

The final routing score $\hat{s}(\bar{p})$ is calculated as:
$$\hat{s}(\bar{p}) = \sum_{k=0}^{K-1} \hat{m}_{\hat{\phi}}(\bar{p})[k] \cdot (CVFL[k,k] - \mathbb{1}[\hat{y}_{loc} = k])$$
where $\hat{m}_{\hat{\phi}}$ is the calibrated posterior and $\hat{y}_{loc}$ is the Round 1 prediction. A sample is escalated if this score is positive (or exceeds a budget-controlled threshold).

Key Contributions

1. Inference-Time Selective Escalation: The work formalizes the problem of selectively invoking a costly collaborative VFL route in a two-round protocol, balancing accuracy against communication budgets.
2. Analytical Routing Rule: The authors propose a method that estimates the expected gain of escalation using only the pooled Round 1 posterior and classwise reliability statistics. This approach requires no dedicated routing network training, making it interpretable and lightweight.
3. Empirical Validation: The method is validated on multi-view classification benchmarks (CIFAR-10, CIFAR-100, ModelNet40), including scenarios with controlled test-time view degradation (Gaussian noise).

Results
Experiments demonstrate that the proposed analytical router achieves a superior accuracy-escalation trade-off compared to several baselines:

• vs. Confidence-based: The proposed method outperforms simple confidence thresholds, indicating that confidence alone is insufficient to predict the benefit of escalation.
• vs. Learned Gain (DIFF-01): The analytical approach consistently outperforms a learned router trained to regress the realized gain, suggesting that the analytical approximation captures the essential routing signal effectively.
• vs. Learning-to-Defer (L2D): At matched escalation budgets, the proposed router achieves higher accuracy and better escalation precision (EP), capturing a larger fraction of beneficial escalations while avoiding neutral ones.
• vs. Oracle: While the proposed method does not reach the theoretical upper bound of an oracle (which knows the true gain), it significantly narrows the gap, particularly in moderate escalation regimes (40–60%).

The results also highlight that a large portion of test samples are "neutral" regarding escalation (i.e., Round 1 and Round 2 yield the same correctness). The proposed router effectively filters these out, whereas aggressive baselines often over-escalate, incurring communication costs without accuracy gains.

Significance and Claims
The paper claims that treating escalation as an expected-gain decision allows for a practical, interpretable, and efficient routing mechanism in VFL. By relying on calibration data rather than complex learned routers, the method offers a straightforward deployment path. The authors emphasize that their approach improves the communication-accuracy trade-off over existing confidence-based, learned-gain, and deferral-based baselines.

The work acknowledges limitations, specifically that the classwise reliability approximation ignores within-class variation in the Round 2 model's behavior and that performance depends on the representativeness of the calibration set. However, the authors assert that the method remains robust even under distribution shifts (simulated by view degradation), provided the calibration data is reasonably representative. The study concludes that selective escalation is a viable strategy for reducing inference costs in VFL without sacrificing predictive performance.