TLDR: Cocoon is a new hardware-software framework designed to make differentially private (DP) machine learning training with correlated noises more efficient and scalable. While correlated noises improve model accuracy over traditional DP methods, they introduce significant memory and computational overheads for large models and embedding tables. Cocoon addresses these issues with Cocoon-Emb, which pre-computes and coalesces noises for embedding tables, and Cocoon-NMP, a near-memory processing device for CXL memory that accelerates noise generation for large models. The framework achieves significant speedups, ranging from 1.55x to 10.82x over baseline approaches.
In the rapidly evolving world of machine learning, models are becoming increasingly powerful, but this power comes with a significant challenge: data privacy. Machine learning models, especially large ones, have a tendency to ‘memorize’ their training data, which can lead to sensitive user information being leaked. This is a serious concern, particularly when models are trained on personal data like medical records, audio inputs, or chat logs.
To combat this, a technique called Differential Privacy (DP) has gained considerable attention. DP provides a mathematical guarantee that the outcome of a training algorithm won’t significantly change if a single data entry in the input dataset is altered, thus limiting what an adversary can infer about individual training samples. A well-known algorithm in this space is DP-SGD, which adds a random Gaussian noise to gradients during each training iteration to ensure privacy. However, this noise can unfortunately degrade the accuracy of the trained model.
More recent advancements have introduced a new family of approaches that use ‘correlated noises’. Instead of adding independent noise at each step, these methods add carefully designed noises that are correlated across iterations. The clever part is that these correlated noises are designed to partially cancel each other out over time, leading to improved model accuracy compared to DP-SGD. While promising, these correlated noise mechanisms have largely been overlooked by the system and architecture communities, meaning their practical implications and potential bottlenecks haven’t been thoroughly studied.
This is where a new research paper, titled “Cocoon: A System Architecture for Differentially Private Training with Correlated Noises”, steps in. Authored by Donghwan Kim, Xin Gu, Jinho Baek, Timothy Lo, Younghoon Min, Kwangsik Shin, Jongryool Kim, Jongse Park, and Kiwan Maeng, the paper presents an extensive study of these emerging DP training methods. The researchers found that while correlated noises offer accuracy benefits, they introduce non-negligible overheads, especially for models with large embedding tables (common in recommendation systems) and for massive, billion-parameter models like Large Language Models (LLMs).
The core problem lies in managing the ‘noise history’ – the record of past noises needed to generate new correlated noises. This history can become enormous, often exceeding the memory capacity of GPUs or even main system memory. When memory runs out, data must be offloaded to slower storage, causing significant performance slowdowns. Additionally, the process of mixing these past noises (a computation called GEMV) also adds a computational burden.
To address these challenges, the researchers propose **Cocoon**, a novel hardware-software co-designed framework. Cocoon aims to make differentially private training with correlated noises efficient and scalable. It tackles the memory and compute overheads through two main innovations:
Cocoon-Emb: Optimizing for Embedding Tables
Embedding tables are crucial components in models like Deep Learning Recommendation Models (DLRMs), converting categorical features into dense vector representations. They can be very large, but only a small subset of their entries are used in each training iteration. Traditional correlated noise mechanisms add noise to all entries, even unused ones, leading to inefficiency. Cocoon-Emb addresses this by pre-computing all the correlated noises for embedding tables before training begins. It then stores these pre-computed noises in a compact, ‘coalesced’ format, exploiting the sparsity of gradient updates. This means it only stores aggregated noise right before an entry is actually accessed, significantly reducing memory footprint. Furthermore, Cocoon-Emb intelligently splits embedding entries into ‘hot’ (frequently accessed) and ‘cold’ (rarely accessed) categories, applying coalescing only to cold entries to maximize efficiency. This approach can improve performance by 2.33 to 10.82 times for models with large embedding tables.
Also Read:
- Securing Spiking Neural Network Inference with Homomorphic Encryption
- NANOMIND: Bringing Advanced Multimodal AI to Small, Battery-Powered Devices
Cocoon-NMP: Near-Memory Processing for Large Models
For large-scale models where the noise history is too vast to fit into main memory and must be offloaded to secondary memory like CXL (Compute Express Link) memory, data transfer becomes a major bottleneck. Cocoon-NMP introduces a custom near-memory processing (NMP) device integrated into the CXL memory controller. This device allows the GEMV computation (the mixing of past noises) to happen directly within the CXL memory controller, much closer to where the noise history is stored. By performing computation near the data, Cocoon-NMP drastically reduces the need for frequent and slow data transfers between CXL memory and the CPU or GPU. This hardware-software co-design can lead to performance improvements of 1.55 to 3.06 times, especially for large, billion-parameter models.
Cocoon represents a significant step forward in making differentially private training with correlated noises practical and efficient for real-world applications. By systematically identifying and addressing the system-level bottlenecks, this framework enables the use of more accurate privacy-preserving techniques without prohibitive performance costs. For more in-depth technical details, you can read the full research paper here.
