Gradient Transformer: Learning to Generate Updates for LLMs

Binh-Nguyen Nguyen, Khang Tran, Hai Phan, Issa Khalil

TL;DR: We propose a data-free knowledge distillation framework to generate updates for large language models based on tiny language models fine-tuned on private data.

Abstract: Many organizations lack computational resources to fine-tune large language models (LLMs) on private (unshareable) data for better utility, while fine-tuning tiny language models (TinyLMs) alone performs poorly. To address this bottleneck, we propose a data-free knowledge distillation framework that generates LLM update vectors based on TinyLMs fine-tuned on private data. An update vector is a vector of parameter changes from an initial model to its fine-tuned version on a dataset, capturing the effect of cumulative gradient steps during fine-tuning. The key idea of our framework is a novel **Gradient Transformer** that transforms TinyLM's update vectors into LLM's update vectors. As derived from shadow datasets, $\texttt{Grad-Transformer}$ captures the correlation between TinyLM and LLM update vectors, enabling third-party providers to generate LLM update vectors given the organization's TinyLM update vectors without accessing the organization's private data. The framework supports multi-organization collaboration to jointly update LLMs, improving performance and cost-efficiency. Extensive experiments across language modeling and reasoning tasks show that $\texttt{Grad-Transformer}$ remarkably outperforms state-of-the-art knowledge distillation baselines, even under strict differential privacy protection.

Lay Summary:

Many organizations handle sensitive data, such as patient records or private business information, that they cannot share with outside parties. This makes it difficult for them to benefit from the most powerful AI language models, which require expensive computing resources and access to the data itself to be fine-tuned for specific tasks.

We developed a framework called Grad-Transformer that lets organizations improve large AI models without ever sharing their data. The key idea is that an organization first trains a much smaller, cheaper AI model on their private data locally. They then send only a compact summary of what the model learned, a "difference vector" capturing how training changed the model's parameters, to a third-party service provider. Our Grad-Transformer, a trained translator model, converts this small-model summary into an equivalent update for a much larger, more capable AI model.

Crucially, the private data never leaves the organization. Multiple organizations can even collaborate this way, pooling their updates to jointly improve a shared large model. Our experiments show that this approach substantially closes the performance gap between small and large models, even when strict mathematical privacy guarantees are applied to the local training process.

Link To Code: https://github.com/nguyenrtm/grad-transformer

Primary Area: General Machine Learning->Transfer, Multitask and Meta-learning

Keywords: Knowledge Distillation, Data Privacy, Transformers

Submission Number: 28042