Triển khai Speculative RAG với Transformers

Mô hình ngôn ngữ lớn (LLM) thường chưa đáp ứng tốt với các nhiệm vụ đòi hỏi nhiều tri thức, cần thông tin chính xác và cập nhật. Đây là lúc retrieval-augmented generation (RAG) phát huy tác dụng, kết hợp khả năng sinh của LLM với các cơ sở tri thức bên ngoài để tăng độ chính xác và mức độ liên quan.

Dù hiệu quả, các hệ thống RAG truyền thống gặp khó khăn khi xử lý tài liệu dài và phức tạp, dẫn đến độ trễ tăng và đôi khi kém chính xác. Để khắc phục, khái niệm speculative RAG đã xuất hiện như một giải pháp đầy hứa hẹn. Hãy cùng tìm hiểu thêm.

Speculative RAG là gì?

Speculative retrieval-augmented generation (RAG) là một khung phương pháp mới nhằm cải thiện hiệu năng và độ chính xác của hệ RAG bằng cách chia quy trình sinh thành hai bước riêng biệt:

1. Soạn thảo: Bước đầu sử dụng một mô hình ngôn ngữ nhỏ, chuyên biệt gọi là RAG Drafter. Mô hình này tạo nhiều bản nháp câu trả lời từ các tập con tài liệu được truy xuất. Mỗi tập con được chọn kỹ để đại diện cho các góc nhìn đa dạng, giảm trùng lặp và nâng chất lượng bản nháp.
2. Kiểm chứng: Các bản nháp do RAG Drafter tạo ra được chuyển đến một LM lớn hơn, đa năng, gọi là RAG Verifier. Mô hình này đánh giá các bản nháp, tập trung chọn ra bản chính xác và đáng tin cậy nhất dựa trên lập luận đi kèm. Quá trình kiểm chứng tận dụng khả năng mô hình ngôn ngữ của LM đa năng để đảm bảo đầu ra cuối vừa chính xác vừa phù hợp ngữ cảnh.

Hãy hiểu rõ hơn về speculative RAG qua sơ đồ này:

Nguồn: Wang và cộng sự, 2024

Cùng đi qua sơ đồ này từng bước. Nó cho thấy các cách dùng thông tin bên ngoài để hỗ trợ LLM trả lời câu hỏi, tập trung vào phương pháp mới gọi là Speculative RAG. Hãy xem các phương pháp ở hàng đầu:

• (a) RAG tiêu chuẩn: Cách đơn giản nhất. Tất cả tài liệu liên quan (1, 2... trong sơ đồ) được đưa trực tiếp cho LLM ("Generalist LM") cùng với câu hỏi (Q) để tạo câu trả lời (A).
• (b) RAG tự phản tư: Ở đây, LLM được huấn luyện để phân tích tài liệu và câu trả lời của chính nó, tạo các thẻ đặc biệt ("Need," "Relevant," "Supported") nhằm cải thiện suy luận.
• (c) Corrective RAG: Phương pháp này dùng một mô hình riêng (mô hình "NLI") để kiểm tra xem các tài liệu truy xuất có thực sự liên quan đến câu hỏi hay không trước khi đưa cho LLM.

Giờ hãy xem kỹ (d) Speculative RAG:

1. RAG Drafter chuyên biệt: Thay vì dồn gánh nặng lên LLM lớn, một mô hình "soạn thảo" nhỏ hơn sẽ dẫn dắt. Nó xem câu hỏi (Q) và tài liệu (1, 2...) rồi tạo nhiều bản nháp câu trả lời (α1, α2, α3...). Mỗi bản nháp dựa trên một tập con tài liệu khác nhau.
2. Lập luận: Mô hình soạn thảo cũng cung cấp lời giải thích (β1, β2, β3...) về lý do chọn các tài liệu cụ thể cho từng bản nháp. Hãy coi như nó "trình bày cách làm".
3. Đánh giá: Lúc này LLM chính ("Generalist LM") vào cuộc. Nó xem từng bản nháp và phần giải thích, rồi chọn câu trả lời (A) tốt nhất dựa trên đánh giá đó.

Tựu trung, speculative RAG giống như có một đội chuyên gia: mô hình soạn thảo đề xuất phương án, và LLM đa năng đưa ra quyết định cuối. Quy trình vì thế nhanh hơn và tiềm năng chính xác hơn.

Ưu điểm của Speculative RAG

Speculative RAG mang lại một số lợi thế chính so với hệ RAG truyền thống:

1. Chính xác hơn

Nhờ dùng RAG Drafter chuyên biệt để tạo nhiều bản nháp từ các tập con tài liệu đa dạng, Speculative RAG nâng cao độ chính xác của đầu ra cuối. Cách tiếp cận này cho phép hệ thống cân nhắc nhiều góc nhìn, giảm khả năng tạo ra câu trả lời sai hoặc thiên lệch. Trong thực nghiệm, Speculative RAG cho thấy cải thiện độ chính xác tới 12,97% trên các chuẩn như PubHealth so với hệ RAG thông thường.

2. Giảm độ trễ

Một thách thức lớn của RAG truyền thống là độ trễ tăng do phải xử lý tài liệu dài và phức tạp. Speculative RAG giải quyết bằng cách chuyển bước soạn thảo sang một LM nhỏ, hiệu quả hơn, có thể tạo bản nháp song song. Cách này rút ngắn tổng thời gian xử lý, với cải thiện độ trễ lên đến 51% trong một số trường hợp.

Độ trễ Speculative RAG so với RAG tiêu chuẩn trên các tập dữ liệu khác nhau. Nguồn: Wang và cộng sự, 2024

3. Hiệu quả tài nguyên

Speculative RAG tối ưu sử dụng tài nguyên bằng cách giao nhiệm vụ soạn thảo cho một LM nhỏ, ít tốn tính toán hơn LM đa năng lớn thường dùng trong RAG truyền thống. Sự phân công này không chỉ tăng tốc độ sinh mà còn giảm tải tính toán tổng thể, giúp giải pháp hiệu quả hơn về tài nguyên.

4. Khả năng mở rộng và linh hoạt

Khung Speculative RAG có tính thích ứng cao và có thể áp dụng cho nhiều nhiệm vụ đòi hỏi tri thức mà không cần tinh chỉnh quá nhiều. Khả năng mở rộng này khiến nó phù hợp với nhiều ứng dụng, từ hỏi đáp đến phân tích tài liệu phức tạp.

Cách triển khai Speculative RAG

Để triển khai phiên bản đơn giản của Speculative RAG cho tác vụ hỏi đáp cơ bản, chúng ta sẽ dùng một tập dữ liệu nhỏ và thiết lập đơn giản với hai mô hình—một mô hình nhỏ để soạn thảo và một mô hình lớn hơn để kiểm chứng.

Đơn giản hóa, chúng ta sẽ dùng thư viện Hugging Face Transformers. Chúng tôi sẽ minh họa cách dùng một mô hình ngôn ngữ nhỏ để tạo nhiều bản nháp và một mô hình lớn hơn để kiểm chứng và chọn bản tốt nhất.

Bước 1: Cài đặt các thư viện cần thiết

Trước tiên, hãy cài đặt các thư viện cần thiết. Bạn sẽ cần khởi động lại phiên làm việc sau khi cài thư viện transformers từ HuggingFace. Đảm bảo bạn đã thêm thông tin xác thực HuggingFace dưới dạng khóa bí mật trong Collab hoặc trình soạn thảo khác.

!pip install transformers torch datasetsĐược hỗ trợ bởi

Bước 2: Nạp tập dữ liệu

Trong ví dụ này, chúng ta sẽ dùng SQuAD (Stanford Question Answering Dataset) có sẵn trong thư viện datasets của Hugging Face.

from datasets import load_dataset
# Load the SQuAD dataset
dataset = load_dataset("squad", split="train[:100]")  # Using a small subsetĐược hỗ trợ bởi

Bước 3: Khởi tạo mô hình

Chúng ta sẽ dùng hai mô hình tiền huấn luyện từ thư viện Hugging Face. Một mô hình nhỏ (distilbert-base-uncased-distilled-squad) sẽ đóng vai RAG Drafter và một mô hình lớn hơn (bert-large-uncased-whole-word-masking-finetuned-squad) sẽ đóng vai RAG Verifier.

Sau đó, chúng ta sẽ xây dựng pipeline soạn thảo và kiểm chứng với các mô hình đã chọn.

from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

# Initialize the smaller model (RAG Drafter)
drafter_model_name = "distilbert-base-uncased-distilled-squad"
drafter_model = AutoModelForQuestionAnswering.from_pretrained(drafter_model_name)
drafter_tokenizer = AutoTokenizer.from_pretrained(drafter_model_name)

# Initialize the larger model (RAG Verifier)
verifier_model_name = "bert-large-uncased-whole-word-masking-finetuned-squad"
verifier_model = AutoModelForQuestionAnswering.from_pretrained(verifier_model_name)
verifier_tokenizer = AutoTokenizer.from_pretrained(verifier_model_name)

# Set up pipelines
drafter_pipeline = pipeline("question-answering", model=drafter_model, tokenizer=drafter_tokenizer)
verifier_pipeline = pipeline("question-answering", model=verifier_model, tokenizer=verifier_tokenizer)Được hỗ trợ bởi

Bước 4: Định nghĩa hàm soạn thảo

Tiếp theo, chúng ta sẽ định nghĩa một hàm để tạo nhiều bản nháp bằng mô hình nhỏ với pipeline Drafter đã định.

def generate_drafts(question, context, num_drafts=3):
   drafts = []
   for _ in range(num_drafts):
       draft = Drafter_pipeline(question=question, context=context)
       drafts.append(draft)
   return draftsĐược hỗ trợ bởi

Bước 5: Định nghĩa hàm kiểm chứng

Giờ chúng ta sẽ định nghĩa hàm verify_drafts() dùng mô hình lớn hơn để kiểm chứng và chọn bản nháp tốt nhất dựa trên điểm tin cậy.

Nó hoạt động bằng cách ánh xạ vị trí ký tự bắt đầu và kết thúc của mỗi bản nháp sang vị trí token tương ứng trong ngữ cảnh đã token hóa thông qua offset mapping. Mô hình kiểm chứng sau đó chấm điểm mỗi bản nháp dựa trên mức độ tin cậy rằng các vị trí này là câu trả lời đúng. Bản nháp có điểm tin cậy cao nhất được chọn làm câu trả lời tốt nhất, đảm bảo đầu ra cuối chính xác và đáng tin cậy.

Việc xử lý cẩn thận căn chỉnh token và chấm điểm trong hàm này là chìa khóa cho hiệu năng được cải thiện của Speculative RAG.

def verify_drafts(question, context, drafts):
    best_draft = None
    highest_score = 0

    # Tokenize the context using the verifier's tokenizer, keeping track of offsets
    inputs = verifier_tokenizer(question, context, return_tensors="pt", return_offsets_mapping=True)
    offset_mapping = inputs['offset_mapping'][0]  # This will give us the character-to-token mapping
    input_ids = inputs['input_ids'][0]

    for draft in drafts:
        start_char = draft['start']
        end_char = draft['end']

        # Find the corresponding token positions using offset mapping
        start_index = None
        end_index = None

        for idx, (start, end) in enumerate(offset_mapping):
            if start_index is None and start_char >= start and start_char < end:
                start_index = idx
            if end_index is None and end_char > start and end_char <= end:
                end_index = idx
            if start_index is not None and end_index is not None:
                break

        # Ensure indices were found and are within bounds
        if (start_index is None or end_index is None or
                start_index >= len(input_ids) or end_index >= len(input_ids)):
            print(f"Draft skipped: Out of bounds or no matching tokens. "
                  f"Start Index: {start_index}, End Index: {end_index}")
            continue

        # Get the confidence score using the larger model
        outputs = verifier_model(input_ids=input_ids.unsqueeze(0))
        score = (outputs.start_logits[0, start_index].item() +
                 outputs.end_logits[0, end_index].item())

        if score > highest_score:
            highest_score = score
            best_draft = draft

    if best_draft is None:
        print("No valid draft found after verification.")

    return best_draftĐược hỗ trợ bởi

Bước 6: Chạy và đánh giá quy trình Speculative RAG

Bây giờ, hãy áp dụng quy trình Speculative RAG của chúng ta cho 10 câu hỏi mẫu từ tập dữ liệu.

correct = 0
total = 10  # Evaluate on 10 samples for simplicity

for i in range(total):
    sample = dataset[i]
    question = sample['question']
    context = sample['context']
    drafts = generate_drafts(question, context)
    best_answer = verify_drafts(question, context, drafts)

    print(f"Q: {question}")

    if best_answer is not None:
        print(f"A: {best_answer['answer']}\n")
        # For simplicity, compare with the first answer (gold) provided in the dataset
        if best_answer['answer'].lower() in sample['answers']['text'][0].lower():
            correct += 1
    else:
        print("No valid draft found.\n")

accuracy = correct / total * 100
print(f"Accuracy: {accuracy}%")Được hỗ trợ bởi

Đoạn mã này đánh giá độ chính xác của phần triển khai speculative RAG giản lược trên 10 mẫu từ một tập dữ liệu. Với mỗi mẫu, nó tạo nhiều bản nháp câu trả lời dựa trên câu hỏi và ngữ cảnh cho trước, sau đó chọn bản tốt nhất bằng hàm kiểm chứng. Nếu tìm được câu trả lời hợp lệ, nó so sánh với câu trả lời chuẩn (gold) trong tập dữ liệu. Nếu khớp, biến correct được tăng lên.

Cuối cùng, mã tính và in ra độ chính xác dưới dạng phần trăm số câu trả lời đúng trên tổng số mẫu. Mã này cho độ chính xác 90% trên tập dữ liệu giản lược này.

Thách thức của Speculative RAG

Dù speculative RAG cho kết quả tốt, nó vẫn tồn tại một số thách thức:

• Độ phức tạp khi triển khai:
• Cần phối hợp nhiều mô hình.
• Tăng độ phức tạp khi tích hợp vào hệ thống hiện có.
• Chi phí huấn luyện:
• Huấn luyện RAG Drafter là một bước bổ sung.
• Cần tài nguyên và thời gian.
• Làm phức tạp việc triển khai, đặc biệt trong môi trường yêu cầu triển khai nhanh.
• Phụ thuộc vào chất lượng truy xuất:
• Hiệu quả phụ thuộc mạnh vào chất lượng tài liệu truy xuất.
• Truy xuất kém dẫn đến bản nháp chưa tối ưu.

Ứng dụng của Speculative RAG

Speculative RAG có thể áp dụng rộng rãi trong các lĩnh vực cần truy xuất thông tin chính xác và hiệu quả. Một số ứng dụng tiềm năng:

• Hỏi đáp đòi hỏi nhiều tri thức:
• Speculative RAG hiệu quả trong các nhiệm vụ đề cao độ chính xác và mức độ liên quan.
• Khung này có thể tạo câu trả lời chính xác và có căn cứ tốt cho các truy vấn phức tạp bằng cách tận dụng các tập con tài liệu đa dạng.
• Truy xuất thông tin thời gian thực:
• Speculative RAG có thể nâng tốc độ và độ tin cậy của phản hồi trong các tình huống cần thông tin cập nhật.
• Khả năng giảm độ trễ mà vẫn giữ độ chính xác khiến nó lý tưởng cho ứng dụng thời gian thực.
• Phân tích văn bản y khoa và pháp lý:
• Speculative RAG hữu ích trong các lĩnh vực cần diễn giải chính xác tài liệu phức tạp.
• Khung này đảm bảo quyết định quan trọng dựa trên thông tin đáng tin cậy bằng cách tạo nhiều bản nháp và chọn bản chính xác nhất dựa trên lập luận.
• Công cụ giáo dục:
• Speculative RAG có thể hỗ trợ tạo nội dung chính xác, giàu ngữ cảnh cho các nền tảng e-learning.
• Điều này đảm bảo người học nhận được thông tin đáng tin cậy và toàn diện, nâng cao trải nghiệm học tập.

Kết luận

Speculative RAG nâng cấp retrieval-augmented generation với độ chính xác cao hơn, giảm độ trễ, hiệu quả tài nguyên và khả năng mở rộng. Chúng tôi đã minh họa cách triển khai bằng Hugging Face Transformers.

Dù còn thách thức, tiềm năng của speculative RAG trong nhiều lĩnh vực khiến nó trở nên hứa hẹn cho nghiên cứu tiếp theo.

Nếu bạn muốn tìm hiểu thêm về triển khai RAG, tôi khuyến nghị các hướng dẫn sau:

• Triển khai Corrective RAG (CRAG) với LangGraph
• RankGPT làm tác nhân xếp hạng lại cho RAG (Hướng dẫn)
• RAG với Llama 3.1 8B, Ollama và Langchain