New pretraining tasks enable better document understanding

DocFormerV2 makes sense of documents using local features, outperforming much bigger models.

In the digital era, when documents are generated and distributed at unprecedented rates, automatically understanding them is crucial. Consider the tasks of extracting payment information from invoices or digitizing historical records, where layouts and handwritten notes play an important role in understanding context. These scenarios highlight the complexity of document understanding, which requires not just recognizing text but also interpreting visual elements and their spatial relationships.

A mailing label from Harvard University Press, with several preprinted, labeled spaces for shipping data, such as "sold to", "ship to", and "date".
Visual document understanding (VDU): A snippet of a document receipt from the DocVQA dataset. A VDU model might be asked to predict the “sold to” address (visual question answering), to predict all relations (“sold to” → <address>, “ship to” → <address>), or to infer information from the table at the top of the document.

At this year’s meeting of the Association for the Advancement of Artificial Intelligence (AAAI 2024), we proposed a model we call DocFormerv2, which doesn't just read documents but understands them, making sense of both textual and visual information in a way that mimics human comprehension. For example, just as a person might infer a report's key points from its layout, headings, text, and associated tables, DocFormerv2 analyzes these elements collectively to grasp the document's overall message.

Related content
Method preserves knowledge encoded in teacher model’s attention heads even when student model has fewer of them.

Unlike its predecessors, DocFormerv2 employs a transformer-based architecture that excels in capturing local features within documents — small, specific details such as the style of a font, the way a paragraph is arranged, or how pictures are placed next to text. This means it can discern the significance of layout elements with higher accuracy than prior models.

A standout feature of DocFormerv2 is its use of self-supervised learning, the approach used in many of today’s most successful AI models, such as GPT. Self-supervised learning uses unannotated data, which enables training on enormous public datasets. In language modeling, for instance, next-token prediction (used by GPT) or masked-token prediction (used by T5 or BERT) are popular.

A schematic of the DocFormerv2 architecture, which takes as input both images of the document and the associated OCR output, along with the spatial coordinates of text, and which is trained on two tasks, token to line and token to grid.
DocFormerv2 architecture.

For DocFormerv2, in addition to standard masked-token prediction, we propose two additional tasks, token-to-line prediction and token-to-grid assignment. These tasks are designed to deepen the model's understanding of the intricate relationship between text and its spatial arrangement within documents. Let’s take a closer look at them.

Token to line

The token-to-line task trains DocFormerv2 to recognize how textual elements align within lines, imparting an understanding that goes beyond mere words to include the flow and structure of text as it appears in documents. This follows the intuition that most of the information needed for key-value prediction in a form or for visual question answering (VQA) is on either the same line or adjacent lines of a document. For instance, in the diagram below, in order to predict the value for "Total" (box a), the model has to look in the same line (box d, "$4.32"). Through this type of task, the model learns to give importance to information about the relative positions of tokens and its semantic implications.

At left is a store receipt with the labels "state tax", "total", and "change" surrounded by red boxes and labeled, respectively, b, a, and c and the total amount of the charge, $4.32, labeled d. At right is a product order form with a 16-cell red grid superimposed on it, each cell labeled with a blue number (1-16).
Novel document pretraining tasks: token to line and token to grid.

Token to grid

Semantic information varies across a document's different regions. For instance, financial documents might have headers at the top, fillable information in the middle, and footers or instructions at the bottom. Page numbers are usually found at the top or bottom of a document, while company names in receipts or invoices often appear at the top. Understanding a document accurately requires recognizing how its content is organized within a specific visual layout and structure. Armed with this intuition, the token-to-grid task pairs the semantics of texts with their locations (visual, spatial, or both) in the document. Specifically, a grid is superimposed on the document, and each OCR token is assigned a grid number. During training, DocFormerv2 is tasked with predicting the grid number for each token.

Target tasks and impact

On nine different datasets covering a range of document-understanding tasks, DocFormerv2 outperforms previous comparably sized models and even does better than much larger models — including one that is 106 times as big as DocFormerv2. Since text from documents is extracted using OCR models, which do make prediction errors, we also show that DocFormerv2 is more resilient to OCR errors than its predecessors.

One of the tasks we trained DocFormerv2 on is table VQA, a challenging task in which the model must answer questions about tables (with either images, text, or both as input). DocFormerv2 achieved 4.3% absolute performance improvement over the next best model.

A spreadsheet table labeled "FM radio stations" whose column labels include "frequency", "call sign", "name", and "format". The entries in the "call sign" column are "KUSK", "KKYA", "KDAM", "WNAX-FM", and "KVHT". "WNAX-FM" is surrounded by a red box.
For the question "Which of these stations does not have a 'k’ in its call sign?", DocFormerv2 correctly answers "WNAX-FM" (fourth row, second column). This requires reasoning over spatial, visual, and language features.
A spreadsheet table with three columns, labeled "District", "Location", and "Communities served". Four of the eight cells in the "Communities served" column — those whose entries begin "Roman Catholic Diocese of Cleveland" — are surrounded by red boxes.
For the question "How many of the schools serve the Roman Catholic diocese of Cleveland?", DocFormerv2 correctly answers "four". This requires arithmetic counting — a challenging task for machine learning models — and reasoning over multiple rows.
A police boat with the word "Police" written on its hull and, below the picture, the text query "What color is the word 'police' written in?"
In this example, an image and text (from an OCR model) are fed to DocFormerv2 along with the question “What color is the word ‘police’ written in?”. Due to its multimodal nature, DocFormerv2 can “see” the image and correctly answer “white”.

But DocFormerv2 also displayed more-qualitative advantages over its predecessors. Because it’s trained to make sense of local features, DocFormerv2 can answer correctly when asked questions like “Which of these stations do not have a ‘k’ in their call sign?” or “How many of the schools serve the Roman Catholic diocese of Cleveland?” (The second question requires counting — a hard skill to learn.)

In order to show the versatility and generalizability of DocFormerv2, we also tested it on scene-text VQA, a task that’s related to but distinct from document understanding. Again, it significantly outperformed comparably sized predecessors.

While DocFormerv2 has made significant strides in interpreting complex documents, several challenges and exciting opportunities lie ahead, like teaching the model to deal with diverse document layouts and enhancing multimodal integration.

Related content

US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
LU, Luxembourg
The Decision, Science and Technology (DST) team part of the global Reliability Maintenance Engineering (RME) is looking for a Senior Operations Research Scientist interested in solving challenging optimization problems in the maintenance space. Our mission is to leverage the use of data, science, and technology to improve the efficiency of RME maintenance activities, reduce costs, increase safety and promote sustainability while creating frictionless customer experiences. As a Senior OR Scientist in DST you will be focused on leading the design and development of innovative approaches and solutions by leading technical work supporting RME’s Predictive Maintenance (PdM) and Spare Parts (SP) programs. You will connect with world leaders in your field and you will be tackling customer's natural language challenges by carrying out a systematic review of existing solutions. The appropriate choice of methods and their deployment into effective tools will be the key for the success in this role. The successful candidate will be a self-starter comfortable with ambiguity, with strong attention to detail and outstanding ability in balancing technical leadership with strong business judgment to make the right decisions about model and method choices. Key job responsibilities • Provide technical expertise to support team strategies that will take EU RME towards World Class predictive maintenance practices and processes, driving better equipment up-time and lower repair costs with optimized spare parts inventory and placement • Implement an advanced maintenance framework utilizing Machine Learning technologies to drive equipment performance leading to reduced unplanned downtime • Provide technical expertise to support the development of long-term spares management strategies that will ensure spares availability at an optimal level for local sites and reduce the cost of spares A day in the life As a Senior OR Scientist in DST you will be focused on leading the design and development of innovative approaches and solutions by leading technical work supporting RME’s Predictive Maintenance (PdM) and Spare Parts (SP) programs. You will connect with world leaders in your field and you will be tackling customer's natural language challenges by carrying out a systematic review of existing solutions. The appropriate choice of methods and their deployment into effective tools will be the key for the success in this role. About the team Our mission is to leverage the use of data, science, and technology to improve the efficiency of RME maintenance activities, reduce costs, increase safety and promote sustainability while creating frictionless customer experiences. We are open to hiring candidates to work out of one of the following locations: Luxembourg, LUX
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists in the Forecasting, Macroeconomics & Finance field document, interpret and forecast Amazon business dynamics. This track is well suited for economists adept at combining cutting edge times-series statistical methods with strong economic analysis and intuition. This track could be a good fit for candidates with research experience in: macroeconometrics and/or empirical macroeconomics; international macroeconomics; time-series econometrics; forecasting; financial econometrics and/or empirical finance; and the use of micro and panel data to improve and validate traditional aggregate models. Economists at Amazon are expected to work directly with our senior management and scientists from other fields on key business problems faced across Amazon, including retail, cloud computing, third party merchants, search, Kindle, streaming video, and operations. The Forecasting, Macroeconomics & Finance field utilizes methods at the frontier of economics to develop formal models to understand the past and the present, predict the future, and identify relevant risks and opportunities. For example, we analyze the internal and external drivers of growth and profitability and how these drivers interact with the customer experience in the short, medium and long-term. We build econometric models of dynamic systems, using our world class data tools, formalizing problems using rigorous science to solve business issues and further delight customers. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists in the Forecasting, Macroeconomics & Finance field document, interpret and forecast Amazon business dynamics. This track is well suited for economists adept at combining cutting edge times-series statistical methods with strong economic analysis and intuition. This track could be a good fit for candidates with research experience in: macroeconometrics and/or empirical macroeconomics; international macroeconomics; time-series econometrics; forecasting; financial econometrics and/or empirical finance; and the use of micro and panel data to improve and validate traditional aggregate models. Economists at Amazon are expected to work directly with our senior management and scientists from other fields on key business problems faced across Amazon, including retail, cloud computing, third party merchants, search, Kindle, streaming video, and operations. The Forecasting, Macroeconomics & Finance field utilizes methods at the frontier of economics to develop formal models to understand the past and the present, predict the future, and identify relevant risks and opportunities. For example, we analyze the internal and external drivers of growth and profitability and how these drivers interact with the customer experience in the short, medium and long-term. We build econometric models of dynamic systems, using our world class data tools, formalizing problems using rigorous science to solve business issues and further delight customers. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Economists in the Forecasting, Macroeconomics & Finance field document, interpret and forecast Amazon business dynamics. This track is well suited for economists adept at combining cutting edge times-series statistical methods with strong economic analysis and intuition. This track could be a good fit for candidates with research experience in: macroeconometrics and/or empirical macroeconomics; international macroeconomics; time-series econometrics; forecasting; financial econometrics and/or empirical finance; and the use of micro and panel data to improve and validate traditional aggregate models. Economists at Amazon are expected to work directly with our senior management and scientists from other fields on key business problems faced across Amazon, including retail, cloud computing, third party merchants, search, Kindle, streaming video, and operations. The Forecasting, Macroeconomics & Finance field utilizes methods at the frontier of economics to develop formal models to understand the past and the present, predict the future, and identify relevant risks and opportunities. For example, we analyze the internal and external drivers of growth and profitability and how these drivers interact with the customer experience in the short, medium and long-term. We build econometric models of dynamic systems, using our world class data tools, formalizing problems using rigorous science to solve business issues and further delight customers. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Bellevue, WA, USA | Boston, MA, USA | Los Angeles, CA, USA | New York, NY, USA | San Francisco, CA, USA | Seattle, WA, USA | Sunnyvale, CA, USA