X-LoRA

Mixture of LoRA Experts (X-LoRA) es un método de PEFT que permite una mezcla dispersa o densa de expertos LoRA basada en una matriz de escalado de alta granularidad (token, capa, secuencia). Esto aprovecha adaptadores LoRA congelados y un modelo base congelado para reducir drásticamente el número de parámetros que deben ajustarse.

Un aspecto único de X-LoRA es su versatilidad: puede aplicarse a cualquier modelo base de transformers con adaptadores LoRA. Esto significa que, a pesar de la estrategia de mezcla de expertos, no es necesario realizar cambios en el código del modelo.

El gráfico siguiente muestra cómo cambian las escalas para diferentes indicaciones en cada token. Esto resalta la activación de diferentes adaptadores a medida que avanza la generación y la secuencia crea un nuevo contexto.

En cada paso, X-LoRA requiere ejecutar el modelo base dos veces: primero, para obtener los estados ocultos sin ningún adaptador LoRA, y en segundo lugar, se utilizan esos estados ocultos para calcular los escalados que se aplican a los adaptadores LoRA, y el modelo se ejecuta una segunda vez. La salida de la segunda ejecución es el resultado del paso del modelo.

En última instancia, X-LoRA permite que el modelo reflexione sobre su conocimiento gracias al esquema de doble pasada hacia adelante, y reconfigure dinámicamente la arquitectura.

El resumen del artículo es:

Informamos sobre una estrategia de mezcla de expertos para crear modelos de lenguaje grandes ajustados finamente mediante un enfoque profundo a nivel de capa y de token basado en adaptación de bajo rango (LoRA). Partiendo de un conjunto de adaptadores LoRA preentrenados, nuestra estrategia de enrutamiento utiliza los estados ocultos para mezclar dinámicamente las capas adaptadas, lo que permite que el modelo X-LoRA resultante aproveche diferentes capacidades y cree combinaciones profundas a nivel de capa nunca antes utilizadas para resolver tareas. El diseño se inspira en los principios biológicos de universalidad y diversidad, donde los bloques de construcción de las redes neuronales se reutilizan en diferentes manifestaciones jerárquicas. Por lo tanto, el modelo X-LoRA puede implementarse fácilmente para cualquier modelo de lenguaje grande (LLM) existente sin necesidad de modificar la estructura subyacente. Desarrollamos un modelo X-LoRA a medida que ofrece capacidades científicas, incluidas tareas de análisis directo/inverso y una capacidad de razonamiento mejorada, centradas en el análisis de biomateriales, la mecánica y el diseño de proteínas. El impacto de este trabajo incluye el acceso a modelos fácilmente expandibles y adaptables con un sólido conocimiento de dominio y la capacidad de integrarse en diferentes áreas del conocimiento. Al contar con expertos en biología, matemáticas, razonamiento, materiales bioinspirados, mecánica y materiales, química, biofísica de proteínas, mecánica y propiedades moleculares basadas en mecánica cuántica, realizamos una serie de estudios de caso centrados en la física. Examinamos la recuperación de conocimiento, tareas directas/inversas de mecánica de proteínas, diseño de proteínas, modelado agéntico adversarial que incluye la construcción de grafos de conocimiento ontológicos, así como el diseño molecular. El modelo es capaz no solo de realizar predicciones cuantitativas de las propiedades nanomecánicas de las proteínas o de las propiedades moleculares de la mecánica cuántica, sino también de razonar sobre los resultados y predecir correctamente los mecanismos probables que explican los distintos comportamientos moleculares..

Cite X-LoRA como:

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