Scaling Federated Linear Contextual Bandits via Sketching

Abstract: In federated contextual linear bandits, high data dimensionality incurs prohibitive computation and communication costs: local agents perform $O(d<sup>3)$-time determinant computation and upload $O(d<sup>2)$ parameters, making existing algorithms unscalable, where $d$ is the dimension of data. To relieve these scaling bottlenecks, this paper proposes Federated Sketch Contextual Linear Bandits (FSCLB). On the computation side, FSCLB uses SVD to indirectly obtain the determinant required for communication, eliminating the prohibitive cost of direct determinant calculation and cutting complexity from $O(d<sup>3)$ to $O(l<sup>2d)$ per round, where $l&lt; d$ is the sketch size. On the communication side, FSCLB introduces a double-sketch strategy that reduces both upload and download costs from $O(d<sup>2)$ to $O(ld)$. Naively involving sketch update into federated contextual linear bandits can destroy the local increment and invalidate the asynchronous communication condition; FSCLB solves this by replacing the covariance matrix with the sketch matrix when deciding whether to communicate. Theoretically, FSCLB achieves a regret bound of $$\widetilde{O} ((\sqrt{d}+\sqrt{M\varepsilon_l})\sqrt{lT})$$, where $\varepsilon_l$ is the upper bounded by the spectral tail of the covariance matrix; when $l$ exceeds the rank of the covariance matrix, the bound simplifies to $\widetilde{O}(\sqrt{ldT})$, matching the optimal no-sketch regret. Experiments on both synthetic and real-world datasets show that FSCLB significantly reduces computational and communication costs by over 90 \% while sacrificing only a negligible amount of cumulative reward.

• The paper introduces FSCLB, which leverages SVD-based sketching to reduce per-round computational complexity from O(d³) to O(l²d), enabling scalable federated bandit learning.
• It implements a double sketch communication protocol that lowers bandwidth needs from O(d²) to O(ld) while maintaining strong regret guarantees.
• Empirical results show over 90% reduction in computation and communication costs with minimal reward degradation on both synthetic and real-world datasets.

Scaling Federated Linear Contextual Bandits via Sketching

Introduction and Motivation

The federated contextual linear bandit (FCLB) problem is a fundamental framework for online learning, combining the exploration–exploitation challenge of multi-armed bandits with high-dimensional contextual information and distributed privacy-preserving learning. Conventional FCLB algorithms, exemplified by FedLinUCB, suffer from cubic computational costs ($O(d^3)$ per round, due to matrix determinant and inversion operations) and quadratic communication costs ($O(d^2)$ per communication, due to transmission of $d \times d$ matrices), both of which are prohibitive in high-dimensional regimes. This paper introduces the Federated Sketch Contextual Linear Bandits (FSCLB) algorithm, which leverages matrix sketching and efficient SVD-based updates to drastically reduce both computation and communication costs, while retaining strong regret guarantees.

Algorithmic Contributions

FSCLB addresses the scalability bottlenecks of FCLB with three main innovations:

1. SVD-based Determinant Computation via Sketching: Instead of performing $O(d^3)$ determinant computations on full covariance matrices, FSCLB maintains a compact $l \times d$ sketch matrix (with $l \ll d$) constructed using the Spectral Compensation Frequent Directions (SCFD) algorithm. The determinant is approximated using the product of singular values from the sketch, reducing per-round complexity to $O(l^2 d)$. Algorithmic checks guarantee that, for sufficiently small $l/d$, SVD-based computation remains asymptotically dominant over direct determinant calculation.
2. Double Sketch Communication Protocol: FSCLB introduces a double sketching mechanism that reduces both upload and download bandwidth from $O(d^2)$ to $O(ld)$. Locally, each agent maintains and uploads only the low-dimensional sketch and accumulated spectral tail. Upon aggregation, the server merges agent sketches and applies a second round of sketching to maintain a consistent $O(d^2)$0 dimensionality for download, thus avoiding blowup in merged matrix rank and maintaining strictly subquadratic communication.
3. Regret-Optimal and Asynchronous Triggers via SCFD: To enable asynchronous update triggers (communicate only upon sufficient covariance growth), FSCLB employs SCFD. Unlike classical Frequent Directions (which can cause non-monotonicity and invalidate asynchronous triggers), SCFD preserves the necessary monotonic growth by compensating with accumulated spectral mass. Agents trigger upload when the determinant computed from the (compensated) sketch exceeds a threshold; theoretical analysis confirms that this rule preserves the optimality guarantees of FedLinUCB.

Theoretical Analysis

Rigorous analysis establishes that, with proper choice of sketch size $O(d^2)$1 and communication tuning parameter $O(d^2)$2, FSCLB achieves cumulative regret

$O(d^2)$3

where $O(d^2)$4 quantifies the spectral approximation error (spectral tail) of the covariance matrix. Notably, if the sketch size $O(d^2)$5 exceeds the rank of the arm-selection matrix, $O(d^2)$6 and the regret reduces to $O(d^2)$7, matching the optimal no-sketch bound of FedLinUCB. Communication and computation costs per communication round are reduced from $O(d^2)$8 and $O(d^2)$9 to $d \times d$0 and $d \times d$1 respectively. The total number of communications is logarithmic in $d \times d$2 and inversely proportional to $d \times d$3.

Strong technical results in the appendix link the spectral properties of merged sketches to the original covariance matrix, showing that the accumulated sketch error is tightly bounded by the spectral tail of the data.

Empirical Evaluation

FSCLB is evaluated on both synthetic (dimensions $d \times d$4 and $d \times d$5, sketch sizes $d \times d$6) and real-world high-dimensional datasets (SatImage, $d \times d$7; MFeat, $d \times d$8) from OpenML. Numerical comparisons are made with FedLinUCB and a random baseline. Principal findings include:

• Over 90% reduction in both computation and communication costs relative to FedLinUCB for comparable cumulative rewards.
• Marginal reward degradation due to sketching, confirming that spectral error does not significantly degrade bandit performance at practical $d \times d$9 values.
• For real-world binary feedback, FSCLB's regret remains statistically indistinguishable from FedLinUCB, with random baselines lagging substantially behind.

These results hold robustly across both simulated and real data, even under adversarial and sensitive arm-selection settings.

Implications and Future Directions

This work demonstrates that randomized and deterministic matrix sketching, specifically SCFD, is a powerful tool for federated online decision-making under strict system constraints. The FSCLB architecture decouples the cost of exploration–exploitation trade-offs from the ambient data dimensionality, facilitating practical deployment of contextual bandits on massive and distributed datasets where communication and computation are primary bottlenecks.

Theoretically, the results motivate further study into more adaptive sketching strategies (e.g., dynamically tuning $O(d^3)$0 per client or per communication epoch), tighter regret bounds under data-dependent sketching, and more aggressive quantization or compression protocols atop the sketch framework. Practically, the double-sketch protocol provides a template for broader classes of federated convex and non-convex optimization problems, particularly for applications in federated recommendation, personalized advertising, and online collaborative filtering—domains where large-scale, heterogeneous, and privacy-sensitive data are the norm.

Conclusion

The FSCLB algorithm achieves significant advances in scaling federated contextual linear bandits by leveraging advanced matrix sketching and a carefully engineered double-sketch communication protocol. By reducing computation and communication overheads by more than 90% with negligible regret increase, FSCLB makes high-dimensional federated contextual bandits tractable for large-scale real-world systems without sacrificing statistical efficiency. The algorithmic and analytical techniques introduced in this paper set a strong precedent for efficient, communication-aware federated learning in online settings.