1
Federated Martingale Posterior Sampling
arXiv:2605.18554v1 [cs.LG] 18 May 2026
Boning Zhang, Student Member, IEEE, Matteo Zecchin, Member, IEEE, Mingzhao Guo, Student Member, IEEE, Dongzhu Liu, Member, IEEE, Osvaldo Simeone, Fellow, IEEE
Abstract—Federated Bayesian neural networks require fixing a prior on the model parameters together with a likelihood. Eliciting meaningful priors on the weight space of modern overparameterized models is notoriously difficult, and misspecification of either component can severely degrade accuracy and calibration. Motivated by the rapid progress of predictive models such as large language models, the martingale posterior, also known as predictive Bayes, replaces the prior–likelihood pair with a predictive distribution and recovers parameter uncertainty by repeatedly drawing predictive samples and refitting the model. A direct federated implementation, however, would require clients to share the local data sets. This letter proposes federated martingale posterior (FMP) sampling, a one-shot embarrassingly parallel protocol in which each client uploads a small set of trainable data embeddings and the server runs the predictive sampler centrally. Experiments on MNIST, CIFAR-10, and CIFAR-100 show that FMP closely matches the centralized counterpart and significantly improves calibration over consensus-style baselines. Index Terms—Bayesian learning, federated learning, martingale posterior
I. I NTRODUCTION
B
AYESIAN learning provides a principled framework for quantifying predictive uncertainty of machine learning models based on the specification of a prior on the model parameters together with a likelihood [1, 2, 3, 4]. However, it is well known that the resulting inference is sensitive to the misspecification of either the prior or the likelihood [5, 6, 7] and that, for modern overparameterized models, eliciting informative priors on the weight space is infeasible [8]. In contrast, the rapid progress of foundation models, including large language models [9], has shown that, for many data sources of interest, it is much easier to Boning Zhang, Mingzhao Guo, and Dongzhu Liu are with the School of Computing, University of Glasgow, Glasgow G12 8QQ, U.K. (e-mail: [email protected]; [email protected]; [email protected]). Matteo Zecchin is with the Communication Systems Department, EURECOM, 06904 Sophia Antipolis, France (e-mail: [email protected]). Osvaldo Simeone is with the Institute for Intelligent Networked Systems, Northeastern University London, London E1 8PH, U.K. (e-mail: [email protected]). The work of O. Simeone was supported by the European Research Council (ERC) under the European Union’s Horizon Europe Programme (grant agreement No. 101198347), by an Open Fellowship of the EPSRC (EP/W024101/1), and by the EPSRC project (EP/X011852/1). The work of M. Zecchin was supported by the Huawei France-funded Chair towards Future Wireless Networks.
Fig. 1. In the proposed FMP protocol, each client compresses its private local dataset Zm into a compressed set Zem = hϕ (Zm ), which is uploaded to the server in one communication round. The server generates predictive samples via induced self-attention blocks (ISAB), and solves the ERM problem to obtain approximate global MP samples.
obtain a powerful predictor of future observations than to specify a meaningful prior on a parameter or function space. This observation has motivated the martingale posterior (MP) [10], also known as predictive Bayes [11], which starts from a joint predictive distribution over unseen data, and recovers samples from an implicit posterior over parameters by repeatedly drawing predictive samples and fitting a parametric model via empirical risk minimization (ERM). Bayesian learning has also found applications to federated settings [12, 13, 14]. However, a naive distributed implementation of MP would require each client to transmit the entire local data set to the server in order for the latter to be able to reconstruct the MP. This letter addresses this issues by proposing federated MP (FMP) sampling, a one-shot embarrassingly parallel protocol that approximates the centralized MP from compressed local data. The main contributions can be summarized as follows. • We formulate the problem of one-shot federated sampling from a global MP induced by a shared settransformer predictor [15, 16], and introduce as a baseline consensus federated MP (CFMP), obtained by directly applying consensus Monte Carlo [13] to local MP samples. • We propose FMP, in which each client compresses its local dataset into a small set of trainable embeddings via an attention-based pooling block and uploads only this compressed set; the server aggregates the embed-
2
dings and runs the centralized predictive sampler on the resulting summary (see Fig. 1). • We design a meta-training procedure that aligns the FMP samples with those of the centralized MP across a corpus of related tasks, and validate the resulting protocol on MNIST, CIFAR-10 and CIFAR-100 under both homogeneous and heterogeneous client partitions. II. BACKGROUND In order to provide the necessary background, we start by reviewing the martingale posterior (MP) [10], together with its attention-based implementation in [16]. 1) Conventional Bayesian Learning vs Martingale Posterior: Traditional Bayesian neural networks (BNNs) specify a likelihood p(y|x, θ) parameterized by weights θ ∈ Θ, and impose a prior distribution p(θ) over θ [3]. Given a dataset Z = {zi }ni=1 with data points zi , the posterior distribution over model parameters is given by p(θ|Z) ∝ p(θ)
n Y
p(zi |θ) .
(1)
i=1
In practice, approximate inference techniques such as variational Bayes [2] or MCMC [1] are commonly used to approximate p(θ|Z) or to draw approximate samples from it. The MP [10] is an alternative formulation of Bayesian theory that reframes posterior uncertainty about parameters θ as predictive uncertainty on unseen, hypothetical data conditional on the observed data. Formally, the MP specifies a joint predictive density over unseen data ′ Z ′ = {zi′ }ni=1 , denoted by p(Z ′ |Z). Given samples Z ′ ∼ p(Z ′ |Z), MP then solves the empirical risk minimization (ERM) problem X ℓ(z, θ), (2) θMP = arg min θ
z∈Z∪Z ′
where ℓ(z, θ) is a loss function. Note that the parameters θMP in (2) are random variables due to the stochasticity of the unseen data Z ′ ∼ p(Z ′ |Z). The samples θMP are treated as draws from the underlying implicit martingale posterior. The connection between the conventional posterior (1) and the martingale posterior underlying the sequence of samples θMP in (2) is given by De Finetti’s theorem. The latter states that a conditionally exchangeable sequence Z ′ admits a joint distribution consistent with the conventional Bayesian model (1), extended to include the unseen data, as Z Y p(Z ′ |Z) = p(θ|Z) p(z|θ) dθ, (3) z∈Z ′
provided that p(Z ′ |Z) is exchangeable for the given fixed data Z. The conditional distribution p(Z ′ |Z) is
said to be exchangeable if it is invariant under arbitrary permutations of the indices of the unseen data Z ′ , i.e., p(Z ′ |Z) = p(π · Z ′ |Z) for any permutation π of indices [n′ ], where n′ = |Z ′ |. Following the predictive-resampling construction of martingale posteriors (see, e.g., Theorem 1 in [10]), assuming that the predictive distribution p(Z ′ |Z) is conditionally exchangeable given dataset Z, with the choice ℓ(z, θ) = − log p(z|θ) in (2), due to the consistency of maximum likelihood estimator, the distribution of the samples θMP converges in distribution to the posterior p(θ|Z) as n′ → ∞ under weak regularity conditions. 2) Set Transformer–Based Predictive Distribution: In [16], the conditional distribution p(Z ′ |Z) is implemented using a set Transformer with induced selfattention blocks (ISAB) [15]. Specifically, the generator ′ i.i.d. first samples an i.i.d. base set E = {εi }ni=1 ∼ p(ε) typically from a standard Gaussian distribution p(ε) = N (0, Id ). Each data point zi ∈ Z is processed by a feedforward neural network g(·) as g(zi ) ∈ Rd , forming the set R = {g(zi )}ni=1 of inducing points. Then, ISAB produces samples Z ′ by applying a cascade of two multihead attention blocks (MABs) with no masking. Denote an MAB block by MAB(Q, C), which maps a query set Q to an updated set by attending to a set C, representing both keys and values, via multi-head attention. The first MAB block uses the base set E and the inducing points R to produce an intermediate state H = MAB(E, R), and the second MAB block applies attention once more to generate the outputs Z ′ = MAB(E, H). This architecture can be proved to ′ produce exchangeable outputs Z ′ = {zi′ }ni=1 [15], and its input-output operation is denoted as Z ′ = ISAB(Z, E) = MAB(E, MAB(E, R)),
(4)
where we recall that the inducing points R are a function of the data Z. III. P ROBLEM F ORMULATION AND BASELINE A. Setting: As shown in Fig. 1, we consider a setting with M clients, with the m-th client holding a private dataset of the same size Zm = {zm,i = (xm,i , ym,i )}ni=1 . The clients are connected to a central server. In conventional federated Bayesian learning [12], one fixes a prior distribution p(θ), along with a likelihood function p(Z|θ). In contrast, in this work we fix a predictive model p(Z ′ |Z), and our goal is to develop distributed protocols that allow the server to draw approximate samples (2) from the global martin′ gale posterior induced by SMthe predictive p(Z |Z) given the global dataset Z = m=1 Zm . Specifically, we assume a fixed predictive mechanism described by the set transformer (4) with a given neural
3
network g(·), which generates inducing points, and a given MAB mechanism MAB(·, ·). Following (2), the target posterior samples are given by X MP θZ = arg min ℓ(z, θ), (5) θ
z∈Z∪Z ′
′
where Z is obtained via (4). We are specifically interested in embarrassingly parallel schemes, in which the clients communicate only once to the server. B. Consensus Federated Martingale Posterior Sampling: As a simple baseline approach, one can directly apply the consensus Monte Carlo protocol introduced in [13], obtaining a benchmark that we refer to as the Consensus Federated Martingale Posterior (CFMP). In the CFMP setup, each client m generates local unseen ′ ′ data Zm using the predictive distribution p(Zm |Zm ) based on the local data Zm . Then, the client solves problem (2), obtaining the samples X MP θm = arg min ℓ(z, θ) (6) θ
′ ∪Z z∈Zm m
from the local martingale posterior. MP are transmitted to the server, which The samples θm MP M }m=1 to approxaggregates all the local samples {θm MP imate a sample θZ in (5) from the global martingale posterior based on the full dataset Z. Specifically, CFMP applies the weighted sum !−1 M M X X MP CFMP −1 Σ̂−1 (7) θ̂Z = Σ̂m m θm , m=1
m=1
where the covariance matrix Σ̂m is an estimate of the MP MP MP − ])(θm − E[θm true covariance matrix Σm = E (θm MP MP ⊤ from (6). ]) obtained using past samples θm E[θm IV. F EDERATED M ARTINGALE P OSTERIOR S AMPLING Constructing the global martingale posterior in (5) requires each m-th client to transmit its local dataset Zm to the server. Communicating the entire local datasets, however, would incur substantial communication overhead and possibly violate privacy constraints. To address this problem, we propose FMP, a novel federated learning protocol that replaces all uses of the local data Zm at the server with trainable embeddings. Specifically, FMP meta-trains the embedding function on data from multiple tasks [14]. 1) The FMP Protocol: As illustrated in Fig. 1, the FMP protocol leverages an embedding function hϕ (·) shared among clients, which takes as input a local dataset Zm to produce a compressed dataset Z̃m = {z̃m,i }si=1 with s < n fictitious data points z̃m,i ∈ Rd for i = 1, · · · , s. The hyperparameter s controls the communication overhead, with a smaller value of s
implying a smaller communication load. The design of the embedding function is discussed in Sec. IV-2. The server aggregates the compressed datasets SM as a surrogate for the original dataset Z as Z̃ = m=1 Z̃m . Using this compressed dataset in lieu of the original dataset, the server draws predictive samples Z̃ ′ using the predictive model (4), i.e., Z̃ ′ = ISAB(Z̃, E), where ′ E = {εi }ni=1 is an i.i.d base set as defined in Sec. II-2. Then, the server obtains approximate martingale posteFMP rior samples θZ by solving the problem X FMP θZ = arg min ℓ(z, θ). (8) θ
z∈Z̃∪Z̃ ′
In general, due to the discrepancy between the local datasets {Zm }M m=1 and the compressed datasets FMP {Z̃m }M differs m=1 , the distribution of the samples θZ MP from the martingale posterior samples θZ in (5). To mitigate this issue, we propose a methodology to optimize the parameters ϕ of the embedding function hϕ (·) in Sec. IV-3. 2) Embedding Function: The embedding function is instantiated using the Pooling by Multihead Attention (PMA) block [15]. Building on the MAB block, PMA takes a learnable set of s seed vectors Sϕ ∈ Rs×d as queries, while keys and values are obtained from the local dataset Zm . This is done by applying a per-data point feedforward layer fϕ (·) applied to each element of the dataset Zm . Overall, the compressed set Z̃m = {z̃m,i }si=1 is given by Z̃m = hϕ (Zm ) = MAB(Sϕ , fϕ (Zm )).
(9)
3) Meta-training the Embedding Function: In order to optimize the embedding function hϕ (·), we assume the server has access to a meta-training corpus of tasks {T i } drawn i.i.d. from a distribution p(T ). Each task T i is associated with a realization of client datasets Z i = i M }m=1 . The data in our federated setting represents {Zm a new task T new = {Zm }M m=1 sampled from the same distribution p(T ). Given a per-task base set E i used in the predictive mechanisms (4) to produce the unseen data Z ′ and Z̃ ′ , we define the per-task loss as MP FMP ℓi (ϕ; E i ) = ∥θZ ∥, i − θZ i MP FMP where θZ and θZ are the samples obtained by the i i centralized MP scheme via (5) and by FMP via (8), respectively. For meta-training, we sample K tasks {T i }K i=1 i.i.d. from the distribution p(T ), along with their correspondi.i.d ing i.i.d. base sets {E i }K i=1 ∼ p(ε). The meta-training objective is defined as the empirical average:
L(ϕ) =
K K 1 X MP 1 X i FMP ℓ (ϕ; E i ) = ∥θ i − θZ ∥, i K i=1 K i=1 Z (10)
4
where we recall that the parameters ϕ determine both the set of seed vectors Sϕ and the per-sample transformation fϕ (·). During meta-training, only the parameters ϕ of the embedding function hϕ are updated via gradient descent.
TABLE I ACCURACY (ACC) AND ECE UNDER HOMOGENEOUS CLIENT PARTITIONS ON MNIST, CIFAR-10, AND CIFAR-100 (20- WAY ). Method
CIFAR-10
ECE ↓
ACC ↑
LANN LBNN LMP
0.9517 0.9510 0.9643
0.0712 0.0793 0.0683
0.5743 0.5595 0.6896
ANN BNN MP
0.9724 0.9713 0.9733
0.0515 0.0449 0.0213
CANN CBNN CFMP FMP
0.9689 0.9705 0.9700 0.9705
0.0678 0.0621 0.0527 0.0406
V. E XPERIMENTS AND C ONCLUSIONS 1) Baselines: In this section, we consider three classes of protocols, namely local, centralized, and one-shot federated (embarrassingly parallel) protocols. Within each class, methods are further categorized as frequentist or Bayesian. For all Bayesian methods, we adopt standard ensembling using model parameter samples. Local protocols operate independently at each client using only local data, including the artificial neural network (LANN), Bayesian neural network (LBNN), and martingale posterior (LMP) baselines. Conversely, centralized protocols operate on the pooled dataset formed by all clients, yielding the ANN, BNN, and MP baselines [10]. Finally, one-shot federated protocols include the consensus artificial neural network (CANN) as the frequentist baseline, in which the server constructs a global model through a single averaging step over the locally trained parameters [17], as well as the consensus Bayesian neural network (CBNN), in which the server combines local posterior samples using consensus Monte Carlo [13], CFMP, and the proposed FMP (see Sec. IV). 2) Classification Tasks: We evaluate all methods on MNIST, CIFAR-10, and a 20-class subset of CIFAR100, allowing us to assess performance across varying levels of task complexity. We first consider a setting in which all clients share the same label space and data are evenly distributed across clients. We then study a heterogeneous setting in which data are partitioned using a Dirichlet distribution with concentration parameter α ∈ {0.1, 0.5, 1.0, 5.0} [18]. Rather than operating on raw images, we first map each input image x to a frozen feature representation h ∈ Rdx , obtained from a feature extractor trained offline with cross-entropy loss [19]. 3) Experimental Results: Table I reports the classification accuracy (ACC) and expected calibration error (ECE) [20] under homogeneous client partitions on MNIST, CIFAR-10, and CIFAR-100 datasets. Among centralized methods, MP achieves the best overall performance in terms of both accuracy and calibration. Among one-shot federated methods, FMP closely matches the centralized MP with only a small performance gap, while outperforming other federated baselines. The advantage becomes more pronounced on more challenging datasets such as CIFAR-10 and CIFAR-100. Notably, FMP achieves substantially lower ECE across all datasets, indicating better alignment between predictive confidence and empirical accuracy. In contrast, parameterspace aggregation methods (e.g., CANN and CBNN)
MNIST ACC ↑
CIFAR-100
ECE ↓
ACC ↑
ECE ↓
0.0845 0.0820 0.1053
0.5736 0.5626 0.6625
0.1804 0.1842 0.1990
0.0729 0.0593 0.0342
0.7030 0.7010 0.7095
0.0503 0.0385 0.0274
0.2130 0.1606 0.1422 0.0572
0.6840 0.6950 0.6973 0.7032
0.3158 0.2458 0.2221 0.0423
Local
Centralized 0.7863 0.7836 0.8057 Federated 0.7406 0.7357 0.7740 0.7891
Fig. 2. Accuracy (ACC) and ECE under heterogeneous client partitions with Dirichlet parameters α ∈ {0.1, 0.5, 1.0, 5.0}. Left: ACC versus Dirichlet α. Right: ECE versus Dirichlet α.
exhibit noticeably worse calibration despite achieving competitive accuracy. Figure 2 reports the accuracy and ECE results under heterogeneous client partitions. Across all settings defined by the heterogeneity parameter α, FMP consistently achieves the best overall performance among the federated methods, attaining the highest accuracy together with the lowest ECE. Its advantage is most pronounced under severe heterogeneity, i.e., α = 0.1. In contrast, parameter-space aggregation methods such as CANN and CFMP achieve competitive accuracy at larger values of α, but remain noticeably worse calibrated. 4) Conclusions: Overall, FMP closely tracks the centralized MP benchmark and consistently improves calibration over consensus-style federated baselines, especially under heterogeneous client partitions. Future work may study formal privacy guarantees, adaptive communication-accuracy tradeoffs, and variants that are robust to communication errors. R EFERENCES [1] C. P. Robert and G. Casella, Monte Carlo Statistical Methods. Springer, 1999, vol. 2.
5
[Online]. Available: https://doi.org/10.1007/978-1 -4757-3071-5 [2] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe, “Variational inference: A review for statisticians,” J. Amer. Statist. Assoc., vol. 112, no. 518, pp. 859–877, 2017. [Online]. Available: https: //doi.org/10.1080/01621459.2017.1285773 [3] W. Xu, A. Liu, Y. Zhang, and V. Lau, “Bayesian deep learning via expectation maximization and turbo deep approximate message passing,” IEEE Trans. Signal Process., vol. 72, pp. 3865–3878, 2024. [Online]. Available: https://doi.org/10.1109/ TSP.2024.3442858 [4] O. Simeone, Machine Learning for Engineers. Cambridge University Press, 2022. [Online]. Available: https://www.cambridge.org/highered ucation/books/machine-learning-for-engineers/7 FD8622836CAFCF5EDB169E7DC8A1ED4 [5] P. G. Bissiri, C. C. Holmes, and S. G. Walker, “A general framework for updating belief distributions,” J. Roy. Statist. Soc. Ser. B, vol. 78, no. 5, pp. 1103–1130, 2016. [Online]. Available: https://doi.org/10.1111/rssb.12158 [6] J. Knoblauch, J. Jewson, and T. Damoulas, “An optimization-centric view on bayes’ rule: Reviewing and generalizing variational inference,” J. Mach. Learn. Res., vol. 23, no. 132, pp. 1–109, 2022. [Online]. Available: https://jmlr.org/papers/ v23/19-1047.html [7] M. Zecchin, S. Park, O. Simeone, M. Kountouris, and D. Gesbert, “Robust pac m: Training ensemble models under misspecification and outliers,” IEEE Trans. Neural Netw. Learn. Syst., vol. 35, no. 11, pp. 16 518–16 532, 2023. [Online]. Available: https://doi.org/10.1109/TNNLS.2023.3295168 [8] S. Sun, G. Zhang, J. Shi, and R. Grosse, “Functional variational Bayesian neural networks,” in Proc. Int. Conf. Learn. Represent. (ICLR), 2019. [Online]. Available: https://openreview.net/forum?i d=rkxacs0qY7 [9] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein et al., “On the opportunities and risks of foundation models,” arXiv:2108.07258, 2021. [Online]. Available: https: //arxiv.org/abs/2108.07258 [10] E. Fong, C. Holmes, and S. G. Walker, “Martingale posterior distributions,” J. Roy. Statist. Soc. Ser. B, vol. 85, no. 5, pp. 1357–1391, 2023. [Online]. Available: https://doi.org/10.1093/jrsssb/qkad005 [11] M. Battiston and L. Cappello, “Bayesian predictive inference beyond martingales,” arXiv:2507.21874, 2025. [Online]. Available: https://arxiv.org/abs/25 07.21874 [12] R. Kassab and O. Simeone, “Federated generalized
bayesian learning via distributed stein variational gradient descent,” IEEE Trans. Signal Process., vol. 70, pp. 2180–2192, 2022. [Online]. Available: https://doi.org/10.1109/TSP.2022.3168490 [13] S. L. Scott, A. W. Blocker, F. V. Bonassi, H. A. Chipman, E. I. George, and R. E. McCulloch, “Bayes and big data: The consensus monte carlo algorithm,” Int. J. Manag. Sci. Eng. Manag., vol. 11, no. 2, pp. 78–88, 2016. [Online]. Available: https://doi.org/10.1080/17509653.2016.1142191 [14] M. Zhu, M. Zecchin, S. Park, C. Guo, C. Feng, and O. Simeone, “Federated inference with reliable uncertainty quantification over wireless channels via conformal prediction,” IEEE Trans. Signal Process., vol. 72, pp. 1235–1250, 2024. [Online]. Available: https://doi.org/10.1109/TSP.2024.33586 15 [15] J. Lee, Y. Lee, J. Kim, A. Kosiorek, S. Choi, and Y. W. Teh, “Set transformer: A framework for attention-based permutation-invariant neural networks,” in Proc. Int. Conf. Mach. Learn. (ICML), 2019, pp. 3744–3753. [Online]. Available: https://proceedings.mlr.press/v97/lee19d.html [16] H. Lee, E. Yun, G. Nam, E. Fong, and J. Lee, “Martingale posterior neural processes,” in Proc. Int. Conf. Learn. Represent. (ICLR), 2023. [Online]. Available: https://openreview.net/forum?i d=-9PVqZ-IR [17] M. Zinkevich, M. Weimer, L. Li, and A. Smola, “Parallelized stochastic gradient descent,” Adv. Neural Inf. Process. Syst., vol. 23, 2010. [Online]. Available: https://papers.nips.cc/paper/4006-paralle lized-stochastic-gradient-descent [18] T.-M. H. Hsu, H. Qi, and M. Brown, “Measuring the effects of non-identical data distribution for federated visual classification,” arXiv:1909.06335, 2019. [Online]. Available: https://arxiv.org/abs/19 09.06335 [19] J. Donahue, Y. Jia, O. Vinyals, J. Hoffman, N. Zhang, E. Tzeng, and T. Darrell, “Decaf: A deep convolutional activation feature for generic visual recognition,” in Proc. Int. Conf. Mach. Learn. (ICML), 2014, pp. 647–655. [Online]. Available: https://proceedings.mlr.press/v32/donahue14.html [20] C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger, “On calibration of modern neural networks,” in Proc. Int. Conf. Mach. Learn. (ICML), 2017, pp. 1321–1330. [Online]. Available: https://proceeding s.mlr.press/v70/guo17a.html