Malware detection for documents using knowledge distillation assisted learning
The detection of malicious documents using knowledge distillation assisted learning is disclosed. A document is received for maliciousness determination. A likelihood that the received document represents a threat is determined. The determination is made, at least in part, using a raw bytes model that was trained, at least in part, using image model prediction probabilities. A verdict for the document is provided as output based at least in part on the determined likelihood.
1 . A system, comprising:
a processor configured to:
receive a document for a maliciousness determination;
determine a likelihood that the document received for the maliciousness determination represents a threat, at least in part using a raw bytes model, wherein the raw bytes model was trained, at least in part, using image model prediction probabilities obtained from an image model trained, at least in part, using a plurality of images generated in a pipeline that converts a plurality of documents labeled as malicious into a corresponding plurality of rendered images of those documents, wherein at least some of the corresponding plurality of rendered images are rendered, collectively, from a multi-page document labeled as a malicious document; and
provide as output a verdict for the document based at least in part on the determined likelihood, wherein the verdict is that the received document is benign; and
a memory coupled to the processor and configured to provide the processor with instructions.
2 . The system of claim 1 , wherein determining the likelihood does not require converting a portion of the received document into an image.
3 . The system of claim 1 , wherein, prior to training the image model, an image hash-based filtering operation is performed on at least some of the plurality of images labeled as malicious documents.
4 . The system of claim 3 , wherein filtered images are stored using a TFRecord data format.
5 . The system of claim 1 , wherein the processor is further configured to generate the image model.
6 . The system of claim 1 , wherein the image model is a convolutional neural network model.
7 . The system of claim 1 , wherein the raw bytes model is a convolutional neural network model.
8 . The system of claim 1 , wherein, at least in part in response to receiving an indication of a false positive result, the image model is retrained using a benign data set that includes the false positive result.
9 . The system of claim 1 , wherein the document is a Microsoft Office document.
10 . The system of claim 1 , wherein training the raw bytes model includes using knowledge distillation.
11 . The system of claim 1 , wherein a loss function used in training the raw bytes model comprises both self loss and imitation loss.
12 . A method, comprising:
receiving a document for a maliciousness determination;
determining a likelihood that the document received for the maliciousness determination represents a threat, at least in part by using a raw bytes model, wherein the raw bytes model was trained, at least in part, using image model prediction probabilities obtained from an image model trained, at least in part, using a plurality of images generated in a pipeline that converts a plurality of documents labeled as malicious into a corresponding plurality of rendered images of those documents, wherein at least some of the corresponding plurality of rendered images are rendered, collectively, from a multi-page document labeled as a malicious document; and
providing as output a verdict for the document based at least in part on the determined likelihood, wherein the verdict is that the received document is benign.
13 . A computer program product embodied in a non-transitory computer readable medium and comprising computer instructions for:
receiving a document for a maliciousness determination;
determining a likelihood that the document received for the maliciousness determination represents a threat, at least in part by using a raw bytes model, wherein the raw bytes model was trained, at least in part, using image model prediction probabilities obtained from an image model trained, at least in part, using a plurality of images generated in a pipeline that converts a plurality of documents labeled as malicious into a corresponding plurality of rendered images of those documents, wherein at least some of the corresponding plurality of rendered images are rendered, collectively, from a multi-page document labeled as a malicious document; and
providing as output a verdict for the document based at least in part on the determined likelihood, wherein the verdict is that the received document is benign.
14 . The method of claim 12 , wherein determining the likelihood does not require converting a portion of the received document into an image.
15 . The method of claim 12 , wherein, prior to training the image model, an image hash-based filtering operation is performed on at least some of the plurality of images labeled as malicious documents.
16 . The method of claim 15 , wherein filtered images are stored using a TFRecord data format.
17 . The method of claim 12 , further comprising generating the image model.
18 . The method of claim 12 , wherein the image model is a convolutional neural network model.
19 . The method of claim 12 , wherein the raw bytes model is a convolutional neural network model.
20 . The method of claim 12 , wherein, at least in part in response to receiving an indication of a false positive result, the image model is retrained using a benign data set that includes the false positive result.
21 . The method of claim 12 , wherein the document is a Microsoft Office document.
22 . The method of claim 12 , wherein training the raw bytes model includes using knowledge distillation.
23 . The method of claim 12 , wherein a loss function used in training the raw bytes model comprises both self loss and imitation loss.