-
Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers
Authors:
Rya Sanovar,
Srikant Bharadwaj,
Renee St. Amant,
Victor Rühle,
Saravan Rajmohan
Abstract:
Transformer-based models have emerged as one of the most widely used architectures for natural language processing, natural language generation, and image generation. The size of the state-of-the-art models has increased steadily reaching billions of parameters. These huge models are memory hungry and incur significant inference latency even on cutting edge AI-accelerators, such as GPUs. Specifica…
▽ More
Transformer-based models have emerged as one of the most widely used architectures for natural language processing, natural language generation, and image generation. The size of the state-of-the-art models has increased steadily reaching billions of parameters. These huge models are memory hungry and incur significant inference latency even on cutting edge AI-accelerators, such as GPUs. Specifically, the time and memory complexity of the attention operation is quadratic in terms of the total context length, i.e., prompt and output tokens. Thus, several optimizations such as key-value tensor caching and FlashAttention computation have been proposed to deliver the low latency demands of applications relying on such large models. However, these techniques do not cater to the computationally distinct nature of different phases during inference.
To that end, we propose LeanAttention, a scalable technique of computing self-attention for the token-generation phase (decode-phase) of decoder-only transformer models. LeanAttention enables scaling the attention mechanism implementation for the challenging case of long context lengths by re-designing the execution flow for the decode-phase. We identify that the associative property of online softmax can be treated as a reduction operation thus allowing us to parallelize the attention computation over these large context lengths. We extend the "stream-K" style reduction of tiled calculation to self-attention to enable parallel computation resulting in an average of 2.6x attention execution speedup over FlashAttention-2 and up to 8.33x speedup for 512k context lengths.
△ Less
Submitted 16 May, 2024;
originally announced May 2024.
-
Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing
Authors:
Dujian Ding,
Ankur Mallick,
Chi Wang,
Robert Sim,
Subhabrata Mukherjee,
Victor Ruhle,
Laks V. S. Lakshmanan,
Ahmed Hassan Awadallah
Abstract:
Large language models (LLMs) excel in most NLP tasks but also require expensive cloud servers for deployment due to their size, while smaller models that can be deployed on lower cost (e.g., edge) devices, tend to lag behind in terms of response quality. Therefore in this work we propose a hybrid inference approach which combines their respective strengths to save cost and maintain quality. Our ap…
▽ More
Large language models (LLMs) excel in most NLP tasks but also require expensive cloud servers for deployment due to their size, while smaller models that can be deployed on lower cost (e.g., edge) devices, tend to lag behind in terms of response quality. Therefore in this work we propose a hybrid inference approach which combines their respective strengths to save cost and maintain quality. Our approach uses a router that assigns queries to the small or large model based on the predicted query difficulty and the desired quality level. The desired quality level can be tuned dynamically at test time to seamlessly trade quality for cost as per the scenario requirements. In experiments our approach allows us to make up to 40% fewer calls to the large model, with no drop in response quality.
△ Less
Submitted 22 April, 2024;
originally announced April 2024.
-
LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression
Authors:
Zhuoshi Pan,
Qianhui Wu,
Huiqiang Jiang,
Menglin Xia,
Xufang Luo,
Jue Zhang,
Qingwei Lin,
Victor Rühle,
Yuqing Yang,
Chin-Yew Lin,
H. Vicky Zhao,
Lili Qiu,
Dongmei Zhang
Abstract:
This paper focuses on task-agnostic prompt compression for better generalizability and efficiency. Considering the redundancy in natural language, existing approaches compress prompts by removing tokens or lexical units according to their information entropy obtained from a causal language model such as LLaMa-7B. The challenge is that information entropy may be a suboptimal compression metric: (i)…
▽ More
This paper focuses on task-agnostic prompt compression for better generalizability and efficiency. Considering the redundancy in natural language, existing approaches compress prompts by removing tokens or lexical units according to their information entropy obtained from a causal language model such as LLaMa-7B. The challenge is that information entropy may be a suboptimal compression metric: (i) it only leverages unidirectional context and may fail to capture all essential information needed for prompt compression; (ii) it is not aligned with the prompt compression objective.
To address these issues, we propose a data distillation procedure to derive knowledge from an LLM to compress prompts without losing crucial information, and meantime, introduce an extractive text compression dataset. We formulate prompt compression as a token classification problem to guarantee the faithfulness of the compressed prompt to the original one, and use a Transformer encoder as the base architecture to capture all essential information for prompt compression from the full bidirectional context. Our approach leads to lower latency by explicitly learning the compression objective with smaller models such as XLM-RoBERTa-large and mBERT.
We evaluate our method on both in-domain and out-of-domain datasets, including MeetingBank, LongBench, ZeroScrolls, GSM8K, and BBH. Despite its small size, our model shows significant performance gains over strong baselines and demonstrates robust generalization ability across different LLMs. Additionally, our model is 3x-6x faster than existing prompt compression methods, while accelerating the end-to-end latency by 1.6x-2.9x with compression ratios of 2x-5x.
△ Less
Submitted 19 March, 2024;
originally announced March 2024.
-
Risk-aware Adaptive Virtual CPU Oversubscription in Microsoft Cloud via Prototypical Human-in-the-loop Imitation Learning
Authors:
Lu Wang,
Mayukh Das,
Fangkai Yang,
Junjie Sheng,
Bo Qiao,
Hang Dong,
Si Qin,
Victor Rühle,
Chetan Bansal,
Eli Cortez,
Íñigo Goiri,
Saravan Rajmohan,
Qingwei Lin,
Dongmei Zhang
Abstract:
Oversubscription is a prevalent practice in cloud services where the system offers more virtual resources, such as virtual cores in virtual machines, to users or applications than its available physical capacity for reducing revenue loss due to unused/redundant capacity. While oversubscription can potentially lead to significant enhancement in efficient resource utilization, the caveat is that it…
▽ More
Oversubscription is a prevalent practice in cloud services where the system offers more virtual resources, such as virtual cores in virtual machines, to users or applications than its available physical capacity for reducing revenue loss due to unused/redundant capacity. While oversubscription can potentially lead to significant enhancement in efficient resource utilization, the caveat is that it comes with the risks of overloading and introducing jitter at the level of physical nodes if all the co-located virtual machines have high utilization. Thus suitable oversubscription policies which maximize utilization while mitigating risks are paramount for cost-effective seamless cloud experiences. Most cloud platforms presently rely on static heuristics-driven decisions about oversubscription activation and limits, which either leads to overloading or stranded resources. Designing an intelligent oversubscription policy that can adapt to resource utilization patterns and jointly optimizes benefits and risks is, largely, an unsolved problem. We address this challenge with our proposed novel HuMan-in-the-loop Protoypical Imitation Learning (ProtoHAIL) framework that exploits approximate symmetries in utilization patterns to learn suitable policies. Also, our human-in-the-loop (knowledge-infused) training allows for learning safer policies that are robust to noise and sparsity. Our empirical investigations on real data show orders of magnitude reduction in risk and significant increase in benefits (saving stranded cores) in Microsoft cloud platform for 1st party (internal services).
△ Less
Submitted 13 January, 2024;
originally announced January 2024.
-
Unlocking Spatial Comprehension in Text-to-Image Diffusion Models
Authors:
Mohammad Mahdi Derakhshani,
Menglin Xia,
Harkirat Behl,
Cees G. M. Snoek,
Victor Rühle
Abstract:
We propose CompFuser, an image generation pipeline that enhances spatial comprehension and attribute assignment in text-to-image generative models. Our pipeline enables the interpretation of instructions defining spatial relationships between objects in a scene, such as `An image of a gray cat on the left of an orange dog', and generate corresponding images. This is especially important in order t…
▽ More
We propose CompFuser, an image generation pipeline that enhances spatial comprehension and attribute assignment in text-to-image generative models. Our pipeline enables the interpretation of instructions defining spatial relationships between objects in a scene, such as `An image of a gray cat on the left of an orange dog', and generate corresponding images. This is especially important in order to provide more control to the user. CompFuser overcomes the limitation of existing text-to-image diffusion models by decoding the generation of multiple objects into iterative steps: first generating a single object and then editing the image by placing additional objects in their designated positions. To create training data for spatial comprehension and attribute assignment we introduce a synthetic data generation process, that leverages a frozen large language model and a frozen layout-based diffusion model for object placement. We compare our approach to strong baselines and show that our model outperforms state-of-the-art image generation models in spatial comprehension and attribute assignment, despite being 3x to 5x smaller in parameters.
△ Less
Submitted 28 November, 2023;
originally announced November 2023.
-
Rethinking Privacy in Machine Learning Pipelines from an Information Flow Control Perspective
Authors:
Lukas Wutschitz,
Boris Köpf,
Andrew Paverd,
Saravan Rajmohan,
Ahmed Salem,
Shruti Tople,
Santiago Zanella-Béguelin,
Menglin Xia,
Victor Rühle
Abstract:
Modern machine learning systems use models trained on ever-growing corpora. Typically, metadata such as ownership, access control, or licensing information is ignored during training. Instead, to mitigate privacy risks, we rely on generic techniques such as dataset sanitization and differentially private model training, with inherent privacy/utility trade-offs that hurt model performance. Moreover…
▽ More
Modern machine learning systems use models trained on ever-growing corpora. Typically, metadata such as ownership, access control, or licensing information is ignored during training. Instead, to mitigate privacy risks, we rely on generic techniques such as dataset sanitization and differentially private model training, with inherent privacy/utility trade-offs that hurt model performance. Moreover, these techniques have limitations in scenarios where sensitive information is shared across multiple participants and fine-grained access control is required. By ignoring metadata, we therefore miss an opportunity to better address security, privacy, and confidentiality challenges. In this paper, we take an information flow control perspective to describe machine learning systems, which allows us to leverage metadata such as access control policies and define clear-cut privacy and confidentiality guarantees with interpretable information flows. Under this perspective, we contrast two different approaches to achieve user-level non-interference: 1) fine-tuning per-user models, and 2) retrieval augmented models that access user-specific datasets at inference time. We compare these two approaches to a trivially non-interfering zero-shot baseline using a public model and to a baseline that fine-tunes this model on the whole corpus. We evaluate trained models on two datasets of scientific articles and demonstrate that retrieval augmented architectures deliver the best utility, scalability, and flexibility while satisfying strict non-interference guarantees.
△ Less
Submitted 27 November, 2023;
originally announced November 2023.
-
Hybrid Retrieval-Augmented Generation for Real-time Composition Assistance
Authors:
Menglin Xia,
Xuchao Zhang,
Camille Couturier,
Guoqing Zheng,
Saravan Rajmohan,
Victor Ruhle
Abstract:
Retrieval augmentation enhances performance of traditional language models by incorporating additional context. However, the computational demands for retrieval augmented large language models (LLMs) pose a challenge when applying them to real-time tasks, such as composition assistance. To address this limitation, we propose the Hybrid Retrieval-Augmented Generation (HybridRAG) framework, a novel…
▽ More
Retrieval augmentation enhances performance of traditional language models by incorporating additional context. However, the computational demands for retrieval augmented large language models (LLMs) pose a challenge when applying them to real-time tasks, such as composition assistance. To address this limitation, we propose the Hybrid Retrieval-Augmented Generation (HybridRAG) framework, a novel approach that efficiently combines a cloud-based LLM with a smaller, client-side, language model through retrieval augmented memory. This integration enables the client model to generate effective responses, benefiting from the LLM's capabilities and contextual information. Additionally, through an asynchronous memory update mechanism, the client model can deliver real-time completions swiftly to user inputs without the need to wait for responses from the cloud. Our experiments on five benchmark datasets demonstrate that HybridRAG significantly improves utility over client-only models while maintaining low latency.
△ Less
Submitted 5 February, 2024; v1 submitted 8 August, 2023;
originally announced August 2023.
-
Bayesian Estimation of Differential Privacy
Authors:
Santiago Zanella-Béguelin,
Lukas Wutschitz,
Shruti Tople,
Ahmed Salem,
Victor Rühle,
Andrew Paverd,
Mohammad Naseri,
Boris Köpf,
Daniel Jones
Abstract:
Algorithms such as Differentially Private SGD enable training machine learning models with formal privacy guarantees. However, there is a discrepancy between the protection that such algorithms guarantee in theory and the protection they afford in practice. An emerging strand of work empirically estimates the protection afforded by differentially private training as a confidence interval for the p…
▽ More
Algorithms such as Differentially Private SGD enable training machine learning models with formal privacy guarantees. However, there is a discrepancy between the protection that such algorithms guarantee in theory and the protection they afford in practice. An emerging strand of work empirically estimates the protection afforded by differentially private training as a confidence interval for the privacy budget $\varepsilon$ spent on training a model. Existing approaches derive confidence intervals for $\varepsilon$ from confidence intervals for the false positive and false negative rates of membership inference attacks. Unfortunately, obtaining narrow high-confidence intervals for $ε$ using this method requires an impractically large sample size and training as many models as samples. We propose a novel Bayesian method that greatly reduces sample size, and adapt and validate a heuristic to draw more than one sample per trained model. Our Bayesian method exploits the hypothesis testing interpretation of differential privacy to obtain a posterior for $\varepsilon$ (not just a confidence interval) from the joint posterior of the false positive and false negative rates of membership inference attacks. For the same sample size and confidence, we derive confidence intervals for $\varepsilon$ around 40% narrower than prior work. The heuristic, which we adapt from label-only DP, can be used to further reduce the number of trained models needed to get enough samples by up to 2 orders of magnitude.
△ Less
Submitted 15 June, 2022; v1 submitted 10 June, 2022;
originally announced June 2022.
-
Privacy Regularization: Joint Privacy-Utility Optimization in Language Models
Authors:
Fatemehsadat Mireshghallah,
Huseyin A. Inan,
Marcello Hasegawa,
Victor Rühle,
Taylor Berg-Kirkpatrick,
Robert Sim
Abstract:
Neural language models are known to have a high capacity for memorization of training samples. This may have serious privacy implications when training models on user content such as email correspondence. Differential privacy (DP), a popular choice to train models with privacy guarantees, comes with significant costs in terms of utility degradation and disparate impact on subgroups of users. In th…
▽ More
Neural language models are known to have a high capacity for memorization of training samples. This may have serious privacy implications when training models on user content such as email correspondence. Differential privacy (DP), a popular choice to train models with privacy guarantees, comes with significant costs in terms of utility degradation and disparate impact on subgroups of users. In this work, we introduce two privacy-preserving regularization methods for training language models that enable joint optimization of utility and privacy through (1) the use of a discriminator and (2) the inclusion of a triplet-loss term. We compare our methods with DP through extensive evaluation. We show the advantages of our regularizers with favorable utility-privacy trade-off, faster training with the ability to tap into existing optimization approaches, and ensuring uniform treatment of under-represented subgroups.
△ Less
Submitted 15 April, 2021; v1 submitted 12 March, 2021;
originally announced March 2021.
-
Training Data Leakage Analysis in Language Models
Authors:
Huseyin A. Inan,
Osman Ramadan,
Lukas Wutschitz,
Daniel Jones,
Victor Rühle,
James Withers,
Robert Sim
Abstract:
Recent advances in neural network based language models lead to successful deployments of such models, improving user experience in various applications. It has been demonstrated that strong performance of language models comes along with the ability to memorize rare training samples, which poses serious privacy threats in case the model is trained on confidential user content. In this work, we in…
▽ More
Recent advances in neural network based language models lead to successful deployments of such models, improving user experience in various applications. It has been demonstrated that strong performance of language models comes along with the ability to memorize rare training samples, which poses serious privacy threats in case the model is trained on confidential user content. In this work, we introduce a methodology that investigates identifying the user content in the training data that could be leaked under a strong and realistic threat model. We propose two metrics to quantify user-level data leakage by measuring a model's ability to produce unique sentence fragments within training data. Our metrics further enable comparing different models trained on the same data in terms of privacy. We demonstrate our approach through extensive numerical studies on both RNN and Transformer based models. We further illustrate how the proposed metrics can be utilized to investigate the efficacy of mitigations like differentially private training or API hardening.
△ Less
Submitted 22 February, 2021; v1 submitted 13 January, 2021;
originally announced January 2021.
-
Analyzing Information Leakage of Updates to Natural Language Models
Authors:
Santiago Zanella-Béguelin,
Lukas Wutschitz,
Shruti Tople,
Victor Rühle,
Andrew Paverd,
Olga Ohrimenko,
Boris Köpf,
Marc Brockschmidt
Abstract:
To continuously improve quality and reflect changes in data, machine learning applications have to regularly retrain and update their core models. We show that a differential analysis of language model snapshots before and after an update can reveal a surprising amount of detailed information about changes in the training data. We propose two new metrics---\emph{differential score} and \emph{diffe…
▽ More
To continuously improve quality and reflect changes in data, machine learning applications have to regularly retrain and update their core models. We show that a differential analysis of language model snapshots before and after an update can reveal a surprising amount of detailed information about changes in the training data. We propose two new metrics---\emph{differential score} and \emph{differential rank}---for analyzing the leakage due to updates of natural language models. We perform leakage analysis using these metrics across models trained on several different datasets using different methods and configurations. We discuss the privacy implications of our findings, propose mitigation strategies and evaluate their effect.
△ Less
Submitted 5 August, 2021; v1 submitted 17 December, 2019;
originally announced December 2019.
