Multi-stage machine learning model training for key-value extraction
Techniques for multi-stage training of a machine learning model to extract key-value pairs from documents are disclosed. A system trains a machine learning model using a set of training data including unlabeled documents of various document categories. The initial stage identifies relationships among tokens, or words, numbers, and punctuation, in documents. The system re-trains the machine learning model using a set of training data which includes a particular category of documents while excluding other categories of documents. The second training stage is a supervised machine learning stage in which the training data is labeled to identify key-value pairs in the documents. In the initial training stage, the system sets parameters of the machine learning model to an initial state. In the second stage, the system modifies the parameters of the machine learning model based on the characteristics of the training data set including the documents of the particular category.
1. A non-transitory computer readable medium comprising instructions which, when executed by one or more hardware processors, causes performance of operations comprising:
executing a first training stage for training a machine learning model to generate first vectors encoding semantic information, and at least one of positional and visual information, of first textual content within documents at least by:
accessing a first plurality of training documents associated with a plurality of document categories, the first plurality of training documents including first textual information, first positional information, and first visual information;
training the machine learning model using the first plurality of training documents associated with the plurality of document categories to generate the first vectors encoding the semantic information, and the at least one of the positional and the visual information based on the first textual information, and at least one of the first positional information and the first visual information;
wherein the first training stage generates a first set of parameters for application of the machine learning model;
executing a second training stage for customizing the machine learning model to (a) generate second vectors to encode semantic information, and at least one of positional information and visual information of second textual content within documents of a particular document category and (b) extract key-value pairs at least by:
accessing a second plurality of training documents associated with the particular document category, the second plurality of training documents being tagged with key-value pairs and including second textual information, second positional information, and second visual information;
training the machine learning model using the second plurality of training documents associated with the particular document category to identify key-value pairs based at least in part on vector encodings of the second textual content that encode semantic relationships, and at least one of positional and visual relationships, between components of the key-value pairs,
wherein training the mahcine learning model using the second plurality of training documents results in a trained machien learning model, and
wherein the second training stage fine-tunes the first set of parameters to generate a second set of parameters for application of the machine learning model; and
applying the trained machine learning model with the second set of parameters to a first document of the particular document category to extract a first plurality of key-value pairs from the first document.
2. The non-transitory computer readable medium of claim 1 , wherein training the machine learning model using the second plurality of training documents comprises training the machine learning model to identify the key-value pairs based at least in part on analyzing vector encodings of the second textual content to determine each of: (a) semantic relationships, (b) positional relationships, and (c) visual relationships between components of the key-value pairs.
3. The non-transitory computer readable medium of claim 1 , wherein training the machine learning model using the first plurality of training documents includes generating, by the machine learning model for each document of the first plurality of training documents, a vector encoding the semantic relationships, positional relationships, and visual relationships among the respective textual content in each respective document of the first plurality of training documents.
4. The non-transitory computer readable medium of claim 1 , the operations further comprising:
accessing feedback data corresponding to the key-value pairs extracted from the second plurality of training documents using the trained machine learning model; and
updating the machine learning model based on the feedback data.
5. The non-transitory computer readable medium of claim 4 , wherein accessing the feedback data corresponding to the key-value pairs extracted from the second plurality of training documents using the trained machine learning model comprises:
applying a set of post-model rules to the key-value pairs to group tokens identified as separate values associated with a same key of a key-value pair into a single value of the key-value pair.
6. The non-transitory computer readable medium of claim 5 , wherein grouping the tokens identified as separate values associated with a same key of a key-value pair into the single value of the key-value pair comprises expanding a bounding box associated with the single value to encompass a first token identified by the machine learning model as a first value and a second token identified by the machine learning model as a second value.
7. The non-transitory computer readable medium of claim 4 , wherein obtaining the feedback data corresponding to the key-value pairs extracted from the second plurality of training documents using the trained machine learning model comprises:
applying a set of post-model rules to the key-value pairs to separate a single value of a key-value pair identified by the machine learning model into a first value corresponding to a first token and a second value corresponding to a second token.
8. The non-transitory computer readable medium of claim 1 , wherein the first plurality of training documents includes documents of the particular document category.
9. The non-transitory computer readable medium of claim 1 , wherein the particular document category comprises an invoice category.
10. The non-transitory computer readable medium of claim 1 , wherein the first plurality of training documents includes documents of a different category than the particular document category, and
wherein the second plurality of training documents excludes documents of any category different from the particular document category.
11. The non-transitory computer readable medium of claim 1 , wherein the second positional information includes bounding box location data associated with tokens in the second plurality of training documents, and
wherein the second visual information includes at least one of size data, shape data, and color data associated with the tokens in the second plurality of training documents.
12. The non-transitory computer readable medium of claim 1 , wherein the operations further comprise:
identifying tokens in each document of the first plurality of training documents; and
for each token identified in the first plurality of training documents: extracting text content, and at least one of (a) bounding box location data and (b) shape and color data associated with the token.
13. The non-transitory computer readable medium of claim 1 , wherein obtaining the second plurality of training documents associated with the particular document category comprises:
obtaining a first subset of training documents comprising historical documents of the particular document category; and
obtaining a second subset of training documents comprising synthetic documents of the particular document category.
14. The non-transitory computer readable medium of claim 13 , wherein obtaining the second subset of training documents comprises applying a set of document-generation rules to generate an image of a document of the particular document category and key-value pair labels associated with tokens in the image of the document.
15. The non-transitory computer readable medium of claim 1 , wherein the first plurality of training documents are unlabeled training documents,
wherein the second plurality of training documents are labeled training documents,
wherein training the machine learning model using the first plurality of training documents comprises applying a first algorithm implementing an unsupervised machine learning operation to the first plurality of training documents, and
wherein training the machine learning model using the second plurality of training documents comprises applying a second algorithm implementing a supervised machine learning operation to the second plurality of training documents.
16. A method comprising:
executing a first training stage for training a machine learning model to generate first vectors encoding semantic information, and at least one of positional and visual information, of first textual content within documents at least by:
accessing a first plurality of training documents associated with a plurality of document categories, the first plurality of training documents including first textual information, first positional information, and first visual information;
training the machine learning model using the first plurality of training documents associated with the plurality of document categories to generate the first vectors encoding the semantic information, and the at least one of the positional and the visual information based on the first textual information, and at least one of the first positional information and the first visual information;
wherein the first training stage generates a first set of parameters for application of the machine learning model;
executing a second training stage for customizing the machine learning model to (a) generate second vectors to encode semantic information, and at least one of positional information and visual information of second textual content within documents of a particular document category and (b) extract key-value pairs at least by:
accessing a second plurality of training documents associated with the particular document category, the second plurality of training documents being tagged with key-value pairs and including second textual information, second positional information, and second visual information;
training the machine learning model using the second plurality of training documents associated with the particular document category to identify key-value pairs based at least in part on analyzing vector encodings of the second textual content to determine semantic relationships, and at least one of positional and visual relationships, between components of the key-value pairs,
wherein training the machine learning model using the second plurality of training documents results in a trained machine learning model, and
wherein the second training stage fine-tunes the first set of parameters to generate a second set of parameters for application of the machine learning model; and
applying the trained machine learning model with the second set of parameters to a first document of the particular document category to extract a first plurality of key-value pairs from the first document.
17. The method of claim 16 , wherein training the machine learning model using the second plurality of training documents comprises training the machine learning model to identify the key-value pairs based at least in part on analyzing vector encodings of the second textual content to determine each of: (a) semantic relationships, (b) positional relationships, and (c) visual relationships between components of the key-value pairs.
18. The method of claim 16 , wherein training the machine learning model using the first plurality of training documents includes generating, by the machine learning model for each document of the first plurality of training documents, a vector encoding the semantic relationships, positional relationships, and visual relationships among the respective textual content in each respective document of the first plurality of training documents.
19. The method of claim 16 , further comprising:
accessing feedback data corresponding to the key-value pairs extracted from the second plurality of training documents using the trained machine learning model; and
updating the machine learning model based on the feedback data.
20. A system comprising:
one or more processors; and
memory storing instructions that, when executed by the one or more processors, cause the system to perform operations comprising:
executing a first training stage for training a machine learning model to generate first vectors encoding semantic information, and at least one of positional and visual information, of first textual content within documents at least by:
accessing a first plurality of training documents associated with a plurality of document categories, the first plurality of training documents including first textual information, first positional information, and first visual information;
training the machine learning model using the first plurality of training documents associated with the plurality of document categories to generate the first vectors encoding the semantic information, and the at least one of the positional and the visual information based on the first textual information, and at least one of the first positional information and the first visual information;
wherein the first training stage generates a first set of parameters for application of the machine learning model;
executing a second training stage for customizing the machine learning model to (a) generate second vectors to encode semantic information, and at least one of positional information and visual information of second textual content within documents of a particular document category and (b) extract key-value pairs at least by:
accessing a second plurality of training documents associated with the particular document category, the second plurality of training documents being tagged with key-value pairs and including second textual information, second positional information, and second visual information;
training the machine learning model using the second plurality of training documents associated with the particular document category to identify key-value pairs based at least in part on vector encodings of the second textual content that encode semantic relationships, and at least one of positional and visual relationships, between components of the key-value pairs,
wherein training the machine learning model using the second plurality of training documents results in a trained machine learning model, and
wherein the second training stage fine-tunes the first set of parameters to generate a second set of parameters for application of the machine learning model; and
applying the trained machine learning model with the second set of parameters to a first document of the particular document category to extract a first plurality of key-value pairs from the first document.