Data loss protection (DLP) utilizing distilled large language models (LLMs)
Systems and methods for Data Loss Protection (DLP) utilizing distilled models include receiving a plurality of general data predictions from a teacher model; determining one or more strengths of the teacher model based on the received general data predictions; generating a synthetic dataset based on the one or more strengths of the teacher model; providing the synthetic dataset to the teacher model and receiving a plurality of synthetic data predictions from the teacher model based thereon; and performing knowledge distillation on a student model based on the synthetic data predictions received from the teacher model to produce a distilled model. The distilled model is then used in production for classifying inputs to a DLP system.
1 . A method comprising steps of:
receiving a plurality of general data predictions from a teacher model on a labeled Data Loss Protection (DLP) dataset comprising a plurality of content categories;
determining one or more strengths of the teacher model based on the received general data predictions by computing, for each of the plurality of content categories, a category-specific classification accuracy of the teacher model relative to labels in the labeled DLP dataset and selecting one or more well-performing categories having the category-specific classification accuracy above a threshold;
generating a synthetic dataset based on the one or more strengths of the teacher model by querying a Large Language Model (LLM) to generate category-specific synthetic documents corresponding to the selected one or more well-performing categories;
providing the synthetic dataset to the teacher model and receiving a plurality of synthetic data predictions from the teacher model based thereon; and
performing knowledge distillation on a student model based on the synthetic data predictions received from the teacher model to produce a distilled model for content-category classification in a DLP system.
2 . The method of claim 1 , wherein determining the one or more strengths comprises determining respective strengths of a plurality of teacher models in different categories, and wherein providing the synthetic dataset to the teacher model comprises providing first category-specific synthetic documents to a first teacher model and second category-specific synthetic documents to a second teacher model, and combining corresponding synthetic data predictions from the plurality of teacher models into a single distillation dataset used to perform the knowledge distillation on the student model.
3 . The method of claim 1 , wherein prior to receiving the plurality of general data predictions from the teacher model, the steps comprise providing, to the teacher model, a general DLP dataset.
4 . The method of claim 1 , wherein the plurality of general data predictions and plurality of synthetic data predictions comprise content category classification predictions.
5 . The method of claim 4 , wherein determining the one or more strengths of the teacher model comprises (i) evaluating, for each content category, the teacher model's classification accuracy using labels of the general Data Loss Protection (DLP) dataset, and (ii) determining one or more categories in which the teacher model performs classification with an accuracy above a threshold.
6 . The method of claim 1 , wherein generating the synthetic dataset comprises utilizing an LLM to generate a plurality of inputs associated with the one or more strengths of the teacher model, wherein the synthetic dataset comprises the plurality of inputs.
7 . The method of claim 1 , wherein the steps further comprise:
utilizing the distilled model in production for classifying inputs to a DLP system.
8 . The method of claim 1 , wherein the steps further comprise:
receiving an input comprising data in any of a plurality of formats;
processing the input via the distilled model to classify the input into a category of a plurality of categories; and
providing an indication of the category of the plurality of categories.
9 . The method of claim 8 , wherein the steps further comprise:
processing the input to determine whether or not the data includes sensitive data prior to processing the input for classification.
10 . The method of claim 8 , wherein the plurality of formats comprise text formats, image formats, audio formats, video formats, source code, and a combination thereof.
11 . A non-transitory computer-readable medium comprising instructions that, when executed, cause one or more processors to perform steps of:
receiving a plurality of general data predictions from a teacher model on a labeled Data Loss Protection (DLP) dataset comprising a plurality of content categories;
determining one or more strengths of the teacher model based on the received general data predictions by computing, for each of the plurality of content categories a category-specific classification accuracy of the teacher model relative to labels in the labeled DLP dataset and selecting one or more well-performing categories having the category-specific classification accuracy above a threshold;
generating a synthetic dataset based on the one or more strengths of the teacher model by querying a Large Language Model (LLM) to generate category-specific synthetic documents corresponding to the selected one or more well-performing categories;
providing the synthetic dataset to the teacher model and receiving a plurality of synthetic data predictions from the teacher model based thereon; and
performing knowledge distillation on a student model based on the synthetic data predictions received from the teacher model to produce a distilled model for content-category classification in a DLP system.
12 . The non-transitory computer-readable medium of claim 11 , wherein determining the one or more strengths comprises determining respective strengths of a plurality of teacher models in different categories, and wherein providing the synthetic dataset to the teacher model comprises providing first category-specific synthetic documents to a first teacher model and second category-specific synthetic documents to a second teacher model, and combining corresponding a synthetic data predictions from the plurality of teacher models into a single distillation dataset used to perform the knowledge distillation on the student model.
13 . The non-transitory computer-readable medium of claim 11 , wherein prior to receiving the plurality of general data predictions from the teacher model, the steps comprise providing, to the teacher model, a general DLP dataset.
14 . The non-transitory computer-readable medium of claim 11 , wherein the plurality of general data predictions and plurality of synthetic data predictions comprise content category classification predictions.
15 . The non-transitory computer-readable medium of claim 14 , wherein determining the one or more strengths of the teacher model comprises (i) evaluating, for each content category, the teacher model's classification accuracy using labels of the general Data Loss Protection (DLP) dataset, and (ii) determining one or more categories in which the teacher model performs classification with an accuracy above a threshold.
16 . The non-transitory computer-readable medium of claim 11 , wherein generating the synthetic dataset comprises utilizing an LLM to generate a plurality of inputs associated with the one or more strengths of the teacher model, wherein the synthetic dataset comprises the plurality of inputs.
17 . The non-transitory computer-readable medium of claim 11 , wherein the steps further comprise:
utilizing the distilled model in production for classifying inputs to a DLP system.
18 . The non-transitory computer-readable medium of claim 11 , wherein the steps further comprise:
receiving an input comprising data in any of a plurality of formats;
processing the input via the distilled model to classify the input into a category of a plurality of categories; and
providing an indication of the category of the plurality of categories.
19 . The non-transitory computer-readable medium of claim 18 , wherein the steps further comprise:
processing the input to determine whether or not the data includes sensitive data prior to processing the input for classification.
20 . The non-transitory computer-readable medium of claim 18 , wherein the plurality of formats comprise text formats, image formats, audio formats, video formats, source code, and a combination thereof.