Mô hình và nền tảng AI
Jamba: Mô hình ngôn ngữ lai của AI21 Labs kết hợp Transformer-Mamba
Thêm Unite.AI vào các nguồn ưu tiên của bạn trên Google
Mô hình ngôn ngữ đã chứng kiến sự phát triển nhanh chóng, với kiến trúc Transformer dẫn đầu trong lĩnh vực xử lý ngôn ngữ tự nhiên. Tuy nhiên, khi các mô hình tăng quy mô, thách thức trong việc xử lý ngữ cảnh dài, hiệu quả bộ nhớ và tốc độ đã trở nên rõ ràng hơn.
AI21 Labs đã giới thiệu một giải pháp mới với Jamba, một mô hình ngôn ngữ lớn (LLM) tiên tiến kết hợp điểm mạnh của cả kiến trúc Transformer và Mamba trong một khuôn khổ lai. Bài viết này cung cấp chi tiết về kiến trúc, hiệu suất và ứng dụng tiềm năng của Jamba.
Tổng quan về Jamba
Jamba là một mô hình ngôn ngữ lớn lai được phát triển bởi AI21 Labs, tận dụng sự kết hợp của các lớp Transformer và Mamba, tích hợp với một mô-đun Mixture-of-Experts (MoE). Kiến trúc này cho phép Jamba cân bằng giữa sử dụng bộ nhớ, tốc độ và hiệu suất, làm cho nó trở thành một công cụ mạnh mẽ cho nhiều nhiệm vụ xử lý ngôn ngữ. Mô hình được thiết kế để phù hợp trong một GPU 80GB đơn, cung cấp tốc độ cao và dấu chân bộ nhớ nhỏ trong khi duy trì hiệu suất hàng đầu trên nhiều tiêu chuẩn đánh giá.
Kiến trúc của Jamba
Kiến trúc của Jamba là nền tảng của khả năng của nó. Nó được xây dựng trên một thiết kế lai mới, xen kẽ các lớp Transformer với các lớp Mamba, tích hợp các mô-đun MoE để tăng cường khả năng của mô hình mà không làm tăng đáng kể nhu cầu tính toán.
1. Các lớp Transformer
Kiến trúc Transformer đã trở thành tiêu chuẩn cho các mô hình ngôn ngữ hiện đại do khả năng xử lý song song hiệu quả và bắt được các mối quan hệ dài trong văn bản. Tuy nhiên, hiệu suất của nó thường bị giới hạn bởi nhu cầu bộ nhớ và tính toán cao, đặc biệt khi xử lý ngữ cảnh dài. Jamba giải quyết những hạn chế này bằng cách tích hợp các lớp Mamba, mà chúng ta sẽ khám phá tiếp theo.
2. Các lớp Mamba
Mamba là một mô hình trạng thái không gian (SSM) mới được thiết kế để xử lý các mối quan hệ dài trong các chuỗi một cách hiệu quả hơn so với các RNN truyền thống hoặc thậm chí các Transformer. Các lớp Mamba đặc biệt hiệu quả trong việc giảm dấu chân bộ nhớ liên quan đến việc lưu trữ bộ nhớ đệm khóa-giá trị (KV) trong các Transformer. Bằng cách xen kẽ các lớp Mamba với các lớp Transformer, Jamba giảm sử dụng bộ nhớ tổng thể trong khi duy trì hiệu suất cao, đặc biệt trong các nhiệm vụ yêu cầu xử lý ngữ cảnh dài.
3. Mixture-of-Experts (MoE) Modules
Mô-đun MoE trong Jamba giới thiệu một cách tiếp cận linh hoạt để mở rộng khả năng của mô hình. MoE cho phép mô hình tăng số lượng tham số có sẵn mà không tăng tương ứng số tham số hoạt động trong quá trình suy luận. Trong Jamba, MoE được áp dụng cho một số lớp MLP, với cơ chế định tuyến chọn các chuyên gia hàng đầu để kích hoạt cho mỗi token. Sự kích hoạt chọn lọc này cho phép Jamba duy trì hiệu quả cao trong khi xử lý các nhiệm vụ phức tạp.
Hình ảnh dưới đây minh họa chức năng của đầu cảm biến trong mô hình Attention-Mamba lai, một tính năng quan trọng của Jamba. Trong ví dụ này, đầu cảm biến chịu trách nhiệm dự đoán các nhãn như “Tích cực” hoặc “Tiêu cực” trong các nhiệm vụ phân tích cảm xúc. Các từ được đánh dấu cho thấy cách mô hình tập trung mạnh vào các token nhãn từ các ví dụ vài lần, đặc biệt là vào thời điểm quan trọng trước khi dự đoán nhãn cuối cùng. Cơ chế chú ý này đóng vai trò quan trọng trong khả năng của mô hình để thực hiện việc học trong ngữ cảnh, nơi mô hình phải suy luận nhãn phù hợp dựa trên ngữ cảnh và các ví dụ vài lần.
Cải thiện hiệu suất được cung cấp bởi việc tích hợp Mixture-of-Experts (MoE) với kiến trúc Attention-Mamba lai được nhấn mạnh trong Bảng. Bằng cách sử dụng MoE, Jamba tăng khả năng của mình mà không tăng chi phí tính toán. Điều này đặc biệt rõ ràng trong sự tăng cường hiệu suất đáng kể trên nhiều tiêu chuẩn đánh giá như HellaSwag, WinoGrande và Câu hỏi Tự nhiên (NQ). Mô hình với MoE không chỉ đạt được độ chính xác cao hơn (ví dụ: 66,0% trên WinoGrande so với 62,5% không có MoE) mà còn thể hiện log-probabilities cải thiện trên các lĩnh vực khác nhau (ví dụ: -0,534 trên C4).
Các tính năng kiến trúc chính
- Thành phần lớp: Kiến trúc của Jamba bao gồm các khối kết hợp các lớp Mamba và Transformer trong một tỷ lệ nhất định (ví dụ: 1:7, có nghĩa là một lớp Transformer cho mỗi bảy lớp Mamba). Tỷ lệ này được điều chỉnh cho hiệu suất và hiệu quả tối ưu.
- Tích hợp MoE: Các lớp MoE được áp dụng sau mỗi vài lớp, với 16 chuyên gia có sẵn và hai chuyên gia hàng đầu được kích hoạt cho mỗi token. Cấu hình này cho phép Jamba mở rộng hiệu quả trong khi quản lý các thỏa hiệp giữa sử dụng bộ nhớ và hiệu quả tính toán.
- Bình thường hóa và ổn định: Để đảm bảo ổn định trong quá trình đào tạo, Jamba tích hợp RMSNorm trong các lớp Mamba, giúp giảm thiểu các vấn đề như sự tăng vọt lớn trong hoạt động có thể xảy ra ở quy mô.
Hiệu suất và đánh giá của Jamba
Jamba đã được kiểm tra nghiêm ngặt trên nhiều tiêu chuẩn đánh giá, thể hiện hiệu suất cạnh tranh trên toàn bộ. Các phần sau đây nhấn mạnh một số tiêu chuẩn đánh giá chính mà Jamba đã vượt trội, thể hiện điểm mạnh của nó trong cả nhiệm vụ xử lý ngôn ngữ chung và các tình huống ngữ cảnh dài.
1. Các tiêu chuẩn đánh giá NLP chung
Jamba đã được đánh giá trên nhiều tiêu chuẩn đánh giá học thuật, bao gồm:
- HellaSwag (10-shot): Một nhiệm vụ lý luận thông thường nơi Jamba đạt được điểm hiệu suất là 87,1%, vượt qua nhiều mô hình cạnh tranh.
- WinoGrande (5-shot): Một nhiệm vụ lý luận khác nơi Jamba đạt được điểm 82,5%, một lần nữa thể hiện khả năng của nó trong việc xử lý lý luận ngôn ngữ phức tạp.
- ARC-Challenge (25-shot): Jamba thể hiện hiệu suất mạnh với điểm 64,4%, phản ánh khả năng của nó trong việc quản lý các câu hỏi trắc nghiệm khó.
Trong các tiêu chuẩn đánh giá tổng hợp như MMLU (5-shot), Jamba đạt được điểm 67,4%, cho thấy sự mạnh mẽ của nó trên nhiều nhiệm vụ.
2. Đánh giá ngữ cảnh dài
Một trong những tính năng nổi bật của Jamba là khả năng xử lý ngữ cảnh cực kỳ dài. Mô hình hỗ trợ độ dài ngữ cảnh lên đến 256K token, dài nhất trong số các mô hình công khai. Khả năng này đã được kiểm tra bằng cách sử dụng tiêu chuẩn đánh giá Needle-in-a-Haystack, nơi Jamba thể hiện độ chính xác thu hồi vượt trội trên nhiều độ dài ngữ cảnh, bao gồm cả 256K token.
3. Tốc độ và hiệu quả
Kiến trúc lai của Jamba cải thiện đáng kể tốc độ, đặc biệt với các chuỗi dài.
Trong các thử nghiệm so sánh tốc độ (token mỗi giây) trên các mô hình khác nhau, Jamba liên tục vượt trội so với các mô hình khác, đặc biệt trong các kịch bản liên quan đến kích thước批 lớn và ngữ cảnh dài. Ví dụ, với ngữ cảnh 128K token, Jamba đạt được tốc độ gấp ba lần so với Mixtral, một mô hình so sánh.
Sử dụng Jamba: Python
Đối với các nhà phát triển và nhà nghiên cứu muốn thử nghiệm với Jamba, AI21 Labs đã cung cấp mô hình trên các nền tảng như Hugging Face, làm cho nó có sẵn cho nhiều ứng dụng. Đoạn mã sau minh họa cách tải và tạo văn bản bằng Jamba:
<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p> <p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1") tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p> <p>input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]</p> <p>outputs = model.generate(input_ids, max_new_tokens=216)</p> print(tokenizer.batch_decode(outputs))
Đoạn mã này tải mô hình Jamba và bộ mã hóa, tạo văn bản dựa trên một lời nhắc đầu vào và in ra đầu ra được tạo.
Tinh chỉnh Jamba
Jamba được thiết kế như một mô hình cơ sở, có nghĩa là nó có thể được tinh chỉnh cho các nhiệm vụ hoặc ứng dụng cụ thể. Việc tinh chỉnh cho phép người dùng thích ứng mô hình với các miền cụ thể, cải thiện hiệu suất trên các nhiệm vụ chuyên biệt. Ví dụ sau minh họa cách tinh chỉnh Jamba bằng thư viện PEFT:
import torch from datasets import load_dataset from trl import SFTTrainer, SFTConfig from peft import LoraConfig from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments <p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1") model = AutoModelForCausalLM.from_pretrained( "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p> <p>lora_config = LoraConfig(r=8, target_modules=[ "embed_tokens","x_proj", "in_proj", "out_proj", # mamba "gate_proj", "up_proj", "down_proj", # mlp "q_proj", "k_proj", "v_proj" # attention], task_type="CAUSAL_LM", bias="none")</p> <p>dataset = load_dataset("Abirate/english_quotes", split="train") training_args = SFTConfig(output_dir="./results", num_train_epochs=2, per_device_train_batch_size=4, logging_dir='./logs', logging_steps=10, learning_rate=1e-5, dataset_text_field="quote") trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args, peft_config=lora_config, train_dataset=dataset, ) trainer.train()
Đoạn mã này tinh chỉnh Jamba trên một tập dữ liệu về các câu trích dẫn tiếng Anh, điều chỉnh các tham số của mô hình để phù hợp tốt hơn với nhiệm vụ tạo văn bản trong một miền chuyên biệt.
Triển khai và tích hợp
AI21 Labs đã làm cho gia đình Jamba rộng rãi có sẵn thông qua nhiều nền tảng và tùy chọn triển khai:
- Nền tảng đám mây:
- Có sẵn trên các nhà cung cấp đám mây chính bao gồm Google Cloud Vertex AI, Microsoft Azure và NVIDIA NIM (NVDA ).
- Sắp có trên Amazon Bedrock, Databricks Marketplace và Snowflake Cortex.
- Khung phát triển AI:
- Tích hợp với các khung phát triển phổ biến như LangChain và LlamaIndex (sắp ra mắt).
- AI21 Studio:
- Truy cập trực tiếp thông qua nền tảng phát triển của AI21.
- Hugging Face:
- Mô hình có sẵn để tải xuống và thử nghiệm.
- Triển khai tại chỗ:
- Tùy chọn triển khai riêng, tại chỗ cho các tổ chức có nhu cầu bảo mật hoặc tuân thủ cụ thể.
- Giải pháp tùy chỉnh:
- AI21 cung cấp dịch vụ tùy chỉnh mô hình và tinh chỉnh cho khách hàng doanh nghiệp.
Tính năng thân thiện với nhà phát triển
Mô hình Jamba đi kèm với một số khả năng tích hợp sẵn làm cho chúng đặc biệt hấp dẫn đối với các nhà phát triển:
- Gọi hàm: Tích hợp dễ dàng các công cụ và API bên ngoài vào các công việc AI của bạn.
- Đầu ra JSON có cấu trúc: Tạo các cấu trúc dữ liệu sạch, có thể phân tích trực tiếp từ đầu vào ngôn ngữ tự nhiên.
- Đọc đối tượng tài liệu: Xử lý và hiểu các cấu trúc tài liệu phức tạp một cách hiệu quả.
- Tối ưu hóa RAG: Tích hợp các tính năng để tăng cường các đường ống tạo Retrieval-Augmented.
Những tính năng này, kết hợp với cửa sổ ngữ cảnh dài và xử lý hiệu quả của mô hình, làm cho Jamba trở thành một công cụ đa năng cho nhiều kịch bản phát triển.
Các xem xét đạo đức và AI có trách nhiệm
Mặc dù khả năng của Jamba là ấn tượng, nhưng việc sử dụng nó cần được tiếp cận với một tư duy AI có trách nhiệm. AI21 Labs nhấn mạnh một số điểm quan trọng:
- Bản chất mô hình cơ sở: Mô hình Jamba 1.5 là mô hình cơ sở được tiền huấn luyện mà không có sự điều chỉnh hoặc huấn luyện cụ thể.
- Thiếu cơ chế kiểm duyệt tích hợp: Các mô hình không có cơ chế kiểm duyệt nội tại.
- Triển khai cẩn thận: Việc thích ứng và kiểm duyệt bổ sung nên được thực hiện trước khi sử dụng Jamba trong môi trường sản xuất hoặc với người dùng cuối.
- Bảo mật dữ liệu: Khi sử dụng triển khai dựa trên đám mây, hãy chú ý đến việc xử lý và tuân thủ dữ liệu.
- Nhận thức về thiên vị: Giống như tất cả các mô hình ngôn ngữ lớn, Jamba có thể phản ánh các thiên vị có trong dữ liệu đào tạo của nó. Người dùng nên nhận thức được điều này và thực hiện các biện pháp giảm thiểu phù hợp.
Bằng cách giữ những yếu tố này trong tâm trí, các nhà phát triển và tổ chức có thể tận dụng khả năng của Jamba một cách có trách nhiệm và đạo đức.
Một chương mới trong phát triển AI?
Sự ra mắt của gia đình Jamba bởi AI21 Labs đánh dấu một cột mốc quan trọng trong sự tiến hóa của các mô hình ngôn ngữ lớn. Bằng cách kết hợp điểm mạnh của cả Transformer và mô hình trạng thái không gian, tích hợp kỹ thuật Mixture-of-Experts và đẩy ranh giới của độ dài ngữ cảnh và tốc độ xử lý, Jamba mở ra những khả năng mới cho các ứng dụng AI trên nhiều ngành.
Khi cộng đồng AI tiếp tục khám phá và xây dựng dựa trên kiến trúc sáng tạo này, chúng ta có thể mong đợi thấy những tiến bộ hơn nữa trong hiệu quả mô hình, hiểu biết ngữ cảnh dài và triển khai AI thực tế. Gia đình Jamba đại diện không chỉ là một tập hợp mô hình mới, mà còn là một sự thay đổi tiềm năng trong cách chúng ta tiếp cận thiết kế và triển khai các hệ thống AI quy mô lớn.
