X-LoRA

Mixture of LoRA Experts (X-LoRA) é um método PEFT que permite uma mistura esparsa ou densa de especialistas LoRA com base em uma matriz de escalonamento de alta granularidade (token, camada, sequência). Isso aproveita adaptadores LoRA congelados e um modelo base congelado para reduzir drasticamente o número de parâmetros que precisam ser ajustados finamente.

Um aspecto único do X-LoRA é sua versatilidade: ele pode ser aplicado a qualquer modelo base do transformers com adaptadores LoRA. Isso significa que, apesar da estratégia de mistura de especialistas, nenhuma alteração no código do modelo precisa ser feita.

O gráfico abaixo demonstra como as escalas mudam para diferentes prompts em cada token. Isso destaca a ativação de diferentes adaptadores à medida que a geração avança e a sequência cria um novo contexto.

Para cada etapa, o X-LoRA exige que o modelo base seja executado duas vezes: primeiro, para obter os estados ocultos sem nenhum adaptador LoRA, e, em seguida, os estados ocultos são usados para calcular escalonamentos que são aplicados aos adaptadores LoRA, e o modelo é executado uma segunda vez. A saída da segunda execução é o resultado da etapa do modelo.

Em última análise, o X-LoRA permite que o modelo reflita sobre seu conhecimento devido ao esquema de dupla passagem de frente e reconfigure dinamicamente a arquitetura.

O resumo do artigo é:

Relatamos uma estratégia de mistura de especialistas para criar modelos de linguagem de grande escala ajustados finamente, utilizando uma abordagem profunda, camada a camada e token a token, baseada em adaptação de baixo rank (LoRA). Partindo de um conjunto de adaptadores LoRA pré-treinados, nossa estratégia de roteamento usa os estados ocultos para misturar dinamicamente as camadas adaptadas, permitindo que o modelo X-LoRA resultante aproveite diferentes capacidades e crie combinações profundas, camada a camada, nunca antes utilizadas para resolver tarefas. O design é inspirado nos princípios biológicos de universalidade e diversidade, nos quais os blocos de construção de redes neurais são reutilizados em diferentes manifestações hierárquicas. Portanto, o modelo X-LoRA pode ser implementado facilmente para qualquer modelo de linguagem de grande escala (LLM) existente, sem necessidade de modificações na estrutura subjacente. Desenvolvemos um modelo X-LoRA personalizado que oferece capacidades científicas, incluindo tarefas de análise direta/inversa e capacidade de raciocínio aprimorada, com foco em análise de biomateriais, mecânica e design de proteínas. O impacto deste trabalho inclui o acesso a modelos prontamente expansíveis e adaptáveis, com forte conhecimento de domínio e a capacidade de integrar diferentes áreas do conhecimento. Com especialistas em biologia, matemática, raciocínio, materiais bioinspirados, mecânica e materiais, química, biofísica de proteínas, mecânica e propriedades moleculares baseadas em mecânica quântica, conduzimos uma série de estudos de caso focados em física. Examinamos a recuperação de conhecimento, tarefas diretas/inversas de mecânica de proteínas, design de proteínas, modelagem agêntica adversarial, incluindo a construção de grafos de conhecimento ontológicos, bem como o design molecular. O modelo é capaz não apenas de fazer previsões quantitativas das propriedades nanomecânicas de proteínas ou das propriedades moleculares da mecânica quântica, mas também de raciocinar sobre os resultados e prever corretamente os mecanismos prováveis que explicam os distintos comportamentos 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)