X-LoRA

Mixture of LoRA Experts (X-LoRA) là một phương pháp PEFT cho phép kết hợp các chuyên gia LoRA theo dạng thưa hoặc đặc dựa trên ma trận mở rộng có độ hạt mịn cao (token, layer, sequence). Phương pháp này tận dụng các adapter LoRA đóng băng và một mô hình cơ sở đóng băng để giảm đáng kể số lượng tham số cần tinh chỉnh.

Một điểm đặc biệt của X-LoRA là tính linh hoạt của nó: nó có thể được áp dụng cho bất kỳ mô hình cơ sở nào của transformers có bộ thích ứng LoRA. Điều này có nghĩa là, bất chấp chiến lược hỗn hợp chuyên gia, không cần phải thực hiện bất kỳ thay đổi nào đối với mã mô hình.

Biểu đồ bên dưới minh họa cách các hệ số mở rộng thay đổi theo từng token cho các prompt khác nhau. Điều này làm nổi bật sự kích hoạt của các adapter khác nhau khi quá trình sinh tiến triển và chuỗi tạo ra ngữ cảnh mới.

Với mỗi bước, X-LoRA yêu cầu chạy mô hình cơ sở hai lần: đầu tiên, để lấy các trạng thái ẩn (hidden states) mà không có bất kỳ bộ thích ứng (adapter) LoRA nào, và thứ hai, các trạng thái ẩn này được sử dụng để tính toán các hệ số tỷ lệ (scalings) được áp dụng cho các bộ thích ứng LoRA, sau đó mô hình được chạy lần thứ hai. Đầu ra của lần chạy thứ hai chính là kết quả của bước mô hình.

Cuối cùng, X-LoRA cho phép mô hình phản ánh lại kiến thức của nó nhờ vào cơ chế hai lần truyền trực tiếp (dual forward pass), và cấu hình lại kiến trúc một cách động.

Tóm tắt từ bài báo là:

Chúng tôi báo cáo một chiến lược hỗn hợp chuyên gia để tạo ra các mô hình ngôn ngữ lớn được tinh chỉnh bằng cách tiếp cận theo cấp độ token từng lớp sâu dựa trên thích ứng hạng thấp (LoRA). Bắt đầu từ một bộ adapter LoRA được huấn luyện trước, chiến lược cổng của chúng tôi sử dụng các trạng thái ẩn để trộn động các lớp đã thích ứng, cho phép mô hình X-LoRA kết quả khai thác các khả năng khác nhau và tạo ra các tổ hợp theo cấp độ lớp sâu chưa từng được sử dụng trước đó để giải quyết các tác vụ. Thiết kế này lấy cảm hứng từ các nguyên tắc sinh học về tính phổ quát và tính đa dạng, trong đó các khối xây dựng của mạng nơ-ron được tái sử dụng trong các biểu hiện phân cấp khác nhau. Do đó, mô hình X-LoRA có thể dễ dàng triển khai cho bất kỳ mô hình ngôn ngữ lớn (LLM) hiện có nào mà không cần sửa đổi cấu trúc nền tảng. Chúng tôi phát triển một mô hình X-LoRA được thiết kế riêng, cung cấp các khả năng khoa học bao gồm các tác vụ phân tích thuận nghịch và khả năng suy luận được tăng cường, tập trung vào phân tích vật liệu sinh học, cơ học và thiết kế protein. Tác động của công việc này bao gồm việc tiếp cận các mô hình có thể mở rộng và thích ứng dễ dàng với kiến thức chuyên ngành mạnh mẽ và khả năng tích hợp xuyên suốt các lĩnh vực kiến thức. Với các chuyên gia trong sinh học, toán học, suy luận, vật liệu lấy cảm hứng từ sinh học, cơ học và vật liệu, hóa học, sinh lý học protein, cơ học và các tính chất phân tử dựa trên cơ học lượng tử, chúng tôi tiến hành một loạt các nghiên cứu tình huống tập trung vào vật lý. Chúng tôi xem xét khả năng ghi nhớ kiến thức, các tác vụ cơ học protein thuận nghịch, thiết kế protein, mô hình hóa tác nhân đối kháng bao gồm xây dựng đồ thị kiến thức luận lý, cũng như thiết kế phân tử. Mô hình không chỉ có khả năng đưa ra các dự đoán định lượng về các tính chất nano cơ học của protein hoặc các tính chất phân tử cơ học lượng tử, mà còn suy luận về các kết quả và dự đoán chính xác các cơ chế có khả năng giải thích các hành vi phân tử khác biệt..

Vui lòng trích dẫn X-LoRA như sau:

@article{10.1063/5.0203126,
    author = {Buehler, Eric L. and Buehler, Markus J.},
    title = "{X-LoRA: Mixture of low-rank adapter experts, a flexible framework for large language models with applications in protein mechanics and molecular design}",
    journal = {APL Machine Learning},
    volume = {2},
    number = {2},
    pages = {026119},
    year = {2024},
    month = {05},
    abstract = "{We report a mixture of expert strategy to create fine-tuned large language models using a deep layer-wise token-level approach based on low-rank adaptation (LoRA). Starting with a set of pre-trained LoRA adapters, our gating strategy uses the hidden states to dynamically mix adapted layers, allowing the resulting X-LoRA model to draw upon different capabilities and create never-before-used deep layer-wise combinations to solve tasks. The design is inspired by the biological principles of universality and diversity, where neural network building blocks are reused in different hierarchical manifestations. Hence, the X-LoRA model can be easily implemented for any existing large language model without a need for modifications of the underlying structure. We develop a tailored X-LoRA model that offers scientific capabilities, including forward/inverse analysis tasks and enhanced reasoning capability, focused on biomaterial analysis, protein mechanics, and design. The impact of this work includes access to readily expandable and adaptable models with strong domain knowledge and the capability to integrate across areas of knowledge. Featuring experts in biology, mathematics, reasoning, bio-inspired materials, mechanics and materials, chemistry, protein biophysics, mechanics, and quantum-mechanics based molecular properties, we conduct a series of physics-focused case studies. We examine knowledge recall, protein mechanics forward/inverse tasks, protein design, adversarial agentic modeling including ontological knowledge graph construction, and molecular design. The model is capable not only of making quantitative predictions of nanomechanical properties of proteins or quantum mechanical molecular properties but also reasoning over the results and correctly predicting likely mechanisms that explain distinct molecular behaviors.}",
    issn = {2770-9019},
    doi = {10.1063/5.0203126},
    url = {https://doi.org/10.1063/5.0203126},
    eprint = {https://pubs.aip.org/aip/aml/article-pdf/doi/10.1063/5.0203126/19964043/026119\_1\_5.0203126.pdf},
}

API

XLoraConfig

XLoraModel

>>> import torch
>>> from transformers import AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
>>> from peft import XLoraConfig, get_peft_model, prepare_model_for_kbit_training

>>> model_config = AutoConfig.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
>>> config = XLoraConfig(
...     task_type="CAUSAL_LM",
...     hidden_size=model_config.hidden_size,
...     xlora_depth=4,
...     adapters={
...         "adapter_1": "./path/to/the/checkpoint/",
...         "adapter_2": "./path/to/the/checkpoint/",
...         "adapter_n": "./path/to/the/checkpoint/",
...     },
... )
>>> int8_config = BitsAndBytesConfig(load_in_8bit=True)
>>> model = AutoModelForCausalLM.from_pretrained(
...     "mistralai/Mistral-7B-Instruct-v0.1",
...     trust_remote_code=True,
...     attn_implementation="flash_attention_2",
...     device_map="cuda:0",
...     torch_dtype=torch.bfloat16,
...     quantization_config=int8_config,
... )
>>> model = prepare_model_for_kbit_training(model)
>>> xlora_model = get_peft_model(model, config)