-
Accelerating Production LLMs with Combined Token/Embedding Speculators
Authors:
Davis Wertheimer,
Joshua Rosenkranz,
Thomas Parnell,
Sahil Suneja,
Pavithra Ranganathan,
Raghu Ganti,
Mudhakar Srivatsa
Abstract:
This technical report describes the design and training of novel speculative decoding draft models, for accelerating the inference speeds of large language models in a production environment. By conditioning draft predictions on both context vectors and sampled tokens, we can train our speculators to efficiently predict high-quality n-grams, which the base model then accepts or rejects. This allow…
▽ More
This technical report describes the design and training of novel speculative decoding draft models, for accelerating the inference speeds of large language models in a production environment. By conditioning draft predictions on both context vectors and sampled tokens, we can train our speculators to efficiently predict high-quality n-grams, which the base model then accepts or rejects. This allows us to effectively predict multiple tokens per inference forward pass, accelerating wall-clock inference speeds of highly optimized base model implementations by a factor of 2-3x. We explore these initial results and describe next steps for further improvements.
△ Less
Submitted 6 June, 2024; v1 submitted 29 April, 2024;
originally announced April 2024.
-
Automated Code generation for Information Technology Tasks in YAML through Large Language Models
Authors:
Saurabh Pujar,
Luca Buratti,
Xiaojie Guo,
Nicolas Dupuis,
Burn Lewis,
Sahil Suneja,
Atin Sood,
Ganesh Nalawade,
Matthew Jones,
Alessandro Morari,
Ruchir Puri
Abstract:
The recent improvement in code generation capabilities due to the use of large language models has mainly benefited general purpose programming languages. Domain specific languages, such as the ones used for IT Automation, have received far less attention, despite involving many active developers and being an essential component of modern cloud platforms. This work focuses on the generation of Ans…
▽ More
The recent improvement in code generation capabilities due to the use of large language models has mainly benefited general purpose programming languages. Domain specific languages, such as the ones used for IT Automation, have received far less attention, despite involving many active developers and being an essential component of modern cloud platforms. This work focuses on the generation of Ansible-YAML, a widely used markup language for IT Automation. We present Ansible Wisdom, a natural-language to Ansible-YAML code generation tool, aimed at improving IT automation productivity. Ansible Wisdom is a transformer-based model, extended by training with a new dataset containing Ansible-YAML. We also develop two novel performance metrics for YAML and Ansible to capture the specific characteristics of this domain. Results show that Ansible Wisdom can accurately generate Ansible script from natural language prompts with performance comparable or better than existing state of the art code generation models. In few-shot settings we asses the impact of training with Ansible, YAML data and compare with different baselines including Codex-Davinci-002. We also show that after finetuning, our Ansible specific model (BLEU: 66.67) can outperform a much larger Codex-Davinci-002 (BLEU: 50.4) model, which was evaluated in few shot settings.
△ Less
Submitted 23 May, 2023; v1 submitted 2 May, 2023;
originally announced May 2023.
-
Study of Distractors in Neural Models of Code
Authors:
Md Rafiqul Islam Rabin,
Aftab Hussain,
Sahil Suneja,
Mohammad Amin Alipour
Abstract:
Finding important features that contribute to the prediction of neural models is an active area of research in explainable AI. Neural models are opaque and finding such features sheds light on a better understanding of their predictions. In contrast, in this work, we present an inverse perspective of distractor features: features that cast doubt about the prediction by affecting the model's confid…
▽ More
Finding important features that contribute to the prediction of neural models is an active area of research in explainable AI. Neural models are opaque and finding such features sheds light on a better understanding of their predictions. In contrast, in this work, we present an inverse perspective of distractor features: features that cast doubt about the prediction by affecting the model's confidence in its prediction. Understanding distractors provide a complementary view of the features' relevance in the predictions of neural models. In this paper, we apply a reduction-based technique to find distractors and provide our preliminary results of their impacts and types. Our experiments across various tasks, models, and datasets of code reveal that the removal of tokens can have a significant impact on the confidence of models in their predictions and the categories of tokens can also play a vital role in the model's confidence. Our study aims to enhance the transparency of models by emphasizing those tokens that significantly influence the confidence of the models.
△ Less
Submitted 3 March, 2023;
originally announced March 2023.
-
VELVET: a noVel Ensemble Learning approach to automatically locate VulnErable sTatements
Authors:
Yangruibo Ding,
Sahil Suneja,
Yunhui Zheng,
Jim Laredo,
Alessandro Morari,
Gail Kaiser,
Baishakhi Ray
Abstract:
Automatically locating vulnerable statements in source code is crucial to assure software security and alleviate developers' debugging efforts. This becomes even more important in today's software ecosystem, where vulnerable code can flow easily and unwittingly within and across software repositories like GitHub. Across such millions of lines of code, traditional static and dynamic approaches stru…
▽ More
Automatically locating vulnerable statements in source code is crucial to assure software security and alleviate developers' debugging efforts. This becomes even more important in today's software ecosystem, where vulnerable code can flow easily and unwittingly within and across software repositories like GitHub. Across such millions of lines of code, traditional static and dynamic approaches struggle to scale. Although existing machine-learning-based approaches look promising in such a setting, most work detects vulnerable code at a higher granularity -- at the method or file level. Thus, developers still need to inspect a significant amount of code to locate the vulnerable statement(s) that need to be fixed.
This paper presents VELVET, a novel ensemble learning approach to locate vulnerable statements. Our model combines graph-based and sequence-based neural networks to successfully capture the local and global context of a program graph and effectively understand code semantics and vulnerable patterns. To study VELVET's effectiveness, we use an off-the-shelf synthetic dataset and a recently published real-world dataset. In the static analysis setting, where vulnerable functions are not detected in advance, VELVET achieves 4.5x better performance than the baseline static analyzers on the real-world data. For the isolated vulnerability localization task, where we assume the vulnerability of a function is known while the specific vulnerable statement is unknown, we compare VELVET with several neural networks that also attend to local and global context of code. VELVET achieves 99.6% and 43.6% top-1 accuracy over synthetic data and real-world data, respectively, outperforming the baseline deep-learning models by 5.3-29.0%.
△ Less
Submitted 12 January, 2022; v1 submitted 20 December, 2021;
originally announced December 2021.
-
Data-Driven AI Model Signal-Awareness Enhancement and Introspection
Authors:
Sahil Suneja,
Yufan Zhuang,
Yunhui Zheng,
Jim Laredo,
Alessandro Morari
Abstract:
AI modeling for source code understanding tasks has been making significant progress, and is being adopted in production development pipelines. However, reliability concerns, especially whether the models are actually learning task-related aspects of source code, are being raised. While recent model-probing approaches have observed a lack of signal awareness in many AI-for-code models, i.e. models…
▽ More
AI modeling for source code understanding tasks has been making significant progress, and is being adopted in production development pipelines. However, reliability concerns, especially whether the models are actually learning task-related aspects of source code, are being raised. While recent model-probing approaches have observed a lack of signal awareness in many AI-for-code models, i.e. models not capturing task-relevant signals, they do not offer solutions to rectify this problem. In this paper, we explore data-driven approaches to enhance models' signal-awareness: 1) we combine the SE concept of code complexity with the AI technique of curriculum learning; 2) we incorporate SE assistance into AI models by customizing Delta Debugging to generate simplified signal-preserving programs, augmenting them to the training dataset. With our techniques, we achieve up to 4.8x improvement in model signal awareness. Using the notion of code complexity, we further present a novel model learning introspection approach from the perspective of the dataset.
△ Less
Submitted 7 January, 2022; v1 submitted 10 November, 2021;
originally announced November 2021.
-
Software Vulnerability Detection via Deep Learning over Disaggregated Code Graph Representation
Authors:
Yufan Zhuang,
Sahil Suneja,
Veronika Thost,
Giacomo Domeniconi,
Alessandro Morari,
Jim Laredo
Abstract:
Identifying vulnerable code is a precautionary measure to counter software security breaches. Tedious expert effort has been spent to build static analyzers, yet insecure patterns are barely fully enumerated. This work explores a deep learning approach to automatically learn the insecure patterns from code corpora. Because code naturally admits graph structures with parsing, we develop a novel gra…
▽ More
Identifying vulnerable code is a precautionary measure to counter software security breaches. Tedious expert effort has been spent to build static analyzers, yet insecure patterns are barely fully enumerated. This work explores a deep learning approach to automatically learn the insecure patterns from code corpora. Because code naturally admits graph structures with parsing, we develop a novel graph neural network (GNN) to exploit both the semantic context and structural regularity of a program, in order to improve prediction performance. Compared with a generic GNN, our enhancements include a synthesis of multiple representations learned from the several parsed graphs of a program, and a new training loss metric that leverages the fine granularity of labeling. Our model outperforms multiple text, image and graph-based approaches, across two real-world datasets.
△ Less
Submitted 7 September, 2021;
originally announced September 2021.
-
Probing Model Signal-Awareness via Prediction-Preserving Input Minimization
Authors:
Sahil Suneja,
Yunhui Zheng,
Yufan Zhuang,
Jim Laredo,
Alessandro Morari
Abstract:
This work explores the signal awareness of AI models for source code understanding. Using a software vulnerability detection use case, we evaluate the models' ability to capture the correct vulnerability signals to produce their predictions. Our prediction-preserving input minimization (P2IM) approach systematically reduces the original source code to a minimal snippet which a model needs to maint…
▽ More
This work explores the signal awareness of AI models for source code understanding. Using a software vulnerability detection use case, we evaluate the models' ability to capture the correct vulnerability signals to produce their predictions. Our prediction-preserving input minimization (P2IM) approach systematically reduces the original source code to a minimal snippet which a model needs to maintain its prediction. The model's reliance on incorrect signals is then uncovered when the vulnerability in the original code is missing in the minimal snippet, both of which the model however predicts as being vulnerable. We measure the signal awareness of models using a new metric we propose- Signal-aware Recall (SAR). We apply P2IM on three different neural network architectures across multiple datasets. The results show a sharp drop in the model's Recall from the high 90s to sub-60s with the new metric, highlighting that the models are presumably picking up a lot of noise or dataset nuances while learning their vulnerability detection logic. Although the drop in model performance may be perceived as an adversarial attack, but this isn't P2IM's objective. The idea is rather to uncover the signal-awareness of a black-box model in a data-driven manner via controlled queries. SAR's purpose is to measure the impact of task-agnostic model training, and not to suggest a shortcoming in the Recall metric. The expectation, in fact, is for SAR to match Recall in the ideal scenario where the model truly captures task-specific signals.
△ Less
Submitted 22 June, 2021; v1 submitted 25 November, 2020;
originally announced November 2020.
-
Learning to map source code to software vulnerability using code-as-a-graph
Authors:
Sahil Suneja,
Yunhui Zheng,
Yufan Zhuang,
Jim Laredo,
Alessandro Morari
Abstract:
We explore the applicability of Graph Neural Networks in learning the nuances of source code from a security perspective. Specifically, whether signatures of vulnerabilities in source code can be learned from its graph representation, in terms of relationships between nodes and edges. We create a pipeline we call AI4VA, which first encodes a sample source code into a Code Property Graph. The extra…
▽ More
We explore the applicability of Graph Neural Networks in learning the nuances of source code from a security perspective. Specifically, whether signatures of vulnerabilities in source code can be learned from its graph representation, in terms of relationships between nodes and edges. We create a pipeline we call AI4VA, which first encodes a sample source code into a Code Property Graph. The extracted graph is then vectorized in a manner which preserves its semantic information. A Gated Graph Neural Network is then trained using several such graphs to automatically extract templates differentiating the graph of a vulnerable sample from a healthy one. Our model outperforms static analyzers, classic machine learning, as well as CNN and RNN-based deep learning models on two of the three datasets we experiment with. We thus show that a code-as-graph encoding is more meaningful for vulnerability detection than existing code-as-photo and linear sequence encoding approaches. (Submitted Oct 2019, Paper #28, ICST)
△ Less
Submitted 15 June, 2020;
originally announced June 2020.
-
Secure Extensibility for System State Extraction via Plugin Sandboxing
Authors:
Sahil Suneja,
Canturk Isci
Abstract:
We introduce a new mechanism to securely extend systems data collection software with potentially untrusted third-party code. Unlike existing tools which run extension modules or plugins directly inside the monitored endpoint (the guest), we run plugins inside a specially crafted sandbox, so as to protect the guest as well as the software core. To get the right mix of accessibility and constraints…
▽ More
We introduce a new mechanism to securely extend systems data collection software with potentially untrusted third-party code. Unlike existing tools which run extension modules or plugins directly inside the monitored endpoint (the guest), we run plugins inside a specially crafted sandbox, so as to protect the guest as well as the software core. To get the right mix of accessibility and constraints required for systems data extraction, we create our sandbox by combining multiple features exported by an unmodified kernel. We have tested its applicability by successfully sandboxing plugins of an opensourced data collection software for containerized guest systems. We have also verified its security posture in terms of successful containment of several exploits, which would have otherwise directly impacted a guest, if shipped inside third-party plugins.
△ Less
Submitted 20 May, 2019;
originally announced May 2019.