X-LoRA

Mixture of LoRA Experts (X-LoRA) est une méthode PEFT permettant un mélange sparse ou dense d'experts LoRA basé sur une matrice de mise à l'échelle à haute granularité (token, couche, séquence). Cette approche exploite des adaptateurs LoRA figés et un modèle de base figé pour réduire considérablement le nombre de paramètres devant être affinés.

Un aspect unique de X-LoRA est sa polyvalence : il peut être appliqué à tout modèle de base transformers doté d'adaptateurs LoRA. Cela signifie que, malgré la stratégie de mélange d'experts, aucune modification du code du modèle n'est nécessaire.

Le graphique ci-dessous illustre comment les échelles évoluent pour différents prompts à chaque jeton. Cela met en évidence l'activation de différents adaptateurs au fur et à mesure que la génération progresse et que la séquence crée un nouveau contexte.

À chaque étape, X-LoRA nécessite d'exécuter le modèle de base deux fois : d'abord, pour obtenir les états cachés sans aucun adaptateur LoRA, et ensuite, les états cachés sont utilisés pour calculer les échelles qui sont appliquées aux adaptateurs LoRA, et le modèle est exécuté une seconde fois. La sortie de la seconde exécution est le résultat de l'étape du modèle.

En fin de compte, X-LoRA permet au modèle de réfléchir à ses connaissances grâce au schéma de double passe avant, et de reconfigurer dynamiquement l'architecture.

L'abstract de l'article est :

Nous présentons une stratégie de mélange d’experts pour créer de grands modèles de langage affinés à l’aide d’une approche profonde, couche par couche et au niveau des jetons, basée sur l’adaptation de faible rang (LoRA). À partir d’un ensemble d’adaptateurs LoRA pré-entraînés, notre stratégie de routage utilise les états cachés pour mélanger dynamiquement les couches adaptées, permettant au modèle X-LoRA résultant de s’appuyer sur différentes capacités et de créer des combinaisons profondes, couche par couche, inédites pour résoudre des tâches. La conception s’inspire des principes biologiques d’universalité et de diversité, où les briques de base des réseaux de neurones sont réutilisées dans différentes manifestations hiérarchiques. Par conséquent, le modèle X-LoRA peut être facilement implémenté pour n’importe quel grand modèle de langage (LLM) existant, sans nécessiter de modifications de la structure sous-jacente. Nous développons un modèle X-LoRA sur mesure offrant des capacités scientifiques, notamment des tâches d’analyse directe/inverse et une capacité de raisonnement renforcée, axées sur l’analyse des biomatériaux, la mécanique et la conception des protéines. L’impact de ce travail inclut l’accès à des modèles facilement extensibles et adaptables, dotés d’une solide connaissance de domaine et de la capacité à s’intégrer à travers différents domaines de connaissance. En mettant en avant des experts en biologie, mathématiques, raisonnement, matériaux bio-inspirés, mécanique et matériaux, chimie, biophysique des protéines, mécanique et propriétés moléculaires basées sur la mécanique quantique, nous menons une série d’études de cas axées sur la physique. Nous examinons la récupération de connaissances, les tâches directes/inverses de mécanique des protéines, la conception de protéines, la modélisation agentique adversariale incluant la construction de graphes de connaissances ontologiques, ainsi que la conception moléculaire. Le modèle est capable non seulement de faire des prédictions quantitatives des propriétés nanomécaniques des protéines ou des propriétés moléculaires quantiques, mais aussi de raisonner sur les résultats et de prédire correctement les mécanismes probables qui expliquent les comportements moléculaires distincts..

Veuillez citer X-LoRA comme suit :

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