Scaling Transformer to 1M tokens and beyond with RMT
arxiv.org
arxiv.org
We already know Large Language Models (LLMs) can learn at runtime (ie, separately to the training process.) This is called "In Context Learning". See [1], [2] for more details. (BTW, when anyone says "LLMs are stochastic parrots" you know they are ignorant of this)
In context learning is wonderful because it means you can "train" a LLM at run time by filling the context with examples. Traditionally this "context window" has been a few thousand tokens, and GPT-4 recently extended that to 32,000 tokens.
That is useful, but if you wanted to say load all of a companies documents and ask questions it doesn't really work because this overflows the context.
But at 2M tokens there's a whole range of applications that become possible.
[1] Language Models are Few-Shot Learners: https://arxiv.org/abs/2005.14165
[2] Language Models Secretly Perform Gradient Descent as Meta-Optimizers: https://arxiv.org/abs/2212.10559
The “learning” in “In context learning” is not the same concept as learning during the network “training”. The former is loosely named and does not have any impact on what’s actually “learned” by the model. The LLM can, by nature, contradict its training data and whatever you/it prompt(s), and incorporates randomness when completing your prompts, hence the term “stochastic”.
Nevertheless, it’s useful to have bigger context window if it works well.
Actually (and very surprisingly!) it is related.
See the "Similarity of Attention Output Updates" and "Similarity of Attention Map" metrics in https://arxiv.org/abs/2212.10559 (linked above) where they show that in-context learning actually changes the attention map in a similar way to traditional fine tuning(!!).
Obviously it can't modify the weights of the model on disk, but after ICL new passes through the model are treated quite similarly to how they would be if it has been fine tuned (again: this is astonishing!)
The empirical studies are fairly weak as well – none of them suggest that there is actually any gradient descent occurring during in-context learning, only that if you fine-tune a model's attention parameters (and no other parameters) for a set of examples, its internal representations and attention patterns are slightly more similar (and it's really not by much!) to those of in-context learning than before fine-tuning. Is it really that surprising? You could most likely take any two techniques that optimize a model towards some target, artificial restrict both of them to only apply to the same set of parameters, use both, and find increased similarity of the two optimized models compared to the non-optimized model.
None of that would be what the paper claims, i.e., a shared mechanism of gradient descent, just a similarity in outcome – and we all already knew that ICL and fine-tuning have similar results!
- Pre-training (for text prediction, using unsupervised learning)
- Fine-tuning (e.g. to follow instructions or to reject certain queries, using supervised and/or reinforcement learning) (optional)
- in-context learning (using "few-shot prompts" as examples)
Now the last two can have similar effects. For example, you might fine-tune a foundation (only pre-trained) model to follow instructions, or you don't and instead just modify your prompt such that it looks like a dialogue between a human and a helpful chat assistant.
But neither can replace the extensive pre-training phase, which is what gives the model all its intelligence.
One other disanalogy between fine-tuning and in-context learning appears to be that the model can't exactly remember the data it was fine-tuned with, while it "knows" exactly everything in its context window. That is its working memory, so to speak.
One may argue that we'd need a separate channel for the instructions, but it would make harder to do anything but zero shot question and answer.
I think it's likely that systems with metadata for tagging sources, tasks separately from the human and ai interaction will eventually emerge. The gptchat api already works somewhat that way, but it isn't following the structure that we'll, probably because it was a later tune and not something built in the model from the beginning.
Training that encoder sounds hard, though.
Slightly unrelated, but has anyone outside of OpenAI got access to this model yet? While I received API access to GPT4 just a day or two after applying for it, I've yet to succeed to get access to the 32k version, nor has anyone I know, and I have not seen it being used by anyone in the wild either.
First you need to get approved for OpenAI access via Azure Cognitive Services: https://customervoice.microsoft.com/Pages/ResponsePage.aspx?...
Then you need to get approved for GPT4: https://customervoice.microsoft.com/Pages/ResponsePage.aspx?...
Once approved for GPT4, you'll have access to the GPT4-32k model.
Here's the full list of models: https://learn.microsoft.com/en-us/azure/cognitive-services/o...
https://platform.openai.com/docs/model-index-for-researchers
This makes it a very powerful tool in general, since it doesn't suffer from mode collapse. It can be used for anything, though one has to prompt it right, since it doesn't follow instructions by default.
This is literally the use case that every large enterprise wants.
Why? The potential cost savings are enormous. A significant percentage of white collar jobs primarily involve performing repetitive tasks and answering frequently asked questions. If you could automate even 20% of this........
The problem is that I don't think they are ready for what this will cause.
Imagine that Amazon managed to do that and now every employee has access to it. What prevents anyone asking "inconvenient" questions? Think something like:
- Does the Echo dot records audio even when the trigger word isn't used? Show me the relevant source code.
- What is the quality standard expected for an Amazon Basic product?
Companies aren't used to total freedom of information. And I don't think is easy to implement any sort of information access control for any system based on LLMs.
The info lives in a database, not the LLM' model weights, you just use the LLM to understand natural langauge questions, find relevant documents (the user has access to) and providing relevant bits from them to the context to allow responding with them (That's how all systems who show sources work afaik, braid, bing, and mdx docs based stuff)
For outsiders? That's just public FAQ's
You're the CEO? You can ask anything.
There will likely be compartmentalized models, and then the master model that the execs get. No reason you can't have 10, 20 500 different models running.
But - note - the label of "stochastic parrots" is not really relevant to the acquisition of notions and the appearance of some consistency checks, but more to the idea that the ideal system performs the checks regularly and deeply - especially in front of new ideas, of which the ones it outputs are especially relevant (they had been in just before, so a new part of the cognitive set, and you expect them rigorously examined). Whatever you utter, you have supposed to have thought about critically with some depth - structurally.
You acquire a notion, you criticize it, you produce new notions, you criticize them, you build on an epistemic corpus born on critical effort.
In fact, it is a gamechanger to see that "critical thinking" is possible with LLMs - but do not forget that we have seen many examples in months of that not actually happening.
Given that large language models perform some kind of implicit gradient descent during in-context learning, it raises the question of whether they are also doing some form of predictive coding. If so, could this provide insights on how to better leverage stochasticity in language models?
I'm not particularly knowledgeable in the area of probabilistic (variational) inference, I realize that attempting to draw connections to this topic might be a bit of a stretch.
[1] The free-energy principle: a unified brain theory: <https://www.fil.ion.ucl.ac.uk/~karl/The%20free-energy%20prin...>
[2] Predictive Coding: Towards a Future of Deep Learning beyond Backpropagation?: https://arxiv.org/abs/2202.09467
Azure Cognitive Services actually indexes company documents and then searches through them for the relevant information to provide context to ChatGPT to workaround the limited tokens.
Basics of Prompt Engineering with Azure OpenAI Service
You really think Timnit Gebru and Margaret Mitchell, both of whom are cited in the first paper you footnoted, are ignorant of in-context learning?
Also their paper ("On the Dangers of Stochastic Parrots: Can Language Models Be Too Big?") doesn't really define what they mean by a stochastic parrot, but they appear to mean mostly that it is non-grounded (section 6):
> Text generated by an LM is not grounded in communicative intent, any model of the world, or any model of the reader’s state of mind. It can’t have been, because the training data never included sharing thoughts with a listener, nor does the machine have the ability to do that.
I think it is becoming increasingly clear that LLMs do in-fact have a "model of the world". Indeed even old techniques like Word2Vec (2014) showed that RNNs could build a model with meaningful relationships.
I don't think this paper really addresses that, although many of their other criticisms remain valid.
??? One of us is confused here (I'm fully willing to admit that it's me), but AFAIK, the big discovery of in-context learning was described in "Language Models are Few-Shot Learners" - published in 2020[1], vs 2021 for the stochastic parrots paper.
Regardless, it's not like the authors of the stochastic parrots paper have disavowed the term. They're still referring to the paper without correction in the statement they published about the proposed pause last month: https://www.dair-institute.org/blog/letter-statement-March20...
I was looking at getting those transcriptions summarized using a similarily wonderful tool but was kind of stumped when I started seeing 'maximum token length exceeded' errors.
I'm hoping to see some easy plug-n-play solutions to the summarization issue that run completely locally.
Fun times ahead.
Thanks for the feedback.
Edit: Won’t the sliding window solution above introduce a sort of bias?
For instance, with a sliding window of 3 units, unit 1 is captured once whereas units 2 and 3 are captured twice and three times respectively.
> "train" a LLM at run time by filling the context with examples.
Which is precisely what a stochastic parrot needs.
It wouldn't have "intentions" in any human manner. It would have to project a response for each token, and then use that token to generate the next token of response. Almost like it has to rethink things through for each syllable.
One might compare it to giving a human a prompt and letting them add one word, and then giving the same prompt with that word added to a second human and letting them add one word, and so on. It could easily end up with a sensible end-state even though there was no continuous intent from start to end, just aligned intent token by token taken from context.
This is the most interesting response I've gotten that demonstrates gpt-4 viewing its own output as it is built up:
>Yes, I am familiar with the story you're referring to. The title of the short story is "The Machine That Won the War." It was written by Isaac Asimov and first published in 1961. The story is a conversation between three men who played major roles in a war against an alien race, and they discuss the role of a machine called Multivac in winning the war.
>However, it seems that I've made an error in my recollection, as the specific detail you mentioned—refusing to work until the engineer says 'please'—is from a different short story, "Sally," also by Isaac Asimov. In "Sally," autonomous cars stop working until a command is given courteously, using the word 'please.'
This was one large reply, not two separate ones.
Do you? Some of them -at least- understand that the output is conditional on the input.
As an example, tell chatGPT that the Queen of England died (which occurred after the data cut off) and then ask it who the head of state of Australia is.
It's able to infer the head of state of Australia is now Charles III (and gives a good explanation of how this is mostly ceremonial.) See https://twitter.com/nlothian/status/1646699207506685953
At some point the word "stochastic" doesn't really capture that behavior in any useful sense.
It implies that during training, it learned the facts that Elizabeth is queen of the UK, and that Charles is its crown prince; but _also_ the logical rule <IF die(monarch) AND alive(heir_to_the_throne) => transform(heir_to_the_throne, monarch) AND transform(monarch, former_monarch)>, or at least something along those lines that allows similarly powerful entailment. And that in addition to the ability to substitute/reify with the input sequence at inference runtime.
Would be nice to see a rigorous survey of its logical capabilities given some complex Prolog/Datalog/etc knowledge-base as baseline.
It absolutely needed to know who the successor would be via training data.
But to know that "The Queen of England died" also means that the head of state of Australia has changed means that it has an internal representation of those relationships.
(Another way of seeing this is with multi-modal models where the visual concepts and word concepts are related enough it can map between the two.)
Not entirely sure what you mean, but ...show me? Why not just share a link instead of making empty assertions?
https://www.quora.com/Once-Queen-Elizabeth-dies-will-Prince-...
There are loads of articles and discussions online speculating about what “will” happen when Queen Elizabeth dies.
When you have a very, very, very large corpus to sample from, it can look a lot like reasoning.
There's clearly an internal representation of the relationships that is being updated.
If you follow my Twitter thread it shows some temporal reasoning capabilities too. Hard to argue that is just copied from training data: https://twitter.com/nlothian/status/1646699218290225154
> Thinking step by step, if King Charles III and his son William abdicate who will be head of state of Australia.
My response (with chatGPT) is:
Firstly, it's important to note that Australia is a constitutional monarchy, which means that the monarch of the United Kingdom is also the monarch of Australia.
If King Charles III were to abdicate, the next in line to the throne is his eldest son, Prince William, who would become the King of the United Kingdom and the Commonwealth Realms, including Australia.
However, if both King Charles III and Prince William were to abdicate, the next in line to the throne is Prince William's eldest child, Prince George. If Prince George is still a minor at the time of the abdication, a regent would be appointed to serve as head of state until he comes of age.
It's worth noting that abdications of this nature are rare and would likely only occur in exceptional circumstances. Additionally, any decision regarding abdication would be made by the monarch themselves and would be subject to the laws and customs of the country in question.
I tested GPT a lot in other domains, what I found that as long the information explicitly exists (connection between facts) then the responses are fine. I assume that if GPT will reach the state where it can infer new facts, we will be flooded with discoveries that require cross domain knowledge. Nothing like that happened yet.
Feels like we're only one paper away now that the context window has absolutely ballooned.
Inapplicable why exactly? Because you say so? Logic isn't magic. Nor is learning. No (external) grounding is required either: iteratively eliminating inconsistent world models is all you need to converge toward a model of the real world. Nothing especially human or inhuman about it. LLM architecture may not be able to represent a fully recursive backtracking truth maintenance system, but it evidently managed to learn a pretty decent approximation anyway.
Chill my friend, no need to get personal. We are talking about ideas. It’s OK to disagree. I am simply dismissing your initial claim. This usually happens when you present a scientific argument based on personal beliefs. If it’s not magic, then we should be able to doubt and examine it and it should eventually pass scientific muster.
> No grounding is required… It evidently managed to learn a pretty decent approximation.
Well, last time I used an LLM it suggested that I should lift the chair I am sitting in. I guess OpenAI has a lot of work to do. They have to eliminate this inconsistent world model for chairs, tables, floor, My dog, my cat and all the cats living on Mars…
edit: added a missing word.
I'm seeing an emergent capability of encoding higher order logic, and the whole point of such abstractions is to not need to hardcode your weights with the minutiae of cats on Mars. LLMs today are only trained to predict text, so it's hardly surprising that they have some gaps in their understanding of Newtonian physics. But that doesn't mean the innate capability of grasping such logic isn't there, waiting for the right training regime to expose it to its own falling apples, so to speak.
Humans continually update their foundational understanding by assimilating vital information from their "context window" and dumping irrelevant noise. If LLMs could emulate this, it would be a huge win.
Overall, very exciting area of research!
Uh, isn't that how it should always be? If the output isn't conditional on the input, then it is basically noise or worthless. How the conditionals are setup are the key but not really sure how does this relate to the point you are making.
People who say "LLMs are stochastic parrots" may well be aware of effects of conditioning the input.
This in itself does not completely describe the capabilities of a LLM, since their ability to learn and use those new facts to override their previous "beliefs" is not what you'd expect from mere stochastic behavior.
My point is that "stochastic parrot" believers are unaware of this.
The practical limit seems to be where one would have to say that humans are also "stochastic parrots", presumably because that defeats the point of making such an argument in the first place.
If an engine outputs things like
"Charles III is the current King of Britain. He was born in January 26, 1763"
, it seems to be conflating different notions without having used the necessary logic required for vetting the statement: the witness has reasons to say "this seems like stochastic parroting" - two memories are joined through an accidental link that slightly increases Bayesian values. That appears without need of handwaving. And a well developed human intellect will instead go "Centuries old?! Mh.", which makes putting engine and human in the same pot a futile claim - "handwaving".Similarly for nice cases to be kept for historic divulgation, such as "You will not fool me: ten kilos of iron and half a kilo of feathers weigh the same".
When on the other hand the engine «solv[es] complicated tasks», what we want is to understand why. And we do that because, on the engineering side, we want reliable; and on the side of science, we want an increase of understanding, then knowledge, then again possible application.
We want to understand what happens in success cases especially because we see failures (evident parroting counts as failure). And it is engineering, so we want reliable - not just the Key Performance Indicator, but the /Critical/ KPI ("pass|fail") is reliability.
So the problem is not that sometimes they do not act like "stochastic parrots" (this is a "problem" in a different sense, theoretical), but that they even only sometimes do act like "stochastic parrots". You do not want an FPU that steers astray.
Normal practice should be to structure automated tests and see what works, what does not, and assess why.
I don't get what you are trying to say. That's like a property of any useful system.
But it will probably be rejected. The quality bar for NeurIPS is quite high.
Some reasons:
The experiments are very weak: There are just a few figures, basically figure 1 and figure 5, which show some results. There are no tables with numbers.
But more importantly: There are no comparisons (in terms of experiments/numbers) to similar models, like:
- Block-Recurrent Transformers (https://arxiv.org/abs/2203.07852) and related approaches to make the Transformer recurrent, so effectively getting infinite context length.
- All the work on sparse attention, or linear attention, like Longformer, etc, which should also allow for such context lengths.
I don't mean that they just mention them in related work (they partly do, although it looks very short, so I'm quite sure they leave out a lot of other related work, without looking too much into it). I mean that they really run experiments and show numbers. And also, to look at the model definitions of those alternatives, and compare and analyze the differences.
Analysis also seems to be a bit weak. Not sure about novelty.
So, while the presented approach could be interesting, it's hard to tell how good it really performs, and how it compares to alternatives.
(This is now a 10 min review by me. If I had to review this, I usually would spend more like 1-2h at least, so giving much more details. But I don't expect that my initial impression would change too much.)
That paper is written very differently.
* FlashAttention: In my experience, the current best solution for n² attention, but it's very hard to scale up beyond the low tens of thousands of tokens. Memory use is O(n) but compute is O(n²). Code: https://github.com/HazyResearch/flash-attention
* Heinsen Routing: In my experience, the current best solution for n×m attention, i.e., mapping n tokens to m tokens. It's like a souped-up version of attention. I've used it to pull up more than a million tokens as context. Memory use and compute are O(nm). It works, but in my (limited) experience, it doesn't work out-of-the-box as well as FlashAttention for n² attention. Code: https://github.com/glassroom/heinsen_routing
* RWKV: A sort-of-recurrent model which claims to have performance comparable to n² attention in transformers. In my (limited) experience, it doesn't. Others seem to agree: https://twitter.com/arankomatsuzaki/status/16390003799784038... . Code: https://github.com/BlinkDL/RWKV-LM
* RMT (this method): I'm skeptical that the recurrent connections will work as well as n² attention or n×m routing in practice, but I'm going to give it a try. Code: https://github.com/booydar/t5-experiments/tree/scaling-repor...
In addition, the group that developed FlashAttention is working on state-space models (SSMs) that look promising to me. The idea is to approximate n² attention dynamically using only O(n log n) compute. There's no code available, but here's a blog post about it: https://hazyresearch.stanford.edu/blog/2023-03-27-long-learn... [CORRECTION: Code is available. See comment by lucidrains below. I'm hopeful this will go to the top of my list.]
If anyone here has other suggestions for working with long sequences (hundreds of thousands to millions of tokens), I'd love to learn about them.
From my tests, SSMs are a very promising line of research and on my (small) tests on S4, it really has better characteristics than transformers, as it learned faster, a larger context and with smaller dataset.
Working on proprietary stuff. Not allowed to share details.
But I'll ask about connecting online :-)
And you need to read to the very and to find out that the contribution is actually: Keep your architecture (but still need to train it on specific tasks)...
It's like a blog-post recipe in research paper form
In the new paper, they plug their memory system into vanilla BERT. This makes the resulting model essentially nothing like Transformer-XL, which was a strictly decoder-only generative language model.
That isn't scaling to 1M tokens - the model should be expected to answer questions about any text within its context window.
do you know of any benchmarks doing this today?
given the acute need to evaluate models on contextual factuality, we're exploring how to create a benchmark for this purpose but prefer existing benchmarks if possible.
openai's truthfulqa[0] is close but does not focus on contextual factuality and targets a much harder problem of absolute truth.
if none exist, and people are interested in contributing, please reach out.
See also the threads here, mentioning other methods for scaling up context length up to the millions of tokens:
I’m going to, hopefully correctly, assume that was a funny joke.
I still don’t understand why this is even wishful thinking, I don’t understand what the point of delaying progress is.
"Progress" doesn't mean "every massive change to the world and humanity is good." There are undeveloped technologies that we are currently not capable of being responsible with.
Google differential technological research.
Yet I have trouble seeing what people find so dangerous. It’s amazingly cool stuff that will create massive productivity gains.
I guess I just lack needed imagination to believe this “most dangerous thing humans have ever done” perspective.
It looks to me the most dangerous thing is actually Gain Of Function research on pathogens. That seems like most dangerous thing. Close 2nd is nuclear weapons.
LLMs seems really nice and fuzzy and warm to me in comparison.
> Intelligence is the only advantage we have over lions, who are otherwise much bigger and stronger and faster than we are. But we have total control over lions, keeping them in zoos [...]
The "intelligence" is a mouthful, but I think the advantage described above is, at its root, caused by a slightly better "predict" step in the observe-predict-act loop. That's it. And LLMs aren't exactly bad predictors, are they? And it looks like while our predictive abilities rise over centuries, theirs rise over decades.
We are so close to something amazing, and scary.
Maybe we do.
I’m not sure how to define it, but new should be a discovery or insight or even relationship that is not explicitly taught in the training data.
If we don’t then it suggests without humans in the loop that the super intelligence is not so close?
Most people forget that there is no memory or thought independent of whatever the output it generates.
But when you ask it to follow a chain of thought and generate it in the output, the eventual conclusions it lands on is scarily human.
People asking ChatGPT for word counts and sighing at how wrong it is. Yeah of course, it has no counter variables to hold numbers and increment them. It has no memory.
But when you ask it to generate a numbered list of all the words in a passage and then output the word count, it gets it right every time. Because you basically gave it a memory by encoding the counter in the generated output.
I have tried this so many times for novel problems and no matter what I do, it eventually recognizes what we're trying to solve and determines it's unsolvable. Since nothing new comes out I can only attribute the times it does seem to do symbolic manipulation as regurgitation.
Probably bugs will happen in the interface between human language and the symbolic reasoning. But that happens with humans too.
https://gist.github.com/int19h/428cea1d87dfc389b99de1b79727f...
What I found really amazing about this particular experiment is that the schema I gave it didn't contain any information that could be used to query for things like distances between places, and yet it came up with the idea of using settlements' culture as a proxy to determine "border fiefs", completely unprompted (and yes, it actually is a very effective proxy for this particular case!).
I wonder now what it would do with Prolog - or maybe Datalog for simplicity? - although that might depend on how much of it there is in the training data. Do you know if anyone tried it yet?
A good example is all those toy examples of "Program a whatever in the style of Shakespeare and David Bowie's love child". This isn't a thing that it has seen in training data.
I saw an example in a Numberphile video where they were able to take the vector of the word "prince", subtract the vector of the word "man", add the vector of the word "woman", and the resulting vector was closest to the word "princess". So in theory there could be a "gender" dimension and a "position of authority" dimension in that vector (or the model might be making other, stranger connections between words that we don't understand).
I think the same thing is happening in your example. The model identifies and produces output that keeps it in the general region in the vector space for the "Shakespearean" and "Bowiean" dimensions while still satisfying other requirements.
It's phenomenally impressive, and may be a stepping stone to models that can come up with new ideas, but I don't think we're there yet.
LLMs seem to be able to capture the idea of Shakespeare and Bowie's styles, and intuit a combination of the two, but when I start asking it questions about what it thinks about the process I don't get the impression of any understanding. It can magic up text from prompts but it doesn't understand what it's doing.
The whole point of the 2017 paper "Attention is all you need" was to show that the recurrent structure previously used (with attention mechanisms) wasn't needed. It's pretty cool if it does turn out that both mechanisms are needed and powerful together. After all, the brain doesn't linearly process the entire book you're reading from page 1 to where you are, for every new word you read, it stashes the important parts into short/long term memory along the way.
Still, passing gradients back through several stages like that is typically a tricky thing to do in practice.
According to Wolfram Alpha that is:
Single-spaced document: 1500 pages
Double-spaced document: 3000 pages
Book: 1028 pages
So around 1-5 books.
I'm assuming they're using OpenAI's tiktoken tokenizer. (??)
I figured that I’d be able to roll my own summarization much like I was able to do so with transcription.
I’m hoping that the findings of this work are open source.
Does anyone have any extra insights?
Fact1: The hallway is east of the bathroom.
Fact2: The bedroom is west of the bathroom.
Question: What is the bathroom east of?
Answer: bedroom
Asked for better understanding, if the two facts given above are embedded in,say the Bible, would the system be able to answer the question correctly? Does that mean it can track the state of an object and determine it at a certain point in time?
This technical report presents the application of a recurrent memory to extend the context length of BERT, one of the most effective Transformer-based models in natural language processing. By leveraging the Recurrent Memory Transformer architecture, we have successfully increased the model's effective context length to an unprecedented two million tokens, while maintaining high memory retrieval accuracy. Our method allows for the storage and processing of both local and global information and enables information flow between segments of the input sequence through the use of recurrence. Our experiments demonstrate the effectiveness of our approach, which holds significant potential to enhance long-term dependency handling in natural language understanding and generation tasks as well as enable large-scale context processing for memory-intensive applications.
Introduction
The Transformer model (Vaswani et al., 2017) has been widely adopted and used in various research areas and industrial applications. The most important issue of the model is quadratic complexity of attention operation, that makes large models increasingly difficult to apply to longer inputs.
This report we show that by using simple token-based memory mechanism introduced in (Bulatov et al., 2022) can be combined with pretrained transformer models like BERT (Devlin et al., 2019) with full attention and full precision operations can be applied to sequences longer than 1 million tokens using a single Nvidia GTX 1080Ti GPU.
Contributions
1. We enhance BERT by incorporating token-based memory storage and segment-level recurrence with recurrent memory (RMT).
2. We demonstrate that the memory-augmented BERT can be trained to tackle tasks on sequences with lengths up to seven times its originally designed input length (512 tokens).
3. We discovered the trained RMT’s capacity to successfully extrapolate to tasks of varying lengths, including those exceeding 1 million tokens with linear scaling of computations required.
4. Through attention pattern analysis, we found the operations RMT employs with memory, enabling its success in handling exceptionally long sequences.
Discussion
The problem of long inputs in Transformers has been extensively researched since the popularization of this architecture. In this work, we demonstrate that applying Transformers to long texts does not necessarily require large amounts of memory. By employing a recurrent approach and memory, the quadratic complexity can be reduced to linear. Furthermore, models trained on sufficiently large inputs can extrapolate their abilities to texts orders of magnitude longer.
Synthetic tasks explored in this study serve as the first milestone for enabling RMT to generalize to tasks with unseen properties, including language modelling. In our future work, we aim to tailor the recurrent memory approach to the most commonly used Transformers to improve their effective context size.