Chain-of-Rank: Enhancing Large Language Models for Domain-Specific RAG in Edge Device

TL;DR

The paper addresses the challenge of performing domain-specific retrieval-augmented generation (RAG) on resource-constrained edge devices by mitigating the costly chain-of-thought (CoT) reasoning in small LLMs. It proposes Chain-of-Rank (CoR), which reframes reasoning as ranking the IDs of retrieved contexts rather than performing elaborate multi-step reasoning, and trains a model with a joint objective under LoRA-based PEFT. Empirical results on HotPotQA and Gorilla API demonstrate that CoR achieves state-of-the-art performance, with higher accuracy (EM, F1, AST) and substantially lower reasoning token cost than CoT- or CoN-based baselines. The work suggests that ranking-based reasoning is both more efficient and effective for domain-specific RAG on edge devices and points to potential generalization to domain-agnostic RAG in future research.

Abstract

Retrieval-augmented generation (RAG) with large language models (LLMs) is especially valuable in specialized domains, where precision is critical. To more specialize the LLMs into a target domain, domain-specific RAG has recently been developed by allowing the LLM to access the target domain early via finetuning. The domain-specific RAG makes more sense in resource-constrained environments like edge devices, as they should perform a specific task (e.g. personalization) reliably using only small-scale LLMs. While the domain-specific RAG is well-aligned with edge devices in this respect, it often relies on widely-used reasoning techniques like chain-of-thought (CoT). The reasoning step is useful to understand the given external knowledge, and yet it is computationally expensive and difficult for small-scale LLMs to learn it. Tackling this, we propose the Chain of Rank (CoR) which shifts the focus from intricate lengthy reasoning to simple ranking of the reliability of input external documents. Then, CoR reduces computational complexity while maintaining high accuracy, making it particularly suited for resource-constrained environments. We attain the state-of-the-art (SOTA) results in benchmarks, and analyze its efficacy.

Figures (2)

• Figure 1: Performance for domain-specific RAG on HotPotQA dataset on LLaMA3-8B with LoRA adapter. The marginal effect of CoT (59.2% $\rightarrow$ 60.6%) is because of the generated incorrect reasoning which severely degrades the performance.
• Figure 2: Illustration of the proposed chain-of-rank for domain-specific RAG. CoR streamlines the reasoning step, which is easier to be learned.