Technical Summary: Data Compression with Stochastic Codes

Problem Statement

The paper addresses a fundamental limitation in traditional lossy source coding: the reliance on quantisation. Standard quantisers map continuous source data to a discrete set of values, forcing the reconstruction distribution to be a Dirac measure (a single point). This constraint limits the flexibility of the compression mechanism, often leading to suboptimal trade-offs between rate, distortion, and "realism" (the statistical fidelity of the output distribution to the source).

The authors propose Relative Entropy Coding (REC) as a general alternative. REC aims to construct stochastic codes that can simulate any desired conditional distribution $P_{\hat{X}|X}$ for reconstruction, rather than being restricted to deterministic quantisation. The core challenge is to achieve this simulation with a codelength that is close to the mutual information $I[X; \hat{X}]$, while maintaining computational feasibility.

Methodology

The paper defines a stochastic code as a triplet $(Z, \text{enc}_Z, \text{dec}_Z)$ where:

1. $Z$ is shared randomness (independent of the input $X$).
2. $\text{enc}_Z$ is an encoder mapping $X$ to a bit string.
3. $\text{dec}_Z$ is a decoder mapping the bit string back to a sample $\hat{X}$.

The goal is to satisfy $\text{dec}_Z(\text{enc}_Z(X)) \sim P_{\hat{X}|X}$ while minimizing the expected codelength. The theoretical lower bound for this codelength is the mutual information $I[X; \hat{X}]$.

The paper details several construction methodologies:

1. Selection Samplers

The authors frame REC as a sampling problem where the decoder must select a sample from a shared sequence of i.i.d. proposals $\{Y_i\} \sim P_Y$ such that the selected sample follows $P_{\hat{X}|X}$.

• Rejection Sampling: The simplest approach. The encoder finds the first index $K$ where a proposal $Y_K$ is accepted based on a density ratio threshold. The index $K$ is encoded. However, standard rejection sampling often yields a codelength of $O(\|r\|_\infty)$, which can be exponentially larger than $I[X; \hat{X}]$ if the density ratio is unbounded or large.
• Poisson Functional Representation (A Coding):* To overcome the inefficiency of standard rejection sampling, the authors describe a method using a Poisson process to simulate proposals in an order that minimizes the index. This "reordered" approach (A* coding) guarantees that the expected codelength is within a logarithmic factor of the mutual information, making it a valid relative entropy code.
• Runtime Limitation: The paper notes that for general selection samplers, the expected runtime scales at least as $E[\|r_X\|_\infty]$, which is often exponential in the information content. This makes general REC algorithms impractically slow compared to entropy coding.

2. Dithered and Layered Quantisers

To achieve computational efficiency, the paper focuses on specific channel structures:

• Dithered Quantisation: For additive uniform noise channels ($Y = X + U$), the authors utilize the "dithering identity" to construct a code where the encoder sends an entropy-coded index of a quantised value. This achieves a codelength of $I[X; Y] + O(1)$ with linear runtime.
• Layered Quantisation: This extends dithered quantisation to unimodal additive noise channels by representing the noise as a scale mixture of shifted uniforms (SMSU). By randomly scaling and shifting the quantisation grid, the method achieves optimal codelengths and fast runtime for one-dimensional unimodal distributions.

3. Approximate and Advanced Methods

• Approximate Codes: To handle unbounded density ratios (e.g., in diffusion models) or reduce runtime, the paper discusses step-limited samplers. These algorithms terminate early if a budget is exceeded, outputting an approximate sample. This trades exact distribution matching for polynomial runtime and parallelizability.
• Branch-and-Bound (BnB) Samplers: For unimodal distributions, BnB search strategies can reduce the expected runtime from exponential to logarithmic in the density ratio, offering an alternative to layered quantisers.
• Channel Codes: Recent work using polar codes is mentioned as a method to construct REC for high-dimensional i.i.d. Bernoulli channels with optimal runtime.

Key Contributions and Applications

The paper positions REC not just as a theoretical curiosity but as a practical tool for learned compression and privacy-preserving data release.

1. Learned Lossy Compression:

   • REC enables end-to-end training of compression systems. Unlike quantisation, which is non-differentiable, REC allows the use of the reparameterization trick to optimize the conditional distribution $P_{\hat{X}|X}$ directly via gradient descent.
   • Case Study (INRs): The authors highlight Bayesian Implicit Neural Representations, where neural network weights are treated as random variables. Instead of quantising weights, the system learns a posterior distribution and uses REC to encode a sample from it.
   • Case Study (Federated Learning): In Federated Probabilistic Mask Training (FedPM), REC reduces communication costs by encoding the index of a selected subnetwork mask rather than the mask itself, achieving up to 82-fold bitrate reductions.

2. Compression with Realism Constraints:

   • REC allows the reconstruction $\hat{X}$ to follow a distribution that matches the source statistics (realism) rather than minimizing pixel-wise distortion. The paper argues that quantisers are inherently disadvantaged here, whereas REC can achieve lower rates for a given realism constraint by utilizing noise mechanisms (e.g., diffusion models).

3. Diffusion Models (DiffC):

   • The paper presents DiffC, a variable-rate compression scheme using Gaussian diffusion models. By treating the diffusion process as a sequence of REC problems along a trajectory, the method achieves "perfectly realistic" reconstructions.
   • Practicality: Since exact REC is infeasible for diffusion models due to unbounded density ratios, the authors advocate for step-limited approximate samplers, which enable parallelization on GPUs, reducing encoding time from hours to seconds.

4. Privacy and Reinforcement Learning:

   • Differential Privacy: REC provides a framework to encode privatized data (e.g., Laplace or Gaussian noise mechanisms) efficiently while maintaining privacy guarantees.
   • Remote Reinforcement Learning: REC allows a controller to transmit a short index identifying a candidate action from the agent's policy, rather than the action itself, significantly reducing communication overhead.

Results and Claims

• Theoretical Bounds: The paper establishes that while general REC algorithms have a codelength lower bound of $I[X; \hat{X}] + \log(I[X; \hat{X}] + 1)$, specific constructions (like dithered quantisers) can achieve $I[X; \hat{X}] + O(1)$.
• Computational Trade-offs: The authors are explicit that general REC is computationally expensive. The runtime of selection samplers is generally exponential in the information content. However, for specific structures (unimodal noise, diffusion models with approximation), efficient algorithms exist.
• Empirical Performance: The paper cites experimental results from referenced works showing:
  • Up to 82-fold bitrate reduction in federated learning (FedPM).
  • 12-fold and 41-fold reductions in remote reinforcement learning compared to transmitting actions or rewards directly.
  • Successful implementation of DiffC on GPUs, reducing encoding time from hours to seconds.

Significance and Limitations

Significance:
The paper argues that REC shifts the paradigm of lossy compression from "choosing a quantiser and analyzing its error" to "choosing a desired reconstruction distribution and implementing it." This flexibility is crucial for modern applications requiring realism (e.g., generative AI) and privacy, where the output distribution matters more than point-wise accuracy. It bridges information theory with machine learning, enabling differentiable compression pipelines.

Limitations:
The authors acknowledge two primary barriers to widespread adoption:

1. Speed: General REC algorithms are significantly slower than quantisation-based methods. While specialized algorithms exist for specific distributions, a universal fast REC algorithm remains an open problem.
2. Synchronization: REC relies on shared randomness (e.g., PRNG seeds). Maintaining perfect synchronization between encoder and decoder is an engineering challenge; desynchronization leads to catastrophic failure.

Conclusion:
The paper concludes that REC is an emerging field with profound theoretical elegance and practical potential, particularly in learned compression and privacy. However, its practical utility is currently constrained by computational complexity and engineering overhead, making the development of fast, special-purpose REC algorithms a critical direction for future research.