The race to build generative AI models with larger context windows is no longer just about scale. AI21 Labs is using Jamba to argue that long-context AI can also become more efficient.
Jamba is a new text-generating and analyzing model from AI21 Labs. It can perform many of the same broad tasks associated with models such as OpenAI’s ChatGPT and Google’s Gemini, while taking a different approach to how it processes long inputs.
Why context windows matter
A context window is the amount of input data a model can consider before it produces output. In text systems, that usually means the words, phrases, instructions, conversation history or documents the model sees before generating more text.
Small context windows create a familiar problem: the model can lose track of information from earlier in a conversation or document. Larger windows reduce that issue because the system can keep more material in view and better follow the flow of the data it receives.
The tradeoff is compute. Models with large context windows tend to demand more processing power and memory. That makes long-context capability useful, but also expensive or difficult to run in many settings.
AI21 Labs’ product lead Or Dagan argues that this cost does not have to be accepted as inevitable. Jamba is the company’s attempt to show that a model can support a large context window while keeping hardware demands within a more practical range.
What Jamba can handle
Jamba can process up to 140,000 tokens. The source describes that as roughly 105,000 words, or 210 pages, which puts the model in the range of a decent-sized novel.
It can do this while running on a single GPU with at least 80GB of memory, such as a high-end Nvidia A100. That detail is central to AI21 Labs’ pitch: the model is not only large-context, but designed to fit on one powerful GPU.
For comparison, Meta’s Llama 2 has a ~4,000-token context window. It also requires less memory, with the source saying it can run on a GPU with ~12GB of memory. That comparison highlights the practical tension in this part of AI development: more context usually means more memory, even when different models make different tradeoffs.
Jamba was trained on a mix of public and proprietary data. It can write in English, French, Spanish and Portuguese.
The architecture behind the claim
At first glance, Jamba might look like another downloadable generative AI model in a crowded field. The source points to other freely available models, including Databricks’ recently released DBRX and Meta’s Llama 2.
What makes Jamba different is the architecture. It combines transformers with state space models, or SSMs.
Transformers are widely used for complex reasoning tasks and power models such as GPT-4 and Google’s Gemini. Their defining feature is an attention mechanism, which lets a model assess how relevant different pieces of input are to one another before generating output.
That attention mechanism is powerful, but it is also part of why large-context models can become compute-intensive. The model has to work across the information in the input and decide what matters for the next generated text.
SSMs take a different path. They combine qualities of older AI model types, including recurrent neural networks and convolutional neural networks, to process long sequences of data more efficiently.
SSMs are not presented as a perfect replacement for transformers. The source notes that they have limitations. But early examples, including an open source model called Mamba from Princeton and Carnegie Mellon researchers, have shown that SSM-based approaches can handle larger inputs than comparable transformer-based models while doing well on language generation tasks.
Jamba uses Mamba as part of its core model. Dagan claims Jamba delivers three times the throughput on long contexts compared with transformer-based models of comparable sizes.
“While there are a few initial academic examples of SSM models, this is the first commercial-grade, production-scale model,” Dagan said in an interview with TechCrunch. “This architecture, in addition to being innovative and interesting for further research by the community, opens up great efficiency and throughput possibilities.”
Why the release is still limited
Jamba has been released under the Apache 2.0 license, an open source license with relatively few usage restrictions. Even so, Dagan describes it as a research release rather than a model intended for commercial use.
The reason is safety and readiness. According to the source, Jamba does not include safeguards to stop it from generating toxic text. It also lacks mitigations aimed at potential bias.
AI21 Labs plans to make a fine-tuned, ostensibly “safer” version available in the coming weeks. Until then, the current release is best understood as a technical demonstration rather than a polished product for deployment.
That distinction matters. A model can be impressive because of its efficiency, context length or architecture, while still being unsuitable for real-world commercial use without further work. Jamba appears to sit in exactly that category.
What Jamba signals for long-context AI
The broader significance of Jamba is not only that it can read and generate across a large amount of text. It is that AI21 Labs is trying to show a different route to long-context capability.
If long-context models keep growing only by becoming more compute-intensive, access will remain constrained by hardware and cost. A model that fits a 140,000-token context window onto a single GPU with at least 80GB of memory suggests a different design target: more context without simply accepting heavier infrastructure as the default.
Dagan’s view is that the architecture is still early and could improve as Mamba receives additional tweaks.
“The added value of this model, both because of its size and its innovative architecture, is that it can be easily fitted onto a single GPU,” he said. “We believe performance will further improve as Mamba gets additional tweaks.”
For now, Jamba’s importance is measured less by immediate commercial availability and more by the question it raises for the AI industry. As context windows expand, efficiency may become just as important as size.