LongRAG, XLand-100B, nabla^2DFT, DeepSeek-Coder-V2 & Long Code Arena

Week 3, June 2024

Greetings,

Welcome to the latest edition of the State of AI. This time, we explore the cutting-edge advancements and intriguing applications that are shaping the future of artificial intelligence. We kick things off with LongRAG, a novel approach to enhancing retrieval-augmented generation using long-context language models. Next, we delve into the expansive XLand-100B dataset, designed to boost multi-task learning in reinforcement learning environments.

In the world of quantum chemistry, we present nabla^2DFT, a universal dataset that will revolutionize the way we approach neural network potentials for drug-like molecules. For those interested in code intelligence, we have DeepSeek-Coder-V2, which is breaking barriers to offer new possibilities in closed-source model environments. Finally, we introduce the Long Code Arena, a comprehensive set of benchmarks that promises to elevate the performance of long-context code models.

Each of these topics promises to provide a deep dive into the diverse and ever-evolving landscape of AI, ensuring an enriching and engaging read. Enjoy!

Best regards,

State of AI

---

Contents

1. LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs

2. XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning

3. nabla^2DFT: A Universal Quantum Chemistry Dataset of Drug-Like Molecules and a Benchmark for Neural Network Potentials

4. DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence

5. Long Code Arena: a Set of Benchmarks for Long-Context Code Models

---

LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs

Authors: Ziyan Jiang, Xueguang Ma, Wenhu Chen

Source and references: https://arxiv.org/abs/2406.15319

---

Introduction

In the ever-evolving landscape of natural language processing and machine learning, the quest to build better models is relentless. Retrieval-Augmented Generation (RAG) methods have long stood as a robust mechanism to supercharge the capabilities of large language models (LLMs). But what if we told you there's a fresh perspective that positively disrupts the existing balance between retrieval unit complexity and reader capacity? Enter LongRAG – a novel framework that is set to redefine how we think about RAG systems. Let's dive into the intriguing journey of how this new architecture works, the challenges it tackles, and the significant performance gains it brings to the table.

The Problem with Traditional RAG

RAG frameworks have traditionally relied on short retrieval units, typically around 100-word paragraphs from sources like Wikipedia. Imagine combing through millions of such tiny fragments to find the precise one that answers your query. This puts immense pressure on the retriever while the reader’s task remains relatively light, simply extracting answers from the snippets passed on. This "heavy retriever, light reader" imbalance could lead to suboptimal performance. Moreover, short units can cause semantic disjunctions, leading to incomplete information retrieval. With this backdrop set, we need something more balanced and efficient.

Introducing LongRAG

The team from the University of Waterloo presents LongRAG, a paradigm shift in the RAG framework. By extending the length of the retrieval units to 4K tokens – a whopping 30 times longer than the traditional 100-word units – the corpus size is dramatically reduced from 22 million to 600,000 units. This reduces the workload of the retriever and allows it to focus more efficiently on fewer, but richer, units of information. The result is a remarkable improvement in answer recall rates, significantly boosting performance metrics like answer recall@1 to 71% on Natural Questions (NQ) and answer recall@2 to 72% on HotpotQA.

The Nuts and Bolts of LongRAG

Long Retrieval Unit

The first pillar of LongRAG is the concept of 'long retrieval units.' By processing entire Wikipedia documents or even aggregating related documents, LongRAG forms units that provide a more comprehensive context. This reduces the sheer number of units the retriever has to scan, from 22 million to 600,000 units. The upside? The retriever can now zero in on rich, meaningful contexts rather than sifting through countless fragments. Moreover, these extended units capture the continuity and completeness of information, which is crucial for the end performance of question-answering tasks.

Long Retriever

To complement the longer retrieval units, LongRAG employs a 'long retriever' designed to work with these 4K-token units. Instead of overwhelming the retriever with the task of pinpointing exact short snippets, LongRAG focuses on identifying broader, coarser information that is still relevant to the query. This shift allows the top 4 to 8 retrieval units to collectively form a dense and informative context, streamlining the retrieval process. And guess what? The unit size guides how many top-k units are considered. More extended units mean fewer top-k units – typically around 4 to 8.

Long Reader

The final gem in LongRAG's crown is the 'long reader.' With queries pulling in retrievals around 30K tokens, LongRAG leverages advanced long-context language models like Gemini or GPT-4o. Without needing any specific training, these readers can handle vast chunks of text, perform intricate reasoning, and extract precise answers. LongRAG's use of powerful long-context LLMs marks a significant departure from traditional models' constraints, where longer contexts were a limitation, not an advantage.

Performance on Natural Questions (NQ)

Natural Questions (NQ) is one of the most challenging datasets designed for end-to-end question answering tasks. From real Google search queries to spans identified in Wikipedia articles, the complexity and real-world application of NQ make it an ideal benchmark for LongRAG. Remarkably, LongRAG achieves an Exact Match (EM) score of 62.7% on the NQ dataset without any fine-tuning – a performance on par with the state-of-the-art models. The eerily efficient retrieval mechanisms of LongRAG have just made one of the hardest tasks look significantly easier.

Tackling HotpotQA

HotpotQA demands an intricate dance between multiple documents to answer multi-hop questions effectively. Traditional RAG methods often struggle with these because they must juggle retrieving and reasoning across several document fragments. Not so with LongRAG. By leveraging long retrieval units that naturally encompass related documents, LongRAG reduces ambiguity and enhances information completeness. This structural advantage is evidenced by an Exact Match (EM) score of 64.3% on the full-wiki HotpotQA – again without any training. The more comprehensive unit size directly translates to more coherent, multi-faceted answers.

Future Roadmap and Challenges

LongRAG’s impact is clear: it's taken fundamental RAG mechanics and supercharged them for the era of long-context LLMs. But as with any technological leap, there's room for improvement. First, there is a need for even stronger long embedding models to better navigate and encapsulate extensive document contexts. Additionally, developing more general methods for formulating long retrieval units beyond hyperlinks will further enhance the versatility and efficacy of the system.

Why This Matters

So, why should you care about LongRAG? In a nutshell, it showcases a transformative step forward in how we balance and leverage a model's retrieval and reading capabilities. As applications for AI expand from simple Q&A to more complex, multi-step reasoning tasks, frameworks like LongRAG could pave the way for more adept and intuitive systems. Whether you're navigating large document repositories, running chatbots, or enhancing search engines, LongRAG offers a blueprint for building more balanced and capable systems.

Final Thoughts

LongRAG represents a paradigm shift in the RAG space, merging the robustness of retrieval with the depth of long-context reading. It's a balanced system designed for efficiency and high performance, and it makes traditional models look like they're playing catch-up. By significantly reducing the burden on the retriever and leveraging advanced long-context LLMs, LongRAG opens up exciting possibilities in the realm of question answering and beyond. Keep an eye on this space; with innovations like LongRAG, the future of NLP looks incredibly promising.

This research brings to light not only the power of long-context processing but also the necessity for evolving traditional frameworks to be more in line with contemporary capabilities. As we march forward into a data-dense future, such innovations will be at the heart of creating smarter, faster, and more intuitive AI systems. Whether you’re a data scientist, an AI enthusiast, or just someone intrigued by the cutting edge of technology, LongRAG offers a tantalizing glimpse into what’s possible when we rethink and re-engineer the tools at our disposal.

---

XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning

Authors: Alexander Nikulin (AIRI), Ilya Zisman (AIRI), Alexey Zemtsov (Tinkoff), Viacheslav Sinii (Tinkoff), Vladislav Kurenkov (AIRI), Sergey Kolesnikov (Tinkoff)

Source and references: https://arxiv.org/abs/2406.08973

---

Introduction

In the fast-evolving world of artificial intelligence, in-context reinforcement learning (RL) is on the rise. However, the lack of challenging benchmarks has held the field back, making it difficult to gauge the true potential of new methods. Enter "XLand-100B," a massive dataset created to push the boundaries of in-context RL. This dataset, inspired by the XLand-MiniGrid environment, is a game-changer, covering nearly 30,000 different tasks with 100 billion transitions and 2.5 billion episodes. The aim? To democratize research in this exciting field and spark the development of even smarter, more adaptable AI.