Universal Sample Coding

Authors: szymon  kobus
deniz  gündüz
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 12:33:53

Abstract: In this work, we study the problem of communicating multiple samples from an unknown probability distribution using as few bits as possible. This is a generalization of the channel simulation problem, which has recently found applications and achieved state of the art results in realistic image compression, neural network compression, and communication-efficient federated learning. In this problem, the transmitter wants the receiver to generate multiple independent and identically distributed (i.i.d.) samples from a target distribution $P$, while the transmitter and the receiver have access to independent samples from a reference distribution $Q$. The core idea is to employ channel simulation in multiple rounds while updating the reference distribution $Q$ after each round in order to reduce the KL-divergence between $P$ and $Q$, thereby reducing the communication cost in subsequent rounds. We derive a lower bound on the expected communication cost and construct a practical algorithm that achieves the lower bound up to a multiplicative constant. We then employ this algorithm in communication-efficient federated learning, in which model updates correspond to samples from a distribution, and achieve a 37% reduction in the communication load. To further highlight the potential of sample communication for generative models, we show that the number of bits needed to communicate samples from a large language model can be reduced by up to 16 times, compared to entropy-based data compression.

Main paper: 
 introduction: 
 Let P be a probability distribution known to an encoder, and we would like to obtain a sample X ∼ P at a remote decoder -what is the minimal number of bits required to communicate from the encoder to the decoder? The obvious procedure is to sample from P at the encoder, compress it, and transmit over the network. The minimum number of bits necessary for lossless transmission of this sample is bounded by the entropy of distribution P , H(P ) bits per sample. However, unlike in classical source coding theory, in our problem, the decoder is interested only in drawing of 'an arbitrary' sample from P , and not the particular sample generated at the encoder. If the encoder and decoder have access to samples from another reference distribution Q, then it is possible for the encoder to communicate a sample from P using approximately D KL (P ∥Q) bits instead. Thus, if we are able to sample from a distribution Q that is close to P , the bandwidth required to communicate a sample from P can be greatly reduced. This is known as 'channel simulation' or 'reverse Shannon coding' (Bennett et al., 2002; Cuff, 2008; Harsha et al., 2010) . This result has recently been used in various neural compression problems, such as compression of images and neural networks (Havasi et al., 2019; Flamich et al., 2020; Agustsson and Theis, 2020; Theis et al., 2022) , where non-differentiable quantization operation is replaced with reverse channel coding, which provides a differentiable step by utilizing the reparameterization trick. Another recent application of channel simulation is in federated learning (FL). Many communication efficient FL methods require sending a sample from a client-only distribution after local training. Instead of deterministic weights, the models can be parameterized with a probability distribution and the central server samples realizations from each client's updated model to construct the global model update. Recently, it was shown in Isik et al. (2023 Isik et al. ( , 2024) ) that the overall communication cost of sending the model updates from the clients to the parameter server can be greatly reduced using channel simulation, achieving state-of-the-art results. In each iteration, the client can enable the parameter server to sample from its locally updated distribution while using the global model distribution as the common reference distribution. As the client's local distribution is typically close to the global model, the communication cost, proportional to the KL-divergence between the two, will will extremely low compared to deterministic alternatives. Channel simulation can also help in enabling differential privacy of the communicated models (Triastcyn et al., 2021; Hasircioglu and Gunduz, 2024; Hegazy et al., 2024; Shahmiri et al., 2024) since it allows for shaping the distribution of the noise. In all the above works, the single-shot version of channel simulation is utilized; that is, a single sample is to be generated at the receiver with the desired distribution. In this work, we are interested in studying the communication cost of sending multiple samples from the same distribution. This problem can directly find an application in the FL framework, where the global model is estimated from samples of local models. A procedure to efficiently communicate more samples would allow the parameter server to approximate the global model more accurately, which would reduce the noise and variance in its estimation, and speed up training.Additionally, we would like to point to a strong conceptual connection between generative models and channel simulation, and propose an avenue of future research. Generative models are one of the most successful examples among the recent advances in machine learning. They include various modalities such as text (Touvron et al., 2023) , images (Ho et al., 2020) , video, and many more. They are trained to mimic some data distribution P based on samples X ∼ P drawn from it, and allow to generate novel samples from the learned distribution. With growing popularity and deployment of generative AI across a wide variety of applications, communicating the outputs of these models (particularly for image, audio and video modalities) will put an increasing burden on the underlying communication network infrastructure. However, instead of generating and then compressing a sample at the cloud server, we can exploit channel simulation to enable the users to locally generate samples at a much lower communication cost. Drawing multiple samples is a common use case for generative models. For instance, in text to image models, such as DALL-E or Midjourney, multiple images are often generated and communicated to the user based on a single prompt to allow the user choose the desired one. Other instances include translation (Lee et al., 2021; Eikema and Aziz, 2022) , code generation (Chen et al., 2021) and planning (Yao et al., 2023) , which can be improved by discriminating among multiple generated samples.The contributions of this work are summarized as follows:• We formulate a novel universal sample coding problem, which defines the problem of communicating multiple samples from an arbitrary discrete distribution unknown to the receiver.• We highlight the relationship between the redundancy in universal source coding and the communication of samples, and show that the optimal redundancy for universal source coding is a lower bound on the communication cost of universal sample coding.• We propose a coding scheme for the universal sample coding problem, providing an upper bound on its communication cost that is within a multiplicative factor from the lower bound.• We employ our algorithm in FL achieving a 37% reduction in the communication load compared to the current state-of-the-art communication-efficient FL algorithm.• We adapt our algorithm for the remote generation problem, where the goal is to sample at a remote user from a generative model hosted on a cloud server, while communicating the least number of bits from the server to the user. We demonstrate the potential benefits of sample communication in this setting through numerical experiments.  
notation: 
 We use X n = (X 1 , X 2 , . . . , X n ) to denote a sequence of n elements and log(•) denotes logarithms of base 2. Let P and Q denote discrete probability distributions. The entropy of P is given byH(P ) ≜ -E x∼P [log P (x)], while the relative entropy from P to Q, or the KL-divergence, is defined asD KL (P ∥Q) ≜ E x∼P log P (x) Q(x) .3 BackgroundThe problem of communicating samples from a desired distribution is a version of the channel simulation problem, also known as the reverse Shannon theorem, or reverse channel coding. Given a sample z from a probability distribution P Z at the encoder, channel simulation entails generating sample x from the conditional distribution P X|Z=z at the decoder, using the fewest number of bits. It was posed in Cuff (2008) , where the solution was characterized in the asymptotic regime, where an infinite sequence of samples is considered, while earlier results appeared in the quantum communication literature (Bennett et al., 2002; Winter, 2002) . The non-asymptotic results were shown in Harsha et al. (2010) using common randomness, and then further refined in Li and Anantharam (2021) , where it was shown that to communicate a sample X from distribution P X|Z=z with joint distribution P XZ , it is sufficient to transmitEQUATIONon average, whereEQUATIONis the mutual information between X and Z. This result is close to optimal as it was shown in Li and El Gamal (2018) , i.e., there exist distributions P X,Z where the minimal number of communication required for reverse channel coding is on average I(X; Z) + log(I(X; Z) + 1) -1 bits.(3)As pointed out and shown in Theis and Yosri (2022) , the sample communication bound does not rely on using the exact marginal distribution P X , and still holds for other reference distributions Q. In general, to generate a sample at the decoder from distribution P , with reference distribution Q, it is enough to communicate an average of D KL (P ∥Q) + log(D KL (P ∥Q) + 1) + 4.732 bits.(4)While efficient in terms of communication cost, the sample communication method in Li and El Gamal (2018) might be prohibitively expensive computationally. There have been increasing efforts in creating algorithms with reduced complexity by constraining the distributions P and Q (Flamich et al., 2022; Flamich, 2023) . An overview of the methods and trade-off between communication, computation, and sample accuracy is provided in Theis and Yosri (2022) .Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The paper addresses its limitations in section 8, highlighting constraints on applying universal sample coding to continuous variables and scalability issues for large distributions.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: The paper includes complete proofs with assumptions for theoretical result. Proofs and necessary theoretical details are provided in section 5 and detailed further in the supplemental material under B. Each theorem and lemma is properly numbered and cross-referenced. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
source coding: 
 Lossless source coding is the problem of describing the realizations of a random variable with the least number of bits on average. Each outcome x ∈ X is assigned a different sequence of bits, called a codeword, with length l(x). The optimal average codeword length is obtained by Huffman coding (Huffman, 1952) , where l(x) ≤ ⌈-log P (x)⌉ andH(P ) ≤ E x∼P [l(x)] ≤ H(P ) + 1.(5)For n independent and identically distributed random variables X n = (X 1 , X 2 , . . . , X n ), X i ∼ P , it is straightforward to show that using a Huffman code, the expected codeword length per symbol isH(P ) ≤ 1 n E x n ∼P ⊗n [l(x n )] ≤ H(P ) + 1 n .As n → ∞ this quantity converges to the entropy, which provides a fundamental limit for compression.We call the code mismatched if it is designed for a source distribution of Q while the true source distribution is P . Then the expected codeword length is given by:EQUATION)For the n-ary case, the rate converges to H(P ) + D KL (P ∥Q) (Cover and Thomas, 2006) , where the KL-divergence is the penalty for using a mismatched source distribution when designing the code.However, there exist source coding methods called universal source coding with an average codeword length that converges to H(P ) for any underlying distribution P not known to the encoder. They can be thought of as empirically estimating (often implicitly) the underlying distribution of the data.The classic and widely used example of such an algorithm is LZ78 (Ziv and Lempel, 1978) . For any distribution P the average number of bits needed to compress n samples is H(P ) + O 1 log n (Savari, 1997) , which converges much slower than Huffman coding for known distributions. This excess factor for any universal source code Γ is known as per letter redundancy, defined as:EQUATIONwhere P k is a k-dimensional probability simplex. It was shown in Davisson et al. (1981) that for a k-dimensional distribution, the minimal per letter redundancy for any universal source code is:EQUATION)The log n factor is the penalty for universality of the code (Krichevsky and Trofimov, 1981) . A more useful quantity for our analysis will be the total redundancy, defined as n r Γ . In the remainder of the paper, when we refer to redundancy, we mean the total redundancy.As mentioned in Section 3, reverse channel coding involves simulating samples from a joint distribution. Sample communication is a similar problem where the aim is to draw samples at the decoder from distribution P , which is only known at the encoder, while both the encoder and the decoder have access to samples from a common distribution Q. This common randomness -in practice, a seed to initialize a pseudorandom number generator -allows both to draw the same sequence of samples from Q. Using any sample communication method with cost specified in Equation (4) a sequence of n independent and identically distributed (i.i.d.) samples X n , X i ∼ P , can be encoded using: D KL (P ⊗n ∥Q ⊗n ) + log(D KL (P ⊗n ∥Q ⊗n ) + 1) + 4.732 = n D KL (P ∥Q) + log(n D KL (P ∥Q) + 1) + 4.732 (10) expected number of bits, and the per sample cost converges to D KL (P ∥Q) as n → ∞.Characterizations of the rate of source coding and sample communication in both matched and mismatched cases are presented in Table 1 , where we can see that 'moving' to source coding adds H(P ) to the required rate, while using a mismatched distribution Q contributes D KL (P ∥Q). The natural question to ask is whether the results of universal source coding can be extended to sample communication. Firstly, does the rate approach 0 when n → ∞, and secondly, what is the total communication cost of communicating n samples? In the next section, we show that the answer to the first question is positive, and that the total communication cost is in the same order as the minimal redundancy, and the two quantities share a common lower bound.Input: P, n ≥ 1 Parameter: c > 0 1: Send sample X 1 ∼ P . 2: Let i = 2, G(1) = 1. 3: while i ≤ ⌈log 1+c (n)⌉ + 1 do 4: Let G(i) = ⌈(1 + c) i-1 ⌉ 5: Let Q i = Qi (X G(i-1) ) be probability estimated with all previous samples (estimator Lemma 5.2).Let g(i) = G(i) -G(i -1).Jointly send g(i) samples XG(i) G(i-1)+1 ∼ P ⊗g(i) using Q ⊗g(i) i(Eqn. (10)).8:i = i + 1 9: end while 5 Universal sample codingThe core idea is that, after observing t samples, the decoder can estimate P , and use this estimate as the reference distribution Q. As more samples are transmitted, the KL-divergence between the decoder's estimate and the true distribution P will diminish, which will translate into a lower communication cost for later samples. The proposed scheme consists of multiple rounds; in each round, a batch of samples are sent jointly, then the reference probability Q is updated. Theorem 5.1 (Universal sample coding). There exists a randomized encoding function f : P k × Z → N, and a randomized decoding function g : N ×Z → X n , such that for any k-dimensional discrete distribution P ∈ P k over alphabet X and a random stringZ ∈ Z = {0, 1} ∞ , such that g(f (P, Z), Z) = X n ∼ P ⊗n and E H(f (P, Z)|Z) = L(n) ≤ V k (c) log(n) + o log(n) ,where the expectation is over all random strings Z, andV k (c) ≜ c ln(1 + c) k -1 2 + ln k-1 2 ln 2 + 1 + 5 ln 2 ln(1 + c) + 1 (11)for any c > 0.The random string Z is the common randomness shared between the encoder and decoder. The expected number of bits needed to communicate n samples from the any distribution P is similar to the redundancy of universal source coding, with some additional factors and per sample costO log n n → 0 as n → ∞.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: The paper provides descriptions for reproducing its experiments. It follows the setup for federated learning from Isik et al. (2024) and specifies any modifications in section 7. For the LLM experiments, which involve only inference, all necessary details are available in appendix D. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Answer: [Yes] Justification: For the federated learning experiments, we report maximum standard deviations for all metrics, which are low, indicating consistency across runs. For experiments involving large language models and the direct evaluation of universal sample coding, we include data for the top and bottom 20th or 25th percentiles.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
probability estimation: 
 As the samples are drawn independently, their order does not reveal any information about P . Thus, any estimator, without loss of optimality, can be based only on the count of each symbol in the observed sequence. One of the canonical estimators is the add-1 or Laplace estimator. As the name suggests, the estimated probability of any symbol is its count in the sequence plus 1, normalized to form a probability distribution. Laplace estimator belongs to a family of add-β estimators, which rely on adding a constant β to each count. Although no add-β estimator achieves the optimal minmax KLdivergence rate (Krichevskiy, 1998) , there exists an estimator (Braess and Sauer, 2004 ) that combines add-1 2 , add-3 4 and add-1 estimators, and achieves the optimal and universal decay of KL-divergence. Lemma 5.2. (Braess and Sauer, 2004) ] Given n i.i.d. samples X n , X i ∼ P from any distribution P ∈ P k , there exists an estimator Q(X n ) such that: EQUATIONc V k (c) k 1 2 Figure 1:The optimal universal sample coding factor inf c V k (c), optimized over the choice of c > 0 for k-dimensional distributions, along with the lower bound factor k-1 2 .2 4 8 16 32 64 128 2 8 2 9 2 10 2 11 k (dimension)1 2 3 4 6 8 10 2 k 1 inf c V k (c) 1 Figure 2: The optimal universal sample coding factor inf c V k (c), normalized by k-12 , optimized over the choice of c > 0 for k-dimensional distributions.If we could communicate a single sample with D KL (P ∥Q) average bits, then alternating between sending a single sample and estimating P would be the best strategy. The algorithm would take n communication rounds and the total expected communication cost of this hypothetical algorithm would be:EQUATIONwhere ⌈log k⌉ is the cost of sending the first sample. However, the real cost of sending a sample, as shown in (4), includes a constant, which would dominate the total cost as n increases, making it linear instead of logarithmic. We propose a solution that requires only O(log n) communication rounds, where in each round, exponentially more samples are sent. Algorithm 1 illustrates the pseudo code for the proposed solution. The proof of Theorem 5.1 is presented in Appendix B. We first upper bound the number of bits communicated in each round. Then we show that the total number of bits communicated over all rounds is V k (c) log(n) + o(log n), where parameter c controls the size of the groups communicated at each round. The optimal choice of c depends on k.We have plotted the infimum ofV k (c) for different values of k in Figure 1. The value of the first term of V k (c), that is c ln(1+c) k-12 , is minimized as c approaches 0 and converges to k-1 2 , which is equal to the optimal redundancy factor. As we show in the next section, it is also the lower bound for any universal sample communication algorithm. The ratio between the optimal value of V k (c) and k-1 2 is shown in Figure 2 , where it starts at around 9 and converges to 1 as k grows; this is the multiplicative gap between the upper bound on the communication cost of the proposed algorithm and the lower bound presented in the next section.The connection between universal source coding and universal sample coding is highlighted in the following theorem, where the lower bound of redundancy of universal source coding is the same as the communication cost of universal sample coding. This is not a coincidence, we can prove this theorem using exactly the same steps as the proof of optimality for the redundancy of universal source coding in Davisson et al. (1981) . Theorem 5.3 (Universal sample coding -lower bound). For any universal sample coding algorithm (Theorem 5.1), there exists k-dimensional distribution P , such that, the expected number of bits L(n) required to communicate n samples from P satisfies: Proof. Theorem 5.3 states that the claim holds for some distribution P , thus it holds for the supremum over all k-dimensional distributions sup P L(n). We will consider P as a random variable distributed according to Ω. Let Ω be any distribution of k-dimensional distributions, then supEQUATIONEQUATIONsince the maximum is always greater than or equal to the average. The right hand side in Equation ( 15) corresponds to the expected number of bits required to communicate a sample X n from P ⊗n , where P is sampled from Ω. This is exactly the reverse channel coding problem, where the lower bound on the communication cost was shown in Harsha et al. (2010) to be the mutual information between the two random variablesEQUATION)This quantity was bounded in Davisson et al. (1981) asEQUATION)To show (17), Davisson et al. (1981) considers a Markov chain P → X n → Q, where Q is a reconstruction of P based on samples X n . By data processing inequality, this Markov chain implies I(P, X n ) ≥ I(P, Q). The quantity I(P, Q) describes the minimum amount of information needed to compress P in lossy source coding, and can be bounded using Shannon lower bound Shannon (1959) . The final lower bound is obtained by choosing an appropriate distribution Ω and a distortion measure between distributions P and Q. 
empirical evaluation: 
 To corroborate the claims about universal sample coding made in Theorem 5.1, we compare the KL-divergence and total communication cost between the theory and numerical evaluations for different number of samples and dimensions. Each experiment is repeated 1000 times, where a probability distribution P is sampled from a k-dimensional Dirichlet distribution with concentration parameters α ∈ R k , α i = 1. Then, n samples are communicated using universal sample coding (Algorithm 1). We consider ordered random coding (Theis and Yosri, 2022) , shown in Appendix C, as the underlying sample communication method (line 7, Algorithm 1). In Figure 3 , we observe that the measured KL-divergence D KL (P ∥ Q) between the true probability P and the estimate Q from the estimator in Lemma 5.2 follows the predicted values across a range of communicated samples.In Figure 4 , we fix the number of samples to n = 128, but vary the dimension k, and show that the KL-divergence lies below the bound from Lemma 5.2. In Figure 5 , we plot the total communication cost for various number of samples, which is contained between the asymptotic upper and lower bounds, V k (c) log(n) (Theorem 5.1) and k-1 2 log(n) (Theorem 5.3), respectively. 
federated learning (fl): 
 To show the efficacy of the universal sample coding in practice, we first apply it to the Federated Probabilistic Mask Training (FedPM) algorithm proposed in Isik et al. (2024) . In FedPM, the weights of the neural network w ∈ R M are randomly initialized, fixed, and known both to the clients and the central server. Each weight w i has an associated probability θ i of that weight being masked. The training of the network consists of finding the suitable mask probabilities θ ∈ [0, 1] M for all the weights using gradient descent. To test the network, a binary mask is sampled, and the effective weights of the network become w i * Bern(θ i ). In a single learning round, the server broadcasts the global mask θ to each client, which then trains an updated version θ ′ using its local data and communicates a sample Bern(θ ′ ) from its updated mask distribution back to the central server. The new global mask probability θ is then estimated from all the received samples. The sample from P = θ ′ is communicated using the coding probability Q = θ, which is known to both the client and the server, thereby achieving a communication cost of approximately D KL (θ ′ ∥θ). To achieve a more accurate estimate of the global mask probability, we propose communicating multiple mask samples per client using universal sample coding, along with the addition of the global prior θ.For the experiments, we follow the same setup as Isik et al. (2024) However, when only 1 out of 10 clients participate in the learning round the accuracy decreases to 75%. By communicating multiple samples in each round, we can achieve test accuracy close to the fully uncongested case. However, as argued in this work, sending multiple samples from the same distribution can be made more communication efficient by estimating the true probability from previous samples. Indeed, we observe a 37% reduction in communication cost by employing universal sample coding with prior θ, compared to only using the prior θ.To incorporate the prior θ into the proposed coding scheme, we use a Bayesian estimator. Although the mask is binary, we describe a general case with k possible outcomes. Let ω ∈ P k be the prior probability, and ω ′ be a random variable distributed according to the Dirichlet distribution with parameters α ∈ R k , for j ∈ {0, . . . , k -1}: α j = µω j + G l=1 1 {l-th sample = j} , where 1{•} is the indicator function, G is the number of samples communicated so far, and µ ∈ R + is a hyperparameter controlling the reliance on the prior. The optimal coding probability Q = E ω ′ ∼P ω ′ ω ′ , i.e., for j ∈ {0, . . . , k -1}, Q j = αj αi , replaces the one in line 5 in Algorithm 1. For the binary mask case, where k = 2, we have ω = [1 -θ, θ], and θ ′ is characterized by a 2-dimensional Dirichlet distribution (or equivalently, a Beta distribution). We propose two potential avenues for extending universal sample coding for continuous variables. Firstly, by imposing additional assumptions on the probability distribution, we could substitute counting-based estimation with a Bayesian approach. The reference distribution Q would be the posterior, updated continually with each new sample. Second, we could explore a model-based approach where a common model describes Q. This model would be incrementally fine-tuned with the communicated samples to more accurately reflect their underlying distribution. As evidence for the potential of this idea, we apply sample communication for generating samples from a large language model on a server, using a smaller reference model at the client. While the space of text is not continuous, it is too large to estimate using count-based methods. The experimental setup and further discussion are detailed in Appendix D. As demonstrated in Table 3 , the application of sample communication results in a 4-to 16-fold reduction in communication rate, depending on the auxiliary model employed. This motivates further exploration into model fine-tuning effects. Although this research currently focuses on text generation, future work will aim to extend these results to image and video generation, which represent a substantial portion of network traffic. Furthermore, determining the optimal use of new samples in the learning process presents a compelling challenge, potentially linking to advancements in online learning methods.The universal sample communication problem, introduced in this work, involves communicating multiple samples drawn from the same probability distribution P , which is unknown to the decoder. An alternative but related problem would involve generating multiple samples, each from a different conditional distribution P X|Z=z . In such a case, the optimal coding distribution would be the marginal distribution Q = P X , and the optimal per-sample communication cost would be I(X; Z). In the universal setting, the decoder would not know this marginal distribution. Thus, a similar estimation and sample communication approach, as demonstrated in this work, could be applied to address this problem. However, further analysis is required to validate the communication performance of such a scheme. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: The paper provides information on the computational resources used for each experiment in the description of each experiment.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
conclusion: 
 As AI tools become ever more prevalent, it becomes increasingly important to consider the resources they consume. This paper analyzes the communication cost of transmitting samples from probability distributions, which can be seen as sending the outcomes of generative models, and show that by leveraging the fact that many generative AI applications only seek to communicate generic samples from a distribution rather than a specific sample, we can reduce the associated communication cost significantly. In this paper, we proposed universal sample coding for transmitting multiple samples from a distribution, which achieves the information theoretic lower bound by up to a multiplicative constant. We also applied it to a FL framework showing a communication cost reduction of 37%, and to remote text generation problem using large language models, showing an up to 16-fold communication cost reduction compared to sample-wise entropy coding methods.Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The study does not involve human subjects or sensitive data, and all computational methods and experiments comply with ethical standards for responsible AI research and development.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: The paper focuses on improving the efficiency of existing technologies and does not introduce new applications; thus, it does not directly discuss broader societal impacts.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper does not release any data or models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We ensure that all external assets used in our research, such as code and datasets, are properly credited to their original creators.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: The paper does not introduce or release any new assets such as datasets, code, or models; therefore, no documentation related to new assets is provided at this stage. However, when the code is released, full documentation will be provided.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] .Justification: The paper does not involve crowdsourcing or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The form of the universal estimator (Braess and Sauer, 2004 ) is provided in Algorithm 2, where k denotes the size of the alphabet and c[i] represents the number of occurrences of symbol i in the observed sequence. Notably, when the sequence is empty, the estimator produces a uniform distribution over all symbols.Require: c // Count of occurrences of symbols in the sequence.1: Initialize array Q of size k 2: for i ∈ {1, 2, . . . , k} do3: if c[i] = 0 then 4: Q[i] ← c[i] + 1 2 5: else if c[i] = 1 then 6: Q[i] ← c[i] + 1 7: else 8: Q[i] ← c[i] + 3 4 9:end if 10: end for11: W ← k i=1 Q[i]// Normalize the distribution. 12: for i ∈ {1, 2, . . . , k} do 13:Q[i] ← Q[i] / W 14: end for 15: return Q B Universal Sample Coding proofProof. Theorem 5.1: Universal Sample Coding Let c > 0, and G(i) be the number of samples communicated up to and including round i ≥ 1:EQUATION)Let g(i) be the number of samples communicated at round i, g(1) = G(1) = 1, for i ≥ 2:EQUATION)At round i, g(i) samples will be coded jointly, using an estimate Q based on previous G(i -1) samples. The cost of sending the first sample is ⌈log k⌉. For rounds i ≥ 2, using cost from Equation (10) (with the constant upper bounded by 5) at round i, with the estimator from Lemma 5.2, the expected number of communicated bits is upper bounded by:EQUATIONEQUATION+ log c + (1 + c) -i+2 + 2 ln 2 k -1 + 5 + o(1).In line (20), we used Lemma 5.2 and Jensen's inequality. In line ( 21), the inequality g(i) G(i-1) ≤ c + (1 + c) -i-2 for i ≥ 2 is used. In ⌈log 1+c (n)⌉ + 1 rounds, there are G(⌈log 1+c (n)⌉ + 1) ≥ n samples communicated, and the total expected number of communicated bits, denoted by L(n, c), is given by:EQUATIONEQUATION+ ⌈log 1+c (n)⌉-1 i=0 (1 + c) -i k -1 2 ln 2 + log c + (1 + c) -i + 2 ln 2 k -1 ≤ log(n) log(1 + c) c k -1 2 ln 2 + log k -1 2 ln 2 + 5 + 1 + 1 c k -1 2 ln 2 (26) + log(n) 1 + ln(1 + 2 ln 2 k-1 ) ln(1 + c) + log 1 + c + 2 ln 2 k -1 + o(log n) = log(n) ln(1 + c) c k -1 2 + ln k -1 2 ln 2 + ln(1 + 2 ln 2 k -1 ) + 5 ln 2 + log(n) + o(log n) (27) = log(n) c ln(1 + c) k -1 2 + ln k-1 2 ln 2 + 1 + 5 ln 2 ln(1 + c) + 1 + o(log n) (28) = V k (c) log(n) + o(log n) (29)where the upper bound of terms dependent on round i in line (25) lemmas B.1,B.2. This proves theorem 5.1.Lemma B.1. For 0 < c, n ≥ 1:EQUATIONThus,⌈log 1+c (n)⌉-1 i=0 (1 + c) -i k -1 2 ln 2 ≤ c -1 + 1 k -1 2 ln 2 (31)Lemma B.2. For 0 < c, n ≥ 1, and k ≥ 2:EQUATIONThe channel simulation method used throughout this work is ordered random coding from Theis and Yosri (2022) reproduced in Algorithm 3 for convenience.Algorithm 3 Ordered Random CodingRequire: P, Q, N 1: t, n, s ⋆ ← 0, 1, ∞ 2: w = min x P (x)/Q(x) 3: repeat 4: z ← sample P 5: v ← N/(N -n + 1) 6: s ← t • P (z)/Q(z) 7: if s < s ⋆ then 8: s ⋆ ← s 9: n ⋆ ← n 10: end if 11: n ← n + 1 12: until s ⋆ ≤ t • w or n > N 13: return n ⋆The idea behind the universal sample communication is based on two ingredients: samples can be efficiently communicated if the decoder has access to a similar distribution, and the distribution can be estimated based on observed samples. Hence, as the receiver acquires more samples, the marginal communication cost decreases. In machine learning, many models, such as classifiers or generative models, describe (conditional) probability distributions. Thus, this recipe straightforwardly translates to such scenarios, where samples from a model can be communicated in the same way, and the estimation step can be replaced by an auxiliary model that learns based on communicated samples. Specifically, let there be a server and a client, where the client wants to use a model located at the server -for instance an artificial neural network. The model represents a conditional probability distribution P X|Z=z , where z is the input to the model. In addition, both the client and the server have access to an auxiliary model describing Q X|Z=z . The client sends z to the server, and then a sample x ∼ P X|Z=z generated by the server's model can be communicated back with approximately D KL (P X|Z=z ∥Q X|Z=z ) bits. Then, the data point (z, x), can be used for updating the common reference model Q X|Z=z . If models P and Q are of the same size, this scheme could be understood as model compression, where P is communicated by using samples from it, and then reconstructed at the client as Q. However, if model P is much bigger than Q, then presumably the clients' model could not fit the same distribution perfectly, and the distributions may differ significantly. As such, there is a trade-off between computation/ complexity at the client and the required communication. Additionally, the model at the server might be capable at the whole range of z ∈ Z, while the client might only be interested in some smaller subset of values of z ∈ Z ′ ⊆ Z. In this case, even less powerful model Q could learn good approximation of P X|Z=z for the desired input space.To demonstrate the validity of the idea, we perform numerical experiments in which we compare the expected number of bits required to communicate a text sample from a large language model (LLM). Instead of the estimation/learning step, we use models of different sizes from the family of Open Pre-train Transformer (OPT) models Zhang et al. (2022) (MIT license). In particular, we use the 13 billion (13B) parameter model to represent the desired conditional probability P , while Q is represented by one of the smaller OPT models with {125M, 350M, 1.3B, 2.7B, 6.7B} parameters, where M denotes million, and B billion. The classical way to communicate such samples is to generate the distribution at the server, sample from it, and then encode and transmit the result -we refer to this method as source coding, and, as mentioned before, it has the minimal cost H(P ), while the sample communication approach would have an approximate cost of D KL (P ∥Q) (per sample).In LLMs, text is represented as a sequence of tokens, where each token is a short string of charactersfor instance a phrase "nocturnal animals" would be parsed into tokens ("no", "ct", "urnal", " animals"). A LLM describes the probability of the next token conditioned on the all the previous ones. This is enough to describe any distribution over a sequence, as it can be factorized in an autoregressive wayP (X n ) = n i=1 P (X i |X i-1 ).To sample from such distributions, we can sample the first element, then conditioned on it sample the second one, and so on so forth. The entropy of a sequence of n random variables can be decomposed as H(P n ) = n i=1 H(P Xi|X i-1 ). Conditioned on a specific x t we can get an exact entropy H(P Xt|X t =x t ) as the model outputs probability of each token given x t . However, the OPT models use 50265 possible tokens, and thus to calculate the entropy of a sequence of length n, one would have to calculate the entropy of 50265 n different sequences which is infeasible. Instead, we estimate the entropy by sampling M = 50000 different random sequences of length N = 128 according to the model. To increase variety and simulate user prompts we condition the answers on questions from OpenbookQA dataset Mihaylov et al. (2018) . Then the entropy of a single token is estimated asEQUATION)We can similarly estimate the KL-divergence asEQUATION)The experiments were performed on 4 Nvidia RTX A6000 GPUs with 48GB of memory and PyTorch version 2.1, totalling 10 hours of wall-clock-time.The estimated expected communication cost per token is shown in Table 3 . Even the smallest 125M model allows us to decrease the communication requirement by more than 4 times. As the auxiliary model gets bigger, the communication cost decreases by up to 16-times. These results are asymptotic, and do not include the overhead for shorter sequences. Thus, we also present upper bounds for communication of groups of B ∈ {1, 2, 4, 8, 16, 32, 64, 128} tokens jointly in Figure 6 . For sample communication, the upper bounds are calculated by including the logarithmic and constant terms from Equation (4). For source coding, we use the upper bound from Equation ( 5). The sample communication approach beats source coding even for a group size B ≥ 2 for any of the auxiliary models. Only for B = 1 source coding is more efficient, as the constant overhead from equation (4) is 5, which is greater than the source coding cost. Even for small group sizes of B = 8, the sample communication approach achieves 2 -4 times reduction in the communication cost.1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The abstract and introduction clearly articulate the main claims and contributions of the paper, including the development of a universal sample coding problem (section 5), the derivation of bounds on communication costs (section 5), and the application of these methods in federated learning (section 7) and remote generation of samples from generative models (section 8 and appendix D).Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper makes a significant advancement on the classical channel simulation problem in which the reference distribution $Q$ is only known via samples. The authors derive information-theoretic bounds on the expected communication cost and also derive a practical algorithm that achieves the lower bound up to a multiplicative constant. Applications to federated learning are presented. This paper makes significant contributions and is enjoyable to read. I believe it will be of interest to the NeurIPS audience.

Reviews: review: paper_summary: paper_summary: The paper proposes a new problem called "Universal Sample Coding". This is related to a problem called reverse channel coding (or channel simulation) where the receiver must generate samples from a target distribution (P) that is known to the sender but not the receiver. In addition, the receiver and sender have shared randomness which is used to generate samples from a reference distribution Q. The goal is to characterize the communication complexity i.e., the number of information bits that must be transmitted from the sender to the receiver to generate the target sample(s). 

As far as I understand the authors propose a new problem setting where  (at least) the decoder also does not know the reference distribution Q, but estimates it. The paper provides upper and lower bounds on the communication complexity in this variant that have the same scaling in terms of the number of target samples.
strength_weakness: strength_weakness: The setup and results need clarification.

1. Can you clarify your problem setup? Is it that both the encoder and decoder do not know the samples?  Who generates the samples $X_1, \ldots, X_n$ generated from the reference distribution and how are they known to the encoder/decoder?   In the reverse channel coding/channel simulation settings that I am familiar with, both the sender and receiver know the reference distribution $Q$ and use common randomness to generate identical samples.  I wasn't sure how the samples are generated.

2. The statement of Thm 5.1 is not clear to me. Usually the rate bound should be on $H(f(P,Z)|Z)$  as $Z$ is known to both the encoder and decoder. Why do you consider $H(f(P,Z))$ as your rate?

3.Theorem 5.1 is for "some c", but in the numerical bounds you are optimizing over $c$. Should the statement be for every $c$?
questions: questions: please see the weakness section above..
limitations: limitations: yes
paper_summary: paper_summary: This paper studies the problem of communicating multiple samples from an unknown distribution using as few bits as possible. The authors provide upper and lower bounds on its communication cost that are within a multiplicative factor from each other. The upper bound is derived from analysing the communication cost of an universal sample coding algorithm, which the authors design based on the reverse channel coding results in information theory. The lower bound is derived based on analysing the connection between universal source coding and universal sample coding. Experiments show that the universal sample coding algorithm can reduce the communication cost in a Federated Learning scenario (up to $37\%$) or for generating samples from a Large Language Model (LLM) on a server (up to $16$ times).
strength_weakness: strength_weakness: + The gap between lower and upper bounds (the ratio between $\inf_c V_k(c)$ and $(k-1)/2$) is quite large for not very large $k$ (See Fig. 1 and Fig. 2). 
+  Algorithm 1 requires a shared source of randomness between the encoder and decoder.
questions: questions: + Please address the weakness comments above. Please explain how to generate the share source of randomness in practice. 
+ What is the distribution of the random string $Z \in \mathcal{Z}=\\{0,1\\}^{\infty}$ in Theorem 5.1?
+ What are the definitions of $I(P,X^n)$ and $I(P,\hat{P})$ in (16)?
limitations: limitations: This is a theoretical research paper, hence the negative society impact of this work is not direct. The authors mention some technical limitations of this work in Section 8.
paper_summary: paper_summary: This paper introduces Universal Sample Coding. This is a simple but significant extension to channel simulation where the sender and receiver communicate $N$ ($N>>1$) samples. The authors prove that the expected codelength per sample will be negligible when $N\rightarrow \infty$, and verify with toy examples. The authors also demonstrate the application of this coding scheme in federated learning and the communication of generated samples from generative models.
strength_weakness: strength_weakness: My major concern with this method is the runtime.
To ensure the per-sample cost is asymptotically 0, the authors propose a solution in which exponentially more samples are sent in each communication round. This means that when the problem scale is large, the $D_{KL}[P||Q]$ in some rounds can be large. 
For order random coding the author used in this paper, to achieve a low biased sample, the total sample size will be $O(2^{D_{KL}[P||Q]})$. 
$D_{KL}[P||Q]$ will, in the end, go to 0, but it will happen only after enough samples are sent. 
Therefore, I kindly ask the author to provide the sample size and the KL divergence for each problem and each round, which can provide more hints on how well this method can scale.



This is overall a good paper and I will be happy to further raise my score if the questions and weakness are addressed.
questions: questions: Besides Weaknesses, I have one further question:
Lemma 2 says an estimator exists for any distribution but does not provide the form for this estimator. 
Do you use this estimator for all of the experiments? In Fig 5, I can observe that the empirical results in this toy example obey the bounds very well. But do you think this will always be the case for the FL and LLM settings? Is there any potential caveat?
limitations: limitations: adequately addressed
paper_summary: paper_summary: The authors consider the problem of efficiently encoding a sequence of $n$ iid realisations $X_i \sim P$ using as few bits as possible. This scenario is different from one-shot channel simulation because it involves multiple rounds of communication. Between these rounds, the encoder and decoder adjust their coding distribution, leading to substantial improvements in compression performance. Using techniques and results from universal source coding, the authors characterise the optimal lower bound for this task and propose a scheme that achieves the lower bound within a multiplicative constant. They conduct experiments on a toy federated learning problem, showing that their technique could make federated learning significantly more robust when only a fraction of the clients participate in each communication round while reducing communication costs.
strength_weakness: strength_weakness: There are three major points that, if clarified, would make the paper significantly stronger:

1. **What is the sample complexity/run time of the authors' proposed solution?** 
Since Algorithm 1 calls an arbitrary channel simulation algorithm a subroutine, it should be relatively easy to calculate the total runtime in terms of the subroutine's runtime and state it as an accompanying result to Thm 5.1. Such a result would benefit people wishing to build on the authors' work, as the general issue with channel simulation algorithms with target $P$ and proposal $Q$ is that their runtime scales with $\Vert dP/dQ \Vert_\infty \geq \exp(KL[P \Vert Q])$.
2. **Clarifying the relationship between universal source coding, universal sample coding and channel simulation.** The authors' framework has key differences compared to channel simulation as well as universal source coding, which should be illustrated better. First, they switch to a streaming/multi-round communication setting, where the encoder and decoder can update their coding distribution, which differs from one-shot channel simulation as, in the latter case, there is only a single round of communication. Second, they fix the target distribution of the elements of the sequence, which is equivalent to assuming a delta distribution on the source in channel simulation. Similarly, the analogous setting in universal source coding would be to have a delta distribution as the source: for an alphabet of $K$ symbols, we wish to encode $n$ copies of the same fixed symbol. This can be done by encoding which symbol is repeated using $\log K$ bits and then encoding the number of repeats, which can be done using $\log n + O(\log\log n)$ bits with Elias delta coding. Therefore, the authors' setting and solution are closer to a universal code (a prefix code on the integers with a certain universality property) than a universal source code (a code which doesn't need to be prefix, but the ratio of the actual and optimal expected codelengths must go to one as the number of symbols increases). I realise that this terminology is annoyingly similar. Still, pointing this out/fixing this would enhance the clarity of the paper, especially given that in the current version, the authors seem to use the terms universal code and universal source code interchangeably. Furthermore, in the same way, that a universal code can be used to create a universal source code, I believe we could use the authors' proposed method to create a closer analogue of universal source coding: Assume we have two correlated variables $X, Y \sim P_{X, Y}$ with unknown marginal $P_Y$, if Alice receives $X \sim P_X$ and wants to communicate $Y \sim P_{Y \mid X}$. We could apply the authors' method of updating the counting distribution such that the approximating distribution converges to the marginal $P_Y$ so that the scheme's expected codelength will asymptotically be $n \cdot I[X; Y]$.
3. **The authors only state results assuming a uniform initial distribution.** There are two reasons why I put this as a major issue: 1) this assumption is not explicitly stated in the paper, but the authors make use of it both in the explanations as well as the proofs. 2) At the end of section 7, the authors suggest using a different initial distribution, whose performance is thus not covered by the theoretical statements, even though this setting is of the greatest practical relevance. I would conjecture that the initial choice should only affect the results up to a constant additive factor.

I have a few more minor issues as well:

1. It would be good to report the scheme's overhead compared to the theoretical lower bound in the FL scenario, akin to Figure 5.
2. I would need some clarification on the FL experimental setup and results: how did the authors set mu in the FL scenario? Why did the authors specifically pick seven samples per client? How does the performance change as a function of the number of samples sent per client and the fraction of participating clients?
3. "We consider ordered random coding Theis and Ahmed (2022) as the underlying sample communication method" - The authors should clarify (possibly by giving pseudocode in the appendix) what exact implementation of ORC they used. I imagine that the authors did the "right thing" and used the discrete version by drawing samples without replacement, and they either ran ORC until the sample space was exhausted or used a more sophisticated early termination criterion. In any case, it would be crucial to state this, as ORC is usually regarded as an approximate sampling scheme (and is usually used for continuous distributions).

Finally, I found some typos/issues with the notation and writing:

 - line 99: "non-asymptomatic"
 - "the sample communication bound does not rely on using the exact marginal distribution P, and still holds for other reference distributions Q" - I'm not sure what this sentence means.
 - Nitpick: the authors state in section 2 that all logs are in base 2, but then they use ln (eg in eq 11)
 - Theis and Ahmed (2022)—The reference is incorrect; the second author is Noureldin Yosri, so it should be Theis and Yosri (2022).
 - "In a single learning round, the server broadcasts the global mask θ to each client, which trains an updated version θ′ using their local data and then communicates a sample Bern(θ′) from its updated mask distribution to the central server, where the new global mask probability θ is estimated from all the samples" - It is quite difficult to parse this sentence; please split it in two.
 - I believe that $\hat{P}$ (in Section 6 and Figures 3 and 4) and $\hat{Q}$ (everywhere else) mean the same thing. Please update them to be consistent if so.
questions: questions: - Could we perhaps extend the result to continuous spaces by using KDE, or is there something fundamentally different about the continuous case?
 - "However, the real cost of sending a sample, as shown in (4), includes a constant, which would dominate the total cost as n increases, making it linear instead of logarithmic." - Does it? In the authors' case, since the approximating distribution converges to the true target, the ORC index distribution should converge to (approximately) the delta on 1. Hence, I would imagine that the overhead the authors mention would vanish if we used a universal source code to compress the ORC indices.
 - What is the long-term behaviour of the approximating distribution? My initial estimate would be that we have $O(k^{-1}n^{-1/2}) $ convergence to the target by the CLT; could we say something more precise? Is there a good way to speed up the convergence, maybe?
limitations: limitations: The authors should address the three main points I laid out in the Weaknesses section.

Comments: title: Reply to the authors' rebuttal
Follow up question on problem formulation
comment: Thank you very much for your answer to my questions. However, based on your answers, I keep my score unchanged.
Thank you for helping me better understand your contributions. I will keep my score as it is.
The reviewer is right. In fact, this is what we state on page 4 of the paper to derive equation (10). One can directly use any standard channel simulation technique to directly generate $n$ samples. (In the conditional version of the problem suggested by the reviewer, this would be equivalent to the mutual information expression). Our work uses this fact for min-max analysis i.e. guarantees for any distribution $P$.
I thank the authors for their response. They addressed my concerns, and I increased my score accordingly.
Thank you for clarifying the problem formulation. I was wondering why is the problem setup not already a special case of the one-shot setup in channel simulation? 

Here is one specific example of one-shot setting: The encoder observes $X \sim p_X(\cdot)$ and the decoder wants to sample $Y \sim p_{Y|X}(\cdot)$. We know from prior works that the communication rate required is approximately $I(X;Y)$. In your setting you want to sample $Y_1,\ldots, Y_n$ i.i.d from $p_{Y|X}(\cdot)$ then using their scheme and defining $Z=(Y_1, \ldots, Y_n)$ the rate of  $I(X;Y_1, \ldots, Y_n)$ is naturally achievable. 

Even if you don't assume $X$ to be random, a natural connection to one-shot schemes can be made in a similar fashion. Can you provide a comparison between your proposed approach a natural extension of one-shot schemes?
> Empirically, in the data provided in the global response (excluding the last round with fewer samples), the probability that: given KL>25 for i-th round, KL>25 for (i+1)-th round is 37%. Thus in most of the cases, the estimate goes 'on the right track' after diverging.

Thank you for your reply. I think showing these results in some way (for example, adding an extra plot or tab) can make your paper stronger. 

I feel my concerns have been addressed. Thus, I raise my rating from 6 to 7.
For the experiments, we draw $2^{\min(D_{KL}, 25)}$ samples per round. The high KL is caused by the large number of samples sent which is not the case in FL. We agree that this might be an issue for the generative case, however, in those settings greater computational resources are usually available.

We think that using biased samples can delay or prevent convergence the convergence. Empirically, in the data provided in the global response (excluding the last round with fewer samples), the probability that: given KL>25 for i-th round, KL>25 for (i+1)-th round is 37%. Thus in most of the cases,  the estimate goes 'on the right track' after diverging.
Thank you for your reply and your additional results on the KL. The plot clear and I recommend including this plot in your paper later. 

I agree that KL values up to 20 are reasonable. However, I do not necessarily agree that KL up to 30 (ideally, at least 2^30 samples if your KL is in bits) is still computationally feasible, especially considering the setting of FL where each client should not have very strong computational resources.  Could you explain how many samples you use for this KL value? 

>  For the few outliers, we can either accept the high-bias samples or incorporate a simple feedback mechanism 

What did you do in your experiments? If you just accept the high bias, how does this bias influence, e.g., convergence?