Efficient End-to-End Visual Document Understanding with Rationale Distillation (2024)

Abstract: Understanding visually situated language requires interpreting complex layouts of textual and visual elements. Pre-processing tools, such as optical character recognition (OCR), can map document image inputs to textual tokens, then large language models (LLMs) can reason over text.However, such methods have high computational and engineering complexity. Can small pretrained image-to-text models accurately understand visual documents through similar recognition and reasoning steps instead?We propose Rationale Distillation (RD), which incorporates the outputs of OCR tools, LLMs, and larger multimodal models as intermediate “rationales”, and trains a small student model to predict both rationales and answers. On three visual document understanding benchmarks representing infographics, scanned documents, and figures, our Pix2Struct (282M parameters) student model finetuned with RD outperforms the base model by 4-5% absolute accuracy with only 1% higher computational cost.

## introduction information in the digital world is conveyed through text integrated with visual elements, such as complex layouts, figures, and illustrations. answering user questions based on such visual documents requires models to recognize and connect text and layout to the user need. while pretrained image-to-text multimodal models have demonstrated strong performance on visual document understanding (vdu) by directly mapping pixel-level input document images to answers corresponding to user queries @xcite @xcite , stateof-the-art approaches benefit from the use of external tools. tools include ocr systems @xcite @xcite ., * work done during internship at google deepmind small vl model learn rationales from tools (only in training) for vdu on transcribing summarization deplotting programming posters scanned documents charts and plots 2022), structured table source extraction @xcite , and llms reasoning over extracted information and the user query @xcite . additional tools such as image captioning, object classification, and search engines have been used for other multimodal tasks @xcite . however, the accuracy gains from these external components come at the cost of decreased computational efficiency and increased engineering complexity. in this work, we ask whether we can achieve high accuracy and efficiency by teaching a smaller model to learn from short rationales generated by external tools and expensive llms (see figure 1 ). we use a small student image-to-text model to perform vdu tasks by decomposing them into rationale prediction and answer steps, predicting the rationale and answer in sequence. the "rationale" can be any intermediate textual information that helps answer a question correctly: for instance, it could be a subset of relevant text from the image as well as layout, structured information, and rea-ocr tool full ocr: social media platform cheat sheet, pinterest, key differentiating trait : more of a niche platform , despite very … instagram … you need to have lots of high-quality images and videos to support this platform. ## task definition in vdu, a model is given an image i and user question q, and predicts text answer a. we focus on training a single small image-to-text model with parameters θ for this task. prior work in vdu trains such models by maximizing the training data log-likelihood according to an image-to-figure 3 : we first crop along the longer edge of the image to create multiple smaller square images. we generate rationales using the appropriate subset of tools (ocr, llm-summarizer, llm-programmer, plot-to-table) on these images, then categorize the examples and rationales with multimodal-verifier (pali-x). we assume that external tools such as ocr systems, llms, larger image-to-text models, or structured input image source information may be available at training time, but not inference time. we use such tools and metadata to derive rationales r paired with training input-output examples (i, q, a), and train a small student image-to-text model to predict rationales r as an intermediate reasoning step, before predicting the answer a. ## rationale distillation we propose rationale distillation (rd), which distills rationales from a predefined set of tools, and trains a student model to predict the relevant rationales before predicting the answer. rationales are sequences of text tokens of relevant information to arrive at the answer. we consider two kinds of rationales r: natural language text evidence derived from the output of an ocr system; and tabular representation of charts in the input image concatenated with simple custom programs with predefined operations. the tools we leverage are: an ocr tool (google cloud ocr); plot-to-table, a converter that converts charts or plots to structured tables; a llm-summarizer (designed by us), which summarizes ocr text to evidences relevant to the question using a prompted palm 2-l model @xcite ; a llm-programmer (also designed by us and based on palm 2-l), which generates simple programs for numerical reasoning tasks; and a multimodal-verifier based on pali-x (chen et al., 2023b), which verifies the quality of the rationales. we provide detailed descriptions of these tools in appendix a. as the tools add heavy computation (for llms) or engineering complexity (for ocr), we depend on none of them at inference time. in this section, we first discuss the process of generating the two types of rationales from tools ( §3.1). we then describe a data augmentation scheme for increasing the number of examples with rationales, and making student models more robust to potentially noisy rationales ( §3.2). finally, we discuss training and inference for student models to predict the rationale and the answer ( §3.3, 3.4). ## experimental results we study the impact of rationale distillation across three benchmarks, analyze the contribution of each component of our approach, and the extent to which a single student model can match the capabilities of the external tools and llms it learns from. ## related work using tools to augment the input in a prediction problem can be seen as using additional reasoning steps, i.e., calling a tool with a set of arguments and integrating its result with the rest of the context. much prior work on vdu has relied on calling ocr @xcite @xcite , object detector @xcite , or de-plotting tools @xcite . such works have not attempted to recognize text or structured data as an intermediate reasoning step using the same small model. on the other hand, the specific structure of reasoning chains through prompting llms has been shown to have significant impact @xcite @xcite . distilling these text rationales from large teacher models has been shown successful by chain-of-thought distillation works on nlp benchmarks @xcite and scienceqa @xcite . toolformer @xcite trains smaller language models to call tools. generic multimodal tool use solutions based on llms have also been proposed @xcite . however, these works do not replicate the results of tool output and re-place them for efficiency. we marry the powerful ideas of taking intermediate reasoning steps from tools for accuracy, and distilling to small student models for efficiency, as we have proposed in rd. 3 ## conclusions we showed that the visual document understanding ability of small image-to-text models can be improved by our proposed rationale distillation. in rd, we obtain rationales for training examples using external tools and llms, and train small end-to-end student models to predict rationales as intermediate reasoning steps. we demonstrated the importance of designing student training tasks that make the model robust to irrelevant rationales. rd leads to substantial improvements via textual evidence distillation on the text-heavy infovqa & docvqa datasets, and via plot-to-table and program distillation on the numerical reasoningfocused chartqa dataset. analysis shows the gains transfer to stronger models such as matcha (appendix b) and larger pix2struct models. marginalizing over rationales and using a cheap calculator tool at inference time bring additional consistent benefits. controlled experiments show that rd offers a tradeoff between performance and computational cost/engineering complexity, in comparison to systems relying on tool pipelines.