X-LoRA

Mixture of LoRA Experts (X-LoRA) adalah metode PEFT yang memungkinkan campuran LoRA experts secara sparse atau dense berdasarkan matriks penskalaan dengan granularitas tinggi (token, layer, sequence). Metode ini memanfaatkan adapter LoRA yang dibekukan dan model dasar yang dibekukan untuk secara drastis mengurangi jumlah parameter yang perlu dihaluskan.

Aspek unik dari X-LoRA adalah versatilitasnya: dapat diterapkan pada model dasar transformers mana pun dengan adapter LoRA. Artinya, meskipun menggunakan strategi mixture of experts, tidak perlu dilakukan perubahan pada kode model.

Grafik di bawah ini menunjukkan bagaimana penskalaan berubah untuk prompt yang berbeda pada setiap token. Hal ini menyoroti aktivasi adapter yang berbeda seiring kemajuan generasi dan pembentukan konteks baru oleh urutan.

Untuk setiap langkah, X-LoRA memerlukan model dasar untuk dijalankan dua kali: pertama, untuk mendapatkan hidden states tanpa adapter LoRA apa pun, dan kedua, hidden states tersebut digunakan untuk menghitung scalings yang diterapkan pada adapter LoRA, kemudian model dijalankan untuk kedua kalinya. Hasil dari eksekusi kedua adalah hasil dari langkah model.

Pada akhirnya, X-LoRA memungkinkan model untuk merefleksikan pengetahuannya berkat skema dual forward pass, serta secara dinamis mengonfigurasi ulang arsitekturnya.

Abstrak dari makalah tersebut adalah:

Kami melaporkan strategi campuran pakar untuk membuat model bahasa besar yang telah disempurnakan menggunakan pendekatan token per lapisan dalam yang berbasis pada adaptasi peringkat rendah (LoRA). Dimulai dari seperangkat adapter LoRA yang telah dilatih sebelumnya, strategi gating kami menggunakan keadaan tersembunyi untuk mencampur lapisan yang telah diadaptasi secara dinamis, memungkinkan model X-LoRA yang dihasilkan untuk memanfaatkan berbagai kemampuan dan menciptakan kombinasi lapisan dalam yang belum pernah digunakan sebelumnya untuk menyelesaikan tugas. Desain ini terinspirasi oleh prinsip biologi universalitas dan keberagaman, di mana blok bangunan jaringan saraf digunakan kembali dalam manifestasi hierarkis yang berbeda. Oleh karena itu, model X-LoRA dapat dengan mudah diimplementasikan untuk model bahasa besar (LLM) yang sudah ada tanpa perlu modifikasi pada struktur dasarnya. Kami mengembangkan model X-LoRA yang disesuaikan yang menawarkan kemampuan ilmiah termasuk tugas analisis maju/balik dan kemampuan penalaran yang ditingkatkan, dengan fokus pada analisis biomaterial, mekanika dan desain protein. Dampak dari pekerjaan ini termasuk akses ke model yang mudah diperluas dan diadaptasi dengan pengetahuan domain yang kuat serta kemampuan untuk terintegrasi di berbagai bidang pengetahuan. Dengan menghadirkan pakar di bidang biologi, matematika, penalaran, material bio-terinspirasi, mekanika dan material, kimia, biofisika protein, mekanika, dan sifat molekuler berbasis mekanika kuantum, kami melakukan serangkaian studi kasus yang berfokus pada fisika. Kami memeriksa pengingatan pengetahuan, tugas maju/balik mekanika protein, desain protein, pemodelan agen adversarial termasuk konstruksi graf pengetahuan ontologis, serta desain molekuler. Model ini mampu tidak hanya membuat prediksi kuantitatif tentang sifat nanomekanik protein atau sifat molekuler mekanika kuantum, tetapi juga menalar hasil tersebut dan dengan benar memprediksi mekanisme yang mungkin menjelaskan perilaku molekuler yang berbeda..

Silakan sitasi X-LoRA sebagai:

@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)