Model Update Compression Strategies

While some approaches focus on compressing gradients computed during local training, another effective strategy targets the result of that local training: the model update itself. Instead of sending compressed gradients after each local step, clients perform their local training (often for multiple epochs) and then compress the resulting change in model weights, $\Delta w = w_{local} - w_{global}$, or sometimes the entire updated local model $w_{local}$ before transmission. This approach is particularly relevant when clients perform significant local computation between communication rounds.

Compressing the model update $\Delta w$ directly addresses the uplink communication bottleneck. The techniques employed often mirror those used for gradient compression but are applied to the cumulative update vector or tensor. Let's examine some prominent strategies.

Quantization of Model Updates

Similar to gradient quantization, model updates can be quantized to reduce the number of bits needed to represent each parameter change. We replace the high-precision floating-point values in $\Delta w$ with lower-precision representations.

Common techniques include:

• Uniform Quantization: Values are mapped to a fixed set of uniformly spaced levels within a specific range. For instance, scaling the update vector to $[-1, 1]$ and then mapping each value to $b$ bits.
• Stochastic Quantization: Introduces randomness, often quantizing a value $x$ to $\lfloor x \rfloor$ with probability $x - \lfloor x \rfloor$ and to $\lceil x \rceil$ with probability $\lceil x \rceil - x$. This can help preserve expectation.
• Non-uniform Quantization (e.g., Logarithmic): Uses more quantization levels for smaller values and fewer for larger values, potentially better capturing the distribution of parameter updates, which often have many small values and fewer large ones.

The server receives the quantized updates ($\Delta w_{quantized}$) and needs to dequantize them before aggregation. While simple, quantization introduces errors ($\Delta w - \Delta w_{quantized}$) that can affect convergence. Techniques like error compensation, discussed later, can help mitigate this.

Sparsification of Model Updates

Sparsification aims to send only a fraction of the parameters in the model update vector $\Delta w$, setting the rest to zero. This directly reduces the amount of data transmitted, especially if using sparse data formats.

• Top-k Sparsification: The client selects the $k$ parameters in $\Delta w$ with the largest absolute magnitudes and transmits only these values along with their indices. The remaining $d-k$ parameters are treated as zero. The value of $k$ controls the compression ratio; a smaller $k$ means higher compression but potentially more information loss.
• Threshold Sparsification: The client sends only those parameters whose absolute magnitude $|\Delta w_i|$ exceeds a predefined threshold $\tau$. The number of parameters sent can vary between clients and rounds.
• Random-k Sparsification: A random subset of $k$ parameters is selected and transmitted. While simpler, it might not capture the most significant changes as effectively as Top-k.

Here's a illustration of Top-k sparsification for a model update vector:

# Pseudocode for Top-k Sparsification of a model update

def top_k_sparsify(delta_w, k):
  """
  Sparsifies the model update delta_w by keeping only the top k elements by magnitude.

  Args:
    delta_w: The model update vector (e.g., a NumPy array).
    k: The number of elements to keep.

  Returns:
    A sparse representation (e.g., indices and values) of the top k elements.
  """
  if k >= len(delta_w):
    # No sparsification needed if k is larger than or equal to vector size
    indices = np.arange(len(delta_w))
    values = delta_w
    return indices, values

  # Calculate magnitudes and find the threshold for the k-th largest
  magnitudes = np.abs(delta_w)
  threshold = np.sort(magnitudes)[-k] # Find the k-th largest magnitude

  # Select indices where magnitude is greater than or equal to the threshold
  # Handle potential ties carefully to ensure exactly k (or slightly more if ties at threshold)
  mask = magnitudes >= threshold
  indices = np.where(mask)[0]
  values = delta_w[mask]

  # If ties resulted in more than k values, take the top k among them
  if len(indices) > k:
      top_indices_within_mask = np.argsort(np.abs(values))[::-1][:k]
      indices = indices[top_indices_within_mask]
      values = values[top_indices_within_mask]

  return indices, values

# --- Client Side ---
# Assume delta_w is the calculated model update vector
# Assume k is the sparsity parameter

k = int(0.1 * len(delta_w)) # Example: Keep top 10%
sparse_indices, sparse_values = top_k_sparsify(delta_w, k)

# Transmit sparse_indices and sparse_values to the server

# --- Server Side ---
# Receive sparse_indices and sparse_values from a client

# Reconstruct the (sparse) update vector for aggregation
reconstructed_delta_w = np.zeros_like(global_model_weights) # Or appropriate shape
reconstructed_delta_w[sparse_indices] = sparse_values

# Aggregate reconstructed_delta_w with updates from other clients

Sparsification requires careful handling during aggregation. The server needs to reconstruct the sparse vectors into a common dimension (often initializing with zeros) before averaging or summing them. Sending indices adds overhead, but this is usually much smaller than sending the dense, non-zero values.

Structured Updates and Low-Rank Methods

For large models, especially neural networks with large weight matrices, the update $\Delta W$ (now a matrix or tensor) can be approximated using structured representations. A common approach is low-rank approximation.

Instead of transmitting the full $m \times n$ update matrix $\Delta W$, the client computes and sends two smaller matrices, $U$ ($m \times r$) and $V$ ($n \times r$), such that $\Delta W \approx UV^T$. Here, $r$ is the rank of the approximation, chosen such that $r \ll \min(m, n)$. The number of parameters transmitted is reduced from $m \times n$ to $r \times (m + n)$. This can lead to significant savings when $r$ is small.

The server receives $U$ and $V$ from each client, reconstructs the approximate $\Delta W$ (or performs aggregation directly using the low-rank factors, which can be more efficient), and applies the aggregated update. Choosing the rank $r$ involves a trade-off between compression and the fidelity of the update approximation.

Combining Techniques

These strategies are not mutually exclusive. It's often effective to combine them. For example:

1. Calculate the full model update $\Delta w$.
2. Sparsify $\Delta w$ using Top-k to keep the most significant changes.
3. Quantize the remaining non-zero values to further reduce their size.

This layered approach can achieve very high compression ratios.

Handling Compression Error

As mentioned, quantization and sparsification (especially non-Top-k methods) introduce errors. If ignored, these errors can accumulate over rounds and slow down or even prevent convergence. Error compensation or error feedback mechanisms, similar to those used in gradient compression, can be applied.

The basic idea is for the client to remember the compression error made in the current round and add it back to the model update in the next round before applying compression again.

Let $\Delta w_t$ be the true update at round $t$. Let $c(\cdot)$ be the compression operation (e.g., quantization + sparsification). Let $e_t$ be the accumulated error.

1. Update Calculation: Client calculates $\Delta w_t$.
2. Error Feedback: Client computes the value to be compressed: $u_t = \Delta w_t + e_{t-1}$ (where $e_0 = 0$).
3. Compression: Client compresses $u_t$: $\delta_t = c(u_t)$.
4. Error Update: Client calculates the error for this round: $e_t = u_t - \delta_t$. This error $e_t$ is stored locally for the next round $t+1$.
5. Transmission: Client sends the compressed update $\delta_t$ to the server.

This ensures that errors made in one round are corrected for in subsequent rounds, preventing systematic drift and often restoring convergence properties close to the uncompressed case, albeit sometimes at a slightly slower rate.

Trade-offs Summary

Choosing a model update compression strategy requires balancing several factors:

• Communication Savings: How much bandwidth is saved? Sparsification and low-rank methods offer tunable compression, while quantization offers fixed reduction based on bit width.
• Computation Cost: Compression and decompression add computational overhead on clients and potentially the server. Complex quantization or SVD for low-rank approximation can be demanding for resource-constrained clients.
• Impact on Convergence: Aggressive compression can slow down convergence or lead to lower final model accuracy. Error compensation helps but doesn't completely eliminate the impact.
• Implementation Complexity: Implementing error feedback or low-rank factorization adds complexity compared to simple quantization or sparsification.

Relationship between communication cost reduction and potential impact on final model accuracy for different compression strategies. Exact trade-offs depend heavily on the specific task, model, data heterogeneity, and chosen hyperparameters.

Model update compression provides a powerful set of tools alongside gradient compression for making federated learning practical in communication-constrained environments. The optimal choice often depends on the specific system characteristics, the model architecture, and the tolerance for computational overhead and potential accuracy trade-offs.