AI Engineer World's Fair 2024
Training Albatross: An Expert Finance LLM
About this talk
Gradient chief scientist Leo Pekelis explains how the Albatross finance-domain language model adapts a Llama 2 foundation using domain-specific training data, human review, and synthetic augmentation. He contrasts long-context, in-context learning with brittle retrieval pipelines for reducing hallucinations and reports improvements on finance-specific evaluations and NVIDIA's RULER long-context benchmark.
Chapters
- 0:00Introduction: Gradient, open models, and expanding context windows
- 3:42Why financial expertise requires domain-specific training data
- 5:48Llama foundations, data screening, and synthetic augmentation
- 7:53Long context, hallucinations, and limitations of RAG
- 14:23Finance benchmarks, NVIDIA RULER, and closing remarks
Talk transcript
- 0:00
[on-hold music] Hi, everyone. Uh, I'm Leo.
- 0:15
I'm the chief scientist at Gradient. And, uh, today I'll be talking about how we trained, uh, large language models to be finance experts. Um, yeah, let's go ahead and dive right into it.
- 0:29
Uh, so, so before kind of I, I start getting into the, the details here, I wanted to make a couple of observations. And the, the first one is that, uh, foundational models have been growing at an exponential rate.
- 0:42
Uh, right? So not only do kind of bespoke AI companies each have their own foundational models, but data companies, uh, general tech companies, uh, they all have their own flavor of a language model, each with its own features, uh, and use cases.
- 0:59
And, uh, another observation, which is, which is pretty related, is that the context length, right, the, the number of tokens that, that you can fit into a prompt, uh, has increased quite a bit over the past year.
- 1:11
Um, the, the largest context length models about a year ago were something like a hundred K. Um, and, uh, in the past year, they've grown to about forty times that, uh, just in models released in the past few months, including one released by Gradient.
- 1:26
Um, and, and both of these, uh, observations are evidence to kind of one, one point, and that's that large language models are not one size fits all. Um, especially when you get to kind of more complicated use cases, uh, taking a, a generalist language model, uh, or, or a base language model kind of off the shelf is,
- 1:46
isn't really gonna get you too far. Um, and, and I realize I'm, I'm talking at the open models track of a conference. I probably don't need to convince you guys, uh, too much of this statement.
- 1:57
But, um, it is pretty important for us at Gradient, and it was actually our, our foundational thesis for, for what we built, uh, which is an AI foundry. And, uh, for us, what an AI foundry is, is it's a collection of custom language models, uh, as well as a number of workflow primitives.
- 2:14
And what we do is we take all these pieces and components together, uh, to create solutions that are a custom fit, uh, for our customers. And, and today I'm gonna talk about specifically, uh, our solutions for the finance domain, right?
- 2:28
Building financial experts. Um, and, and for those solutions, really, uh, two components have been incredibly useful. Uh, one, uh, should be fairly, fairly straightforward is our, uh, domain-specific finance language model, and the other one is a context length, uh, extension that we've worked on.
- 2:49
Um, and, and so why are these important specifically for finance? Well, uh, a little while ago, we, we got together and wrote down kind of six requirements for finance applications of language models that, that generalist models, uh, tend, tend to lack or fall a bit short on.
- 3:03
Um, you know, i-if you look at these, uh, requirements, they're, they're fairly general. They kind of apply across, uh, industries, but in particular for finance, they seem pretty important.
- 3:13
Um, and today I'm just gonna talk about two of them, uh, that happen to be paired, uh, with, uh, the two solutions, uh, that I also wanna talk about, the finance language model and, and the extended context length.
- 3:28
So, um, jumping, uh, jumping right into it, uh, the first one is the, uh, finance language model. Um, you know, may-- you might be wondering, uh, why, why even have a domain-specific language model?
- 3:42
Uh, why is domain knowledge important? Uh, the, the reason is, is that your, your general purpose language models like, like the GPTs of the world, um, they are, uh, trained on a, on a very broad set of data, uh, kind of broad and not deep, uh, especially in, in kind of like more technical situations like technical financial
- 4:01
information. Um, and as kind of like an illustrative example on why this is important, uh, here's a chart from a recent research paper, and it shows that even for very large models, right, the, the red line at the top there is for a hundred and seventy-six billion parameter model.
- 4:18
Um, you need something on the order of thousands of relevant documents in the model's pre-training, uh, in order for the model to get decent... I mean, here it's even above fifty percent accuracy, uh, on answering a related question, right?
- 4:31
And so, uh, kind of what this implies is that if you ask, uh, a language model questions, uh, about, uh, data that's kind of like in the tails of its training data, um, then it-it's gonna do a poor job at answering those questions, right?
- 4:45
And so, um, you know, the, the natural way to fix this is to ca-- is to say, okay, uh, base model doesn't know a lot about finance. Let's train it some finance.
- 4:56
Um, an issue there, um, and, and here I'm gonna talk about kind of how we trained our, our finance-specific language model, is, uh, so an issue there is that there's, there's a whole lot of financial data out there, right?
- 5:09
Like way more, uh, than you could possibly, uh, review or look at manually. Um, and so that requires creating an automated data pipeline. Um, and we, that's what we did.
- 5:19
We created one. Uh, probably the, the most, uh, compelling or interesting part of this data pipeline is the automated data curation, uh, where we borrowed ideas from, uh, the membership inference literature.
- 5:34
And so what we do is we amass a, a whole large corpus of, of training data. Um, and then, uh, we use techniques to, to try to see if a particular document, uh, if there's a high chance that it was already in the model's training data, right?
- 5:48
So maybe you have like a Llama base model, you have a document, uh, and you can run some of these techniques to, to get a probability of whether or not the model has already seen that data in training.
- 5:57
Uh, so you filter out all the data that the model hasn't seen before.
- 6:00
Uh, what you're left with is a much smaller set of data now that's, uh, manageable to, to look at through human review. Uh, and then finally pass through to a synthetic data augmentation, right?
- 6:11
Both to up-sample data, uh, and to handle, uh, some variations in, in data representation and formatting.
- 6:20
Um, and, and kind of like the, the last part of, of the recipe for, for how to train, um, domain-specific, uh, language models, uh, is to take that data set that you created and, and to pass it through a training pipeline.
- 6:35
Um, uh, I think by now a training pipeline like this is, is fairly standard. Uh, there, there's two main parts. One is the continuous pre-training. So you take that data set, uh, that you created on, on the previous slide, uh, and you do kind of next token prediction, uh, on it, uh, o- off of a, a, a
- 6:54
base existing model, right? So again, uh, we're taking a, a base foundational model like a Llama model, uh, to start with. And then the second part, uh, is you do, um...
- 7:03
i- is you run alignment on the model. Uh, here we ran both supervised fine-tuning a- and preference optimization. Um, and, and kind of the way I like to think about, uh, the division between, uh, these two tasks is pre-training i- is something like if you had a bunch of textbooks and you wanted, uh, a model to read
- 7:22
all those textbooks and, and understand all, all that information or retain all that information. And alignment i- is kind of like, uh, then instructing the model on how to use that information or, or best practices and what to do with that.
- 7:35
Um, and so if pre-training is like reading textbooks, alignment is like maybe like taking an exam, uh, on a class or working on a project.
- 7:44
Right. And, um, that's really all I wanted to say about the domain-specific language model. Um, now I want to talk about the... See how much time I have. Great.
- 7:53
Um, about the other part, which is the, uh, extended context and, uh, how extended context or long context language models, uh, help us address hallucinations.
- 8:04
Right. Uh, to give a quick refresher, what are hallucinations? Well, it's a pretty broad term, uh, and it's used quite frequently nowadays. It's, it's whenever, uh, you run inference on a model wh- when you give it a query and it generates content that is irrelevant or, or made up or inconsistent with the input data.
- 8:21
Um, 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, some causes might be just the training data is outdated, right?
- 8:36
You're asking the model a question on information that is now updated since the training data. Uh, another one is, uh, a lot of the training data, uh, practices require automated data collection, and if there's ever inconsistencies or bugs in that, uh, data collection, um, you can get source reference divergence, right?
- 8:56
So the model's just trained on data that doesn't quite make sense. Uh, and, and there's a few other reasons. All of these, uh, can, uh, encode information in the model's memory banks that there isn't quite accurate, uh, and, and that'll cause the model to hallucinate.
- 9:11
And while, uh, alignment or, or, uh, continued training o- of the model can alleviate hallucinations, um, at Gradient we find that actually in-context learning, so, uh, working directly on the prompt during the execution pipeline, uh, is the most direct and, and sample efficient way to reduce hallucinations, right?
- 9:32
Because, uh, what you can do is you can put in a, a relatively small amount of information directly into the prompt, uh, kind of at inference time, uh, and sort of, uh, plaster over or Band-Aid over, uh, issues, uh, with, with the model's training data.
- 9:47
Um, and so that's great. In-context learning works really well. Um, the issue is it works so well that once you start doing it, you wanna do more and more of it, and, and then kind of you run into the...
- 9:58
one of, one of the biggest pain points, uh, with this practice or one of the biggest bottlenecks, uh, which is the context length. Um, and I'm guessing that this is an issue that, that many of you in this room have, have come across yourselves.
- 10:12
Um, and that's, uh, you, you just run out of prompt, uh, in, in terms of for in-context learning. Um, a few examples, uh, for why that can be an issue.
- 10:22
Uh, if you're trying to put in a few short examples into the prompt, you run out of prompt space before you run out of examples, so now you have to spend a lot of time in choosing the particular example or, or working on some kind of like lossy summarization technique.
- 10:35
Um, for more complex produ-- problems, they may require some brittle, uh, pre-processing pipelines. Each can have errors. Um, and also if you do some kind of external memory management such as RAG, uh, those systems tend to have poor performance when the chunks that get pulled, uh, require them to be interrelated, right?
- 10:56
So if you pull one chunk and another chunk that you need to pull, uh, has to reference a previous chunk to, to know if it needs to get, uh, queried, right?
- 11:05
And RAG does, uh, typically does a pretty poor job at that.
- 11:09
Um, right. So context length, uh, is the bottleneck for this. So the most natural thing to do is just extend the context length. Um, and, and so that's, that's what we did with, with some of our models.
- 11:22
Uh, and here really I just wanted to talk about a couple of examples of what suddenly becomes possible, uh, when you have a context length that, that's sort of in the realm of, of a million tokens long.
- 11:35
Um, here on, on the left-hand side, uh, is an example showing that you can now actually put thousands of examples directly into the prompt. Uh, and that kind of gets you back into this kind of like domain learning regime that I talked about earlier.
- 11:49
Uh, it's just now it is, uh, on the fly and at inference time, right? So it can be very adaptive to the problem. Um, and, uh, you, you do find that, uh, for a lot of tasks out there, this like thousands of examples mark is actually necessary, uh, to get kind of production-grade accuracy or, or dangerous levels
- 12:07
of, of accuracy for a model. Um, and, and the other example is, uh, with the long context length Um, you can leverage what transformer models are, are natively really good at, which is being able to attend to every single token in the prompt.
- 12:23
Um, and by doing that, you can actually have the model perform, uh, fairly complicated reasoning, uh, implicitly just, just in through, through going through its, uh, layers and attention layers.
- 12:34
Um, and an example that, um, that we kind of, uh, cooked up, uh, i-in-house, uh, was we took, uh, books that were written by Mark Twain, the author. Uh, and first, we scrubbed the books of any kind of identifying information, right?
- 12:49
So, so no mention of the author or anything like that. Uh, and then we gave that into the model, uh, into its prompt, into its context, and asked the model to generate, uh, new stories in the same style.
- 13:00
Uh, and after kind of five books of, of reference prompts, the model was, uh, able to generate stories, um, that convinced, uh, a separate critic model, uh, that those short stories could have been actually written, uh, by that same author, right?
- 13:15
Uh, and, and in pretty actually, like, deep and intricate ways, not just kind of like stylistic similarity or language, uh, but down to theme and characters and setting and, and things like that.
- 13:26
So, uh, kind of the punchline is, is that long context language models give you more, uh, grounded and robust systems, and there's fewer moving parts. Much more is contained in the language model, which, which is the thing that we all care about, um, and, and that in turn reduces hallucinations.
- 13:45
Right. So, um, you know, the-- those are basically the, the two components, um, two solutions of our platform that I wanted to, to describe to you all today. Um, one of the things that, that we believe in pretty strongly at, at Gradient is to have transparent and verifiable benchmarks.
- 14:03
Uh, and also we're pretty passionate in giving back to the open source community because a lot of what we've, uh, built our work on are our open source, uh, models and techniques themselves.
- 14:14
Uh, and so for both of those solutions, we've open sourced models, uh, on our, um, company page at Hugging Face. Um, one of them is the, the v-alpha-tross model.
- 14:23
So that's the result of applying our, uh, finance domain training on a Llama 2 base model. Um, and here the, the benchmarks show that after doing that, uh, it ends up being competitive, uh, and actually better, uh, competitive at kind of open LLM general, uh, benchmarks and better at finance-specific benchmarks, uh, to models in, in the same
- 14:45
class to its peers. And the other model is, um, a one million context length extension of, uh, a Llama 3 base model that we released pretty recently. Um, and with it, uh, we were able to get, uh, a hundred percent needle in a haystack scores, actually, uh, above one million context lengths.
- 15:06
That's the first image. Uh, and also had, uh, pretty substantial performance improvement, uh, over the base model on a RULER long context length benchmark. That's a benchmark put out by NVIDIA.
- 15:18
Um, and that brings this model kind of in the realm of, uh, flagship long-context models, uh, like Gemini 1.5 Pro, GPT-4 and, and Commander Plus.
- 15:31
Right? And so the- these models are open source, publicly available. Invite you all to, to go and check them out.
- 15:37
Um, and I have about a, about a minute left, so, uh, I'll finish off, uh, here. Uh, there's of course lots more to building, uh, an AI financial expert.
- 15:49
These are just two pieces of the puzzle, even though they're two important ones. Uh, and if you guys are interested in finding out more, uh, feel free to check us out on our, on our website or reach out and contact us.
- 15:59
Cool. Thank you. [upbeat music]