Reasoning Over Complex Documents with DocLLM with Armineh Nourbakhsh - #672

0:04All right, everyone. Welcome to another episode of the TwiML AI podcast. I am your host Sam Charrington. Today I'm joined by Armina Norbosh. Armina is an executive director at JPMorgan AI Research. Before we get going, be sure to take a moment to hit that follow button on Apple Podcasts, Spotify, or wherever you're listening to today's show. Armina, welcome to the podcast. Thank you for having me. I'm looking forward to digging into our conversation. We'll be talking about DocLLM, a pre-trained model that understands structured documents. But before we do, I'd love to have you share a little bit about your background and how you came to work in AI.

0:43Sure. So my background was actually in unimodal NLP. So very standard kind of NLP tasks, information extraction and sentiment analysis. And I came to work in NLP kind of by accident, by luck, I would like to say. So my undergraduate degree was in computer science. And after I graduated, I applied for scholarship at Nanyang Technological University in Singapore. And the scholarship program was pretty interesting. It was a master's degree, but it was a master's by research. And it was slightly different from standard programs where you apply, you know what program you're going to get into, potentially what faculty you're going to work with.

1:15You would just submit your resume in a cover letter and it would get circulated among different faculty. And they would pick you if your background matched their brand. And I was sure sending in my resume that if I got in, I would work with someone in the computer science department and I would be working on database optimization or some core metric like that. And then I got very lucky. Professor Chris Koo, who used to teach at the School of Journalism and Communication, picked my resume. And he wanted to work with targeted sentiment analysis of political news articles at the time. He really wanted to bring in NLP practices.

1:45How did you feel about that at the time? It was intriguing, I have to say. I think maybe it was because I'd mentioned in my resume that I was very interested in languages. And I actually used to teach as an ESL teacher. But it was very exciting because this is back in 2006. And sentiment is very much an unsolved problem in NLP. So it was exciting to work on that front. And I stayed after graduation, I worked in the industry for a decade, staying in core unimodal NLP before I made the switch to document AI. And really the first time I faced my first kind of serious multimodal document AI challenge was when I was working at S &P Global, trying to bring AI practices into the credit ratings business.

2:23And of course, part of the credit rating business is to have these highly skilled analysts look at sometimes 10, sometimes hundreds of pages of documentation from each client. And we were trying to automate some of that work for them. And it was very exciting because it was kind of at the frontier of my comfort zone. And it allowed me to flex my skills a little bit. And I guess the rest is history. I've been in Document AI ever since. The types of documents that you're focused on, I'm imagining, are things like annual reports and regulatory filings, perhaps news, but maybe less so news. Tell us a little bit about the way you think about the document AI landscape.

3:02Sure. So the way I would frame it is enterprise documents. That's our area of interest because, of course, documents could be very diverse. I mean, a restaurant menu is a document, right? But of course, as you said, we are very much focused on enterprise documents. I think of things that JP Morgan and S &P Global might be interested in. Again, like filings and annual reports might be different from expense reports and invoices and purchase orders, which are also enterprise documents. Is that fair? That is very fair. On the one end, you have these dense documents, like you said, financial reports and credit agreements and term sheets and so on.

3:43On the other end, you have very sparse form-like documents like invoices and receipts and tax forms, kind of carrying more visual structure as well as semantic structure, for sure. I've done a number of interviews on the space over the years, and I kind of would have thought it was a bit of a solved problem at this point. I know a lot of people think, hey, let's just apply OCR and that's it. But that's not the complete solution because of the structure, as you've mentioned, but apparently not. Tell us a little bit about that. Sure. I'm glad you asked that question because I do see that attitude sometimes from, say, maybe more kind of research-focused or academic communities about certain areas of.AI being quote-unquote solved.

4:26I actually started saying sentiment analysis was very much an unsolved problem in 2006. And you'd be surprised how unsolved it remains to this day. And the same would extend to DocAI, I would say, even maybe more so. So something that, for instance, really prompted us to think about DocA11 and became a main motivation for it was by early 2023, we were working in DocAI, my team and I, at JPMorgan for about three years. And we had created a suite of capabilities that actually had pretty competitive performance. But one challenge that we kept running into was the fact that the state of the art was dominated by encoder only architectures.

5:04So architectures that were following BERT and the BERTology suite of models. And of course, they have a stack of transformer encoders and they would adapt the mask language modeling objective to the multimodal domain. So they would have something called, say, mask visual language modeling, where instead of masking tokens in a sequence, you are masking visual tokens in a snippet of a document. or some other models would take a document and split it into patches. And then they would also take the text that was on the page and they would try to match each token to the corresponding patch. And that's how the model learned.

5:36That's kind of what I remember. Just going back a few years, one of the big innovations was, hey, before we were treating these documents as just text, and now we have this kind of Swiss army knife called a CNN, and we can start to apply it in different ways to our documents. And hey, now all of our problems are solved. but not quite solved. And we've advanced the tools that we're applying to the problem, it sounds like. Exactly. And I mean, to be fair, it is a very robust architecture, right? It has pretty decent performance. I think Layoutland version 3, which follows that architecture by Microsoft, is still a state of the art across many VRDU.

6:13VRDU stands for Visually Rich Document Understanding. Across many VRDU tasks, the issue we had, though, was, you know, you have to also to think about things like robustness and scaling and efficiency, right? And the issue we ran into was we would have a pre-trained model. And every time you want to apply it to a new task, you would have to fine-tune it. And one of the challenges these models have is the scale of the fine-tuning data set is much smaller compared to the scale of the pre-training data set. So the model, because it doesn't have a generative nature, it doesn't end up modeling the fine-tuning data set very well or very robustly.

6:48So you end up not having to fine tune it on different tasks, but also potentially having to fine tune it on every new distribution of data that you get. So if you have one pre-trained model and say 200 downstream tasks, you are creating 200 fine tuned versions of this model and deploying them independently, which was a problem at the time. And we want it to be more efficient. And I think it was in the spring 2023, it was already several months after ChatGPT was released and was met with huge success, that Dongsheng Wang, who's the lead author on the DACA alum paper, started asking this interesting question of, what if we had a foundational model?

7:23It could basically work with an individual interface being the prompting interface that you would only need to, at most, instruction to work with a suite of tasks. You would not need to create several copies of this model. And that's where really the team got started on this problem. So can you talk a little bit about how an LLM has been incorporated into the approach? You know, we think about LLMs, at least the kind of initial crop of LLMs were, you know, largely text-based. Now we're seeing more integration into multimodal scenarios. How are LLMs incorporated into what you're doing with Doc LLM?

8:04Sure. So as you said, the state-of-the-art LLMs at the time were mostly unimodal. We were just seeing some multimodal LLMs come out. And usually they follow the most successful, at least LLMs follow that autoregressive objective, right, of next token prediction. And we wanted to marry that objective with the layout information that we were getting from the document. So we started doing a pretty massive deep dive on the state-of-the-art literature, both on VRDU front and the LLM front. And we realized that there were a few gaps in the way the documents were being represented in both directions.

8:43One was, and this was not something that we discovered that was already being discussed in the literature. If you wanted to take a transformer based model or an LLM, whatever unimodal model you have, and you wanted to add multimodal signal to it, you had a choice of directly working with the image of a document if you're dealing with a scanned document, for instance. and a lot of the models out there were using these very fancy vision encoders as their backbone and analyses were starting to show that these vision encoders were very inefficient they were adding a lot of parameters to the size of the model but they weren't really doing much the model still kept relying on text is kind of kind of experiencing mode collapse so to speak and the reason was that these vision encoders their architecture was adapted directly from open domain visual question answering.

9:31And of course, if you want to do open domain VQA, it's a different ballgame. You have to exploit inductive biases that don't make a lot of sense in the.KI space, right? You have a much wider distribution of shading and lighting and colors and so on. You have to pay attention to contouring and all that. Whereas in documents, you have a pretty fixed color template usually, and the shading is pretty fixed. There is a grid structure that you can exploit in many forms that was not really being exploited. So a lot of the models actually started throwing away the vision aspect altogether. And that's actually what we also decided to do with DocLLM.

10:07But in future versions, we're hoping to look back at the vision encoder, adding a vision encoder as well. But we purely looked at the spatial and the textual content. The second aspect was that these multimodal models, they basically take unimodal information, say text and layout and vision, and they have to kind of fuse them somehow. Right. And what was happening was the text, the layout information was usually always treated as a second class citizen in a way. It was basically not modeled independently from any of the other modalities. It was just simply treated as, say, positional encodings and added to the text modality.

10:43And that ended up not being effective very much because the model didn't have an opportunity to learn kind of a mental map, for lack of a better term, of what layout means on its own. So we wouldn't learn disentangled representations of spatial versus textual information. And what we decided to do was basically allow the model to model each modality separately before fusing them together. So we have self-attention applied on the layout information, self-attention applied on text, and then we bring them together. And this is a very nice kind of architectural modification because it's plug and play.

11:19You could add it to any large language model backbone. And we, in fact, showed that you could do it with two different architectures, the Falcon architecture and the Llama 2 architecture. um the third and probably the biggest um contribution was that we realized that it's not so as i said earlier we really wanted a generative model that you could use um and to have a generative model you need to have a decoder architecture or an encoder decoder architecture the issue with these architectures is that they work very well and very robustly but only at massive scale. So at smaller scale, below a billion parameters, studies have shown that the encoder-only architecture is a lot more robust.

12:03And that's actually, that kind of makes sense when you think about it, because intuitively the tasks that encoder-only architectures are performing are a lot, quote-unquote, easier, right? They have access to the full context. They see the left context and the right context, whereas the next token prediction objective doesn't allow the model to look into the future. So the issue we had was in order to train a decoder architecture at very large scale so that it cannot perform the encoder-only architectures, you would need lots of data. And there is a fundamental issue of data availability in Document AI.

12:41The largest publicly available set of documents out there that is not licensed, not copyrighted, it's available for public consumption under fair use is the IIT-CDIP dataset, which is composed of, I believe, six or seven million documents from the tobacco industry. And this is the outcome of a lawsuit of the federal US government against the tobacco industry. So number one, a very biased dataset. Number two, very, very small collection compared to whatever we have available in the Unimodal literature, say the number of webpages on the World Wide Web or the number of Wikipedia entries. I would have thought that SEC report style documents would, you know, there would be an abundance of those as well.

13:28Does that not reach the size of this tobacco data set you're referring to? A very good question. It might. The issue with the SEC filings is that they do, because they are primarily XML based, their format is a lot more homogenous compared to what you get from the tobacco industry. They're not all that interesting because they are all the same. Yeah, they look very much the same, at least in terms of the color palette and content. So how do you, at medium scale, how do you make the model more robust? And this is something that we investigated and realized that and open AI had actually proposed something.

14:08I think there was a paper that came out in 2022 or one, if I'm not mistaken, around fill in the middle objectives. So this is almost like cheating. You have a decoder model, but you help it function as an encoder model. By not asking it to predict the next token, you take a sequence, you take a span in the middle of that sequence, you mask that span, and then move the true span to the end of the sequence. And given the full input sequence, you ask the model to generate the middle span. So the model still has access to the left and right context, but it's still working in a generative setting, which is a very nice setup.

14:46And there was actually, I think there was a GLM paper that came out that was kind of an advanced version of this fill in the middle of task. So instead of having one continuous span, you now mask multiple tokens in the middle of a sequence, move all of them to the end, shuffle them around so that they're not in order. So make the task a little harder for the model, but then still include some positional information so that the model knows, okay, now I I am filling position number three that was masked when I was reading the sequence. And that's something that we adapted to the doc LLM architecture.

15:20And it worked really well. Something that we also added on top of that was we have the multimodal space, right? We have an image of a page. So you can mask different tokens on the page. But usually, especially in sparse documents like forms, the text forms kind of natural blocks, right? Like an address forms a block, right? a paragraph forms a block, you could actually do the masking at the block level instead of the token level. And that makes it harder for the model because the model has to understand the entire block and what it needs to generate for that block. And remember, when it's seeing the sequence, it has no idea how many tokens are in the block.

15:56So it has to guess what the content is, where it is, what it looks like, how many tokens are there. And it kind of makes the task slightly harder, meaning the model learns a slightly more robust representation. You've made a couple of comments that were interesting. One is that you talked about how you wanted to, it was important to you to have a generative aspect to this model. You've also said, I think that you wanted to kind of move beyond next token prediction. Can you kind of frame out the tasks that you are hoping to achieve with what ultimately became DocL? I'm like, what core problems were you going after that you designed around?

16:40Yeah, there were four core VRD tasks that we were pursuing. There are more tasks in this space, but those were out of scope for us. The four core tasks were key information extraction, visual question answering, document classification, and tabular reasoning. And those are the four that we also evaluated the model on. Of course, there is page segmentation, there is optical character recognition, there's other tasks that you could look into, but we considered those out of scope because there are solutions out there that kind of do address our needs as people who need VRDU models for enterprise documents.

17:14But these were really the four core tasks where the state of the art was failing to give us what we needed. In spite of the fact that this is Doc LLM, you're not asking it to generate documents for you. You're not asking it to answer questions about documents. It's not like a rag style thing or anything like that it's more focused on kind of these core tasks that you mentioned? Yeah. So the base model is a generative model. So you could potentially instruction tune it for other tasks, but at least those were the tasks that we were targeting. And as you said, we do support visual question answering over documents, but not in a rag style manner where you are basically asked to retrieve and cobble together information from one or more documents.

18:00So that is one of the research directions that we are exploring at the moment. One of the tricks you employed was randomizing the position of the tokens that you were asking the model to predict during training, but then you also fed in the position information, which seems like it would kind of undermine the randomization. Like, how does that work? I think it teaches the model to consider the correct context when it's generating a certain token because you have to understand, especially when you put it into two-dimensional space in an image, it has to understand, okay, now I'm generating block number three and this is what was around block number three.

18:41So I have to be careful to generate something that is semantically consistent with what I was seeing in that region of the page. Another concept that comes up in this paper is the idea of graph machine learning or kind of this graphical structure of these documents. Is that part of the Doc LLM work? I know there's adjacent work or work kind of in the same family that is focused on that, but does this graph semantics come into play in Doc LLM? That's a very good question. So Doc LLM does not have a graph component, but graphs actually offer a really good opportunity to look at documents, especially the spatial information, because the spatial information can be encoded directly into the structure of the graph.

19:28And in fact, the Google cloud team has done some really good research on graph-based models. I believe form net is still state of the art in many form-based extraction tasks. We did also explore some graph representations for document understanding. There was a paper that came out at CIR last year named doc graph LM, where we try to explore that option and that direction. I think there's lots of room to explore in graph models, and we are hoping to spend more investment in there. The main challenge with graph architectures is that it's a little difficult to marry the local and the global kind of view of the graph, right?

20:09Usually the majority... Yeah, so we have kind of two families of graph learning approaches, right? The pre-deep learning, it was all about spectral qualities of the graph. You kind of try to learn the global topology of a graph and then create node representations that way. And then since the graph convolutional networks were popularized, it's more about the spatial properties where you have a message passing algorithm that basically is propagating information from one node to its neighbors and so on. And that message passing algorithm has a locality problem because, again, something that happens with a lot of sometimes even regular recurrent networks is over very long sequences, you tend to lose that message, right?

20:51So there hasn't been a really good, at least in the VRD literature, a very good kind of model that could balance these two views. And I think that's why first, I mean, it's not really trivial or simple, but for lack of a better term, for simpler tasks, such as key information extraction, where you just need local context, graphs tend to do really well. But when you go to more complex tasks, such as visual question answering, they're not there yet. Block infilling, is that what you described where you were infilling textual tokens, but at the end of the sequence? Or I guess the block had me thinking about spatial information as opposed to textual information.

21:33So we actually mask both. We mask the textual information and the spatial information. So the model has to predict both. But what we do provide is the ordering, right? So the model knows that this is block number five that I'm now generating. So it makes it a little more kind of spatially aware as well. You mentioned that one of the tasks that you focused on was focused on tabular data. Can you talk a little bit about how the model performs in that domain? I know that in previous conversations, one of the big challenges in dealing with enterprise documents is that you'll have tabular data in the middle of flowing text, and it really causes a lot of problems with the model.

22:17Can you talk a little bit about with DocLLM, kind of what you've seen there and anything that you've done to specifically accommodate those kinds of documents? Yeah, that's a really good question. And that's exactly right. The model almost has to context switch when it sees a table in a way that it doesn't need to in any other type of segment, like a bullet list or a paragraph. And the reason is for tables, you kind of need to pay attention in two directions, right? Whereas everywhere else, you're looking one way. So the reading order kind of changes. And each table has a different reading order.

22:50Some are more columnar, some need to be processed row-wise. So it makes it very challenging. And in fact, if you look at the evaluation section, I don't think tabular reasoning was where we did the best. Doc LLM did best in key information extraction tasks. But in visual question answering, I think we were 50-50 compared to state-of-the-art, including GPT-4. And then I think in tabular reasoning, we didn't come on top. And that could be one of the reasons. And that's something that we are actually looking at. Overall, when we did the instruction tuning exercises, we realized that there were a couple of angles that we could improve on and a couple of angles where we did really well on.

23:27So when more spatial reasoning was required over both text and spatial signal, we usually came up either first or second compared to the state of the art. When there was denser paragraphs that you had to do higher level reasoning, usually GPT-4 would come up top because it has very strong verbal reasoning skills. Now, we are doing a deep dive on the tabular reasoning data sets specifically because we're not sure if we were failing because of spatial information or we were failing because of something like quantitative reasoning that is very common in data sets that have tabular reasoning questions in them, right?

24:12Our model at the moment works best. What's an example that kind of illustrates the different failure modes? Sure. So, for example, consider that you have a financial report with an income statement, right? You could ask an extractive question that says, what was the net income in 2022? Or you could ask an abstractive question that says, what was the change in net income between 2022 and 2023? Right. And that would require spatial reasoning, followed by some quantitative calculations on top of it. And that's what Dr. LLM doesn't have out of the box as of now. So it works better for extractive settings and where more spatial reasoning is required rather than higher level reasoning, either verbal or quantitative.

24:51The other analysis that was kind of interesting, though, was that we did much better when the datasets were mostly composed of enterprise documents, the ones that you actually enumerated in the beginning. Whereas when we were working with data sets that included screenshots of web pages or Wikipedia tables or other kind of web-based content, we didn't do as well. And that could be because we were trained on a data set that was very heavily geared towards enterprise collection. So that could be another angle that we explore further. We couldn't find a really good, strong tabular reasoning document collection that was geared towards enterprise documents.

25:33We do have data sets such as FinQA, TACQA, all these quantitative reasoning data sets that are put together, but they're not focused on the spatial reasoning aspect of it. They're mostly focused on the quantitative reasoning aspect of it. So I guess that's something we could look at for future work. To what degree are you incorporating vision or pixels into this architecture and where? We are not. This.QLM works purely with text and bounding boxes as in spatial information. However, those bounding boxes are coming from somewhere, and that somewhere is an OCR engine. But it's like pre-processing as opposed to part of the model.