Training Albatross An Expert Finance LLM: Leo Pekelis

Hi everyone. I'm Leo, I'm the Chief Scientist at Gradient, and today I'll be talking about how we trained large language models to be finance experts. Um, yeah, let's go ahead and diveright into it. So, so before I kind of, I— I start getting into the details here, I wanted to make a couple of observations.

And the first one is that foundational models have been growing at an exponential rate,right? So not only do kind of bespoke AI companies each have their own foundational models, but data companies, general tech companies, they all have their own flavor of a language model, each with its own features and use cases.

And another observation, which is— which is pretty related, is that the context length,right, the number of tokens that you can fit into a prompt, has increased quite a bit over the past year. The largest context-length models about a year ago were something like 100k, and in the past year they've grown to about 40 times that, just in models released in the past few months, including one released by Gradient.

And both of these observations are evidence to kind of one— one point, and that's that large language models are not one-size-fits-all. Especially when you get to kind of more complicated use cases, taking a generalist language model or a base language model kind of off-the-shelf isn't really going to get you too far.

And I realize I'm talking at the Open Models track of a conference, so I probably don't need to convince you guys too much of this statement. But it is pretty important for us at Gradient, and it was actually our foundational thesis for what we built, which is an AI Foundry.

And for us, what an AI Foundry is, is it's a collection of custom language models as well as a number of workflow primitives. And what we do is we take all these pieces and components together to create solutions that are a custom fit for our customers.

And today I'm going to talk about specifically our solutions for the finance domain,right, building financial experts. And for those solutions, really, two components have been incredibly useful. One should be fairly, fairly straightforward: it is our domain-specific finance language model, and the other one is a context-length extension that we've worked on.

And so why are these important specifically for finance? Well, a little while ago we got together and wrote down kind of six requirements for finance applications of language models that generalist models tend— tend to lack or fall a little bit short on.

You know, if you look at these requirements, they're fairly general, they kind of apply across industries, but in particular for finance they seem pretty important. And today I'm just going to talk about two of them that happen to be paired with the two solutions that I also want to talk about: the finance language model and the extended context length.

So jumping— jumpingright into it, the first one is the finance language model. You know, you might be wondering why— why even have a domain-specific language model? Why is domain knowledge important? The reason is that your general-purpose language models, like the GPTs of the world, they are trained on a very broad set of data, kind of broad and not deep, especially in kind of like more technical situations, like technical financial information.

And as kind of like an illustrative example on why this is important, here's a chart from a recent research paper, and it shows that even for very large models,right, the red line at the top there is for a 176 billion parameter model.

You need something on the order of thousands of relevant documents in the model's pre-training in order for the model to get decent— I mean, here it's even above 50% accuracy on answering a related question,right? And so kind of what this implies is that if you ask a language model questions about data that's kind of like in the tails of its training data, then it's going to do a poor job at answering those questions,right?

And so, you know, the natural way to fix this is to say, okay, base model doesn't know a lot about finance, let's train it some finance.

And the issue there, and here I'm going to talk about kind of how we trained our finance-specific language model, is— so an issue there is that there's a whole lot of financial data out there,right, like way more than you could possibly review or look at manually.

And so that requires creating an automated data pipeline. And we— that's what we did, we created one. Probably the most compelling or interesting part of this data pipeline is the automated data curation, where we borrowed ideas from the membership inference literature.

And so what we do is we amass a whole large corpus of training data, and then we use techniques to try to see if a particular document, if there's a high chance that it was already in the model's training data,right?

So maybe you have like a Llama-based model, you have a document, and you can run some of these techniques to get a probability of whether or not the model has already seen that data in training. So you filter out all the data that the model hasn't seen before, what you're left with is a much smaller set of data now that's manageable to look at through human review, and then finally pass through to synthetic data augmentation,right, both to upsample data and to handle some variations in data representation and formatting.

And kind of like the last part of the recipe for how to train domain-specific language models is to take that data set that you created and to pass it through a training pipeline. I think by now a training pipeline like this is fairly standard.

There's two main parts. One is the continuous pre-training. So you take that data set that you created on the previous slide, and you do kind of next token prediction on it off of a base existing model,right? So again, we're taking a base foundational model like a Llama model to start with.

And then the second part is you do— is you run alignment on the model. Here we ran both supervised fine-tuning and preference optimization. And kind of the way I like to think about the division between these two tasks is pre-training is something like if you had a bunch of textbooks and you wanted a model to read all those textbooks and understand all that information or retain all that information.

And alignment is kind of like then instructing the model on how to use that information or best practices and what to do with that. And so if pre-training is like reading textbooks, alignment is like maybe like taking an exam on a class or working on a project.

Right, and that's really all I wanted to say about the domain-specific language model. Now I want to talk about the— let's see how much time I have, great— about the other part, which is the extended context and how extended context or long-context language models help us address hallucinations.

Right, to give a quick refresher, what are hallucinations? Well, it's a pretty broad term and it's used quite frequently nowadays. It's whenever you run inference on a model, when you give it a query, and it generates content that is irrelevant or made up or inconsistent with the input data.

There's been a fair amount of research as to the cause of hallucinations. A lot of that research points to deficiencies in the underlying training data,right? So some causes might be just the training data is outdated,right? You're asking the model a question on information that is now updated since the training data.

Another one is a lot of the training data practices require automated data collection, and if there's ever inconsistencies or bugs in that data collection, you can get source reference divergence,right? So the model is just trained on data that doesn't quite make sense.

And there's a few other reasons. All of these can encode information in the model's memory banks that isn't quite accurate, and that'll cause the model to hallucinate. And while alignment or continued training of the model can alleviate hallucinations, at Gradient we find that actually in-context learning, so working directly on the prompt during the execution pipeline, is the most direct and sample-efficient way to reduce hallucinations,right?

Because what you can do is you can put in a relatively small amount of information directly into the prompt, kind of at inference time, and sort of plaster over or Band-Aid over issues with the model's training data. And so that's great.

In-context learning works really well. The issue is it works so well that when you start doing it, you want to do more and more of it, and then kind of you run into one of the biggest pain points with this practice, or one of the biggest bottlenecks, which is the context length.

And I'm guessing that this is an issue that many of you in this room have come across yourselves, and that's you just run out of prompt in terms of for in-context learning. A few examples for why that can be an issue: if you're trying to put in a few shot examples into the prompt, you run out of prompt space before you run out of examples, and now you have to spend a lot of time in choosing the particular example or working on some kind of like glossy summarization technique.

For more complex problems, they may require some brittle pre-processing pipelines, each can have errors. And also if you do some kind of external memory management such as RAG, those systems tend to have poor performance when the chunks that get pulled require them to be interrelated,right?

So if you pull one chunk and another chunk that you need to pull has to reference a previous chunk to know if it needs to get queried,right? And RAG does typically does a pretty poor job with that. Right, so context length is the bottleneck for this.

So the most natural thing to do is just extend the context length. And so that's what we did with some of our models. And here really I just wanted to talk about a couple of examples of what suddenly becomes possible when you have a context length that's sort of in the realm of a million tokens long.

Here on the left-hand side is an example showing that you can now actually put thousands of examples directly into the prompt, and that kind of gets you back into this kind of like domain learning regime that I talked about earlier.

It's just now it is on the fly and at inference time,right? So it can be very adaptive to the problem. And you do find that for a lot of tasks out there, this like thousands of examples mark is actually necessary to get kind of production-grade accuracy or dangerous levels of accuracy for a model.

And the other example is with a long-context length, you can leverage what transformer models are natively really good at, which is being able to attend to every single token in the prompt. And by doing that, you can actually have the model perform fairly complicated reasoning implicitly, just in through going through its layers and attention layers.

And an example that we kind of cooked up in-house was we took books that were written by Mark Twain, the author, and first we scrubbed the books of any kind of identifying information,right? So no mention of the author or anything like that.

And then we gave that into the model, into its prompt, into its context, and asked the model to generate new stories in the same style. And after kind of five books of reference prompts, the model was able to generate stories that convinced a separate critic model that those short stories could have been actually written by that same author,right?

And in pretty actually like deep and intricate ways, not just kind of like stylistic similarity or language, but down to theme and characters and setting and things like that. So kind of the punchline is that long-context language models give you more grounded and robust systems, and there's fewer moving parts, much more is contained in the language model, which is the thing that we all care about, and that in turn reduces hallucinations.

Right, so you know those are basically the two components, two solutions of our platform that I wanted to describe to you all today. One of the things that we believe in pretty strongly at Gradient is to have transparent and verifiable benchmarks.

And also we're pretty passionate in giving back to the open-source community because a lot of what we've built our work on are open-source models and techniques themselves. And so for both of those solutions, we've open-source models on our company page at Hugging Face.

One of them is the v-alpha-tross model, so that's the result of applying our finance domain training on a Llama 2-based model. And here the benchmarks show that after doing that, it ends up being competitive and actually better competitive at kind of open LLM general benchmarks and better at finance-specific benchmarks to models in the same class, to its peers.

And the other model is a one million context length extension of a Llama 3-based model that we released pretty recently. And with it, we were able to get 100% needle-in-the-haystack scores actually above one million context lengths, that's the first image, and also had pretty substantial performance improvement over the base model on a ruler long-context length benchmark, that's a benchmark put out by Nvidia.

And that brings this model kind of in the realm of flagship long-context models like Gemini 1.5 Pro, GPT-4, and Command R Plus. Right, and so these models are open-source, publicly available, invite you all to go and check them out.

And I have about a minute left, so I'll finish off here. There's of course lots more to building an AI financial expert. These are just two pieces of the puzzle, even though they're two important ones. And if you guys are interested in finding out more, feel free to check us out on our website or reach out and contact us.

Cool, thank you.