Predicción de la eficiencia del ARN con guía CRISPR-Cas9 con modelos optimizados de manera eficiente en Amazon SageMaker

La tecnología de repetición palindrómica corta agrupada y regularmente interespaciada (CRISPR) promete revolucionar las tecnologías de edición genética, lo que transformará la forma en que entendemos y tratamos las enfermedades. Esta técnica se basa en un mecanismo natural que se encuentra en las bacterias y que permite que una proteína acoplada a una sola cadena de ARN guía (ARNg) se ubique y realice cortes en sitios específicos en el genoma objetivo. Poder predecir computacionalmente la eficiencia y especificidad del ARNg es fundamental para el éxito de la edición genética.

El ARN, que se transcribe a partir de secuencias de ADN, es un tipo importante de secuencia biológica de ribonucleótidos (A, U, G, C) que se pliega en una estructura tridimensional. Gracias a los recientes avances en los modelos de lenguaje de gran tamaño (LLM), se pueden resolver diversas tareas de biología computacional mediante el ajuste fino de los LLM biológicos entrenados previamente con miles de millones de secuencias biológicas conocidas. Las tareas posteriores en los ARN están relativamente poco estudiadas.

En esta publicación, adoptamos un LLM genómico preentrenado para la predicción de la eficiencia de gRNA. La idea es tratar un gRNA diseñado por computadora como una oración y ajustar el LLM para realizar tareas de regresión a nivel de oración análogas al análisis de sentimientos. Usamos métodos de ajuste fino de parámetros eficientes para reducir la cantidad de parámetros y el uso de GPU para esta tarea.

Descripción general de la solución

Los modelos de lenguaje grandes (LLM) han despertado mucho interés por su capacidad para codificar la sintaxis y la semántica de los lenguajes naturales. La arquitectura neuronal que sustenta los LLM son transformadores, que se componen de bloques codificadores-decodificadores basados ​​en la atención que generan una representación interna de los datos a partir de los cuales se entrenan (codificador) y son capaces de generar secuencias en el mismo espacio latente que se asemejan a los datos originales (decodificador). Debido a su éxito en el lenguaje natural, trabajos recientes han explorado el uso de los LLM para la información de biología molecular, que es de naturaleza secuencial.

DNABERT es un modelo de transformador entrenado previamente con datos de secuencias de ADN humano no superpuestos. La estructura básica es una arquitectura BERT compuesta por 12 capas de codificación. Los autores de este modelo informan que DNABERT es capaz de capturar una buena representación de características del genoma humano que permite un rendimiento de vanguardia en tareas posteriores como la predicción de promotores y la identificación de sitios de unión/splice. Decidimos utilizar este modelo como base para nuestros experimentos.

A pesar del éxito y la adopción popular de los LLM, ajustar estos modelos puede resultar difícil debido a la cantidad de parámetros y cálculos necesarios para ello. Por este motivo, se han desarrollado métodos de ajuste fino con eficiencia de parámetros (PEFT). En esta publicación, utilizamos uno de estos métodos, llamado LoRA (adaptación de bajo rango). Presentamos el método en las siguientes secciones.

El siguiente diagrama es una representación del mecanismo de diana del ADN de Cas9. El ARNg es el componente que ayuda a dirigirse al sitio de escisión.

El objetivo de esta solución es ajustar un modelo DNABERT base para predecir la eficiencia de la actividad de diferentes candidatos de gRNA. Por ello, nuestra solución primero toma los datos de gRNA y los procesa, como se describe más adelante en esta publicación. Luego, usamos un notebook de amazon SageMaker y la biblioteca PEFT Hugging Face para ajustar el modelo DNABERT con los datos de ARN procesados. La etiqueta que queremos predecir es la puntuación de eficiencia tal como se calculó en condiciones experimentales de prueba con las secuencias de ARN reales en cultivos celulares. Esas puntuaciones describen un equilibrio entre poder editar el genoma y no dañar el ADN que no era el objetivo.

El siguiente diagrama ilustra el flujo de trabajo de la solución propuesta.

Prerrequisitos

Para esta solución, necesita acceso a lo siguiente:

• Una instancia de notebook de SageMaker (entrenamos el modelo en una instancia ml.g4dn.8xlarge con una sola GPU NVIDIA T4)
• transformadores-4.34.1
• peft-0.5.0
• ADNBERT 6

Conjunto de datos

Para esta publicación, utilizamos los datos de gRNA publicados por investigadores en un artículo sobre Predicción de gRNA mediante aprendizaje profundoEste conjunto de datos contiene puntuaciones de eficiencia calculadas para diferentes ARNm. En esta sección, describimos el proceso que seguimos para crear los conjuntos de datos de entrenamiento y evaluación para esta tarea.

Para entrenar el modelo, necesitas una secuencia de ARNg de 30 meros y una puntuación de eficiencia. Un k-mero es una secuencia contigua de k bases de nucleótidos extraídas de una secuencia de ADN o ARN más larga. Por ejemplo, si tienes la secuencia de ADN “ATCGATCG” y eliges k = 3, entonces los k-meros dentro de esta secuencia serían “ATC”, “TCG”, “CGA”, “GAT” y “ATC”.

Puntuación de eficiencia

Empezar con un archivo de Excel 41467_2021_23576_MOESM4_ESM.xlsx del artículo CRISPRon en la sección Datos suplementarios 1. En este archivo, los autores publicaron las secuencias de ARNg (20-mer) y los correspondientes total_indel_eff puntuaciones. Usamos específicamente los datos de la hoja llamada spCas9_eff_D10+dox. Usamos el total_indel_eff columna como puntuación de eficiencia.

Datos de entrenamiento y validación

Dados los 20-meros y las puntuaciones crispron (igual que el total_indel_eff puntuaciones) de antes, complete los siguientes pasos para reunir los datos de entrenamiento y validación:

1. Convierta las secuencias de la hoja “TRAP12K microarray oligos” en un archivo .fa (fasta).
2. Ejecutar el script get_30mers_from_fa.py (desde Repositorio de GitHub de CRISPRon) para obtener todos los 23-meros y 30-meros posibles de las secuencias obtenidas en el Paso 1.
3. Utilice el CRISPRspec_CRISPRoff_pipeline.py script (del repositorio CRISPRon GitHub) para obtener la energía de enlace para los 23-meros obtenidos en el Paso 2. Para obtener más detalles sobre cómo ejecutar este script, consulte el código publicado por los autores del artículo CRISPRon (consulte el script CRISPRon.sh).
4. En este punto, tenemos 23 meros junto con los puntajes de energía de enlace correspondientes y 20 meros junto con los puntajes CRISPRon correspondientes. Además, tenemos los 30 meros del Paso 2.
5. Utilice el script prepare_train_dev_data.py (de nuestro código publicado) para crear divisiones de entrenamiento y validación. Al ejecutar este script se crearán dos archivos: train.csv y dev.csv.

Los datos se parecen a lo siguiente:

id,rna,crisproff_score,crispron_score
seq2875_p_129,GTCCAGCCACCGAGACCCTGTGTATGGCAC,24.74484099890205,85.96491228
seq2972_p_129,AAAGGCGAAGCAGTATGTTCTAAAAGGAGG,17.216228493196073,94.81132075
. . .
. . .

Arquitectura del modelo para la codificación de gRNA

Para codificar la secuencia de ARNg, utilizamos el codificador DNABERT. DNABERT se entrenó previamente con datos genómicos humanos, por lo que es un buen modelo para codificar secuencias de ARNg. DNABERT convierte la secuencia de nucleótidos en k-meros superpuestos, y cada k-mero sirve como una palabra en el vocabulario del modelo DNABERT. La secuencia de ARNg se divide en una secuencia de k-meros, y luego cada k-mero se reemplaza por una incrustación para el k-mero en la capa de entrada. De lo contrario, la arquitectura de DNABERT es similar a la de BERT. Después de codificar el ARNg, utilizamos la representación del token (CLS) como la codificación final de la secuencia de ARNg. Para predecir la puntuación de eficiencia, utilizamos una capa de regresión adicional. La pérdida de MSE será el objetivo de entrenamiento. El siguiente es un fragmento de código de la DNABertForSequenceClassification modelo:

class DNABertForSequenceClassification(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        
        self.bert = BertModel(config)
        classifier_dropout = (
            config.classifier_dropout
            if config.classifier_dropout is not None
            else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional(torch.Tensor) = None,
        attention_mask: Optional(torch.Tensor) = None,
        token_type_ids: Optional(torch.Tensor) = None,
        position_ids: Optional(torch.Tensor) = None,
        head_mask: Optional(torch.Tensor) = None,
        inputs_embeds: Optional(torch.Tensor) = None,
        labels: Optional(torch.Tensor) = None,
        output_attentions: Optional(bool) = None,
        output_hidden_states: Optional(bool) = None,
        return_dict: Optional(bool) = None,
    ) -> Union(Tuple(torch.Tensor), SequenceClassifierOutput):
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `(0, ...,
            config.num_labels - 1)`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        print('bert outputs', outputs)
        pooled_output = outputs(1)
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (
                    labels.dtype == torch.long or labels.dtype == torch.int
                ):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)
        if not return_dict:
            output = (logits,) + outputs(2:)
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

Afinando y estimulando los LLM genómicos

Ajustar todos los parámetros de un modelo es costoso porque el modelo entrenado previamente se vuelve mucho más grande. LoRA es una técnica innovadora desarrollada para abordar el desafío de ajustar con precisión modelos de lenguaje extremadamente grandes. LoRA ofrece una solución al sugerir que los pesos del modelo entrenado previamente permanezcan fijos mientras se introducen capas entrenables (conocidas como matrices de descomposición de rangos) dentro de cada bloque de transformador. Este enfoque reduce significativamente la cantidad de parámetros que se deben entrenar y disminuye los requisitos de memoria de la GPU, porque la mayoría de los pesos del modelo no requieren cálculos de gradiente.

Por lo tanto, adoptamos LoRA como un método PEFT en el modelo DNABERT. LoRA se implementa en la biblioteca PEFT Hugging Face. Al usar PEFT para entrenar un modelo con LoRA, los hiperparámetros del proceso de adaptación de rango bajo y la forma de encapsular los modelos de transformadores base se pueden definir de la siguiente manera:

from peft import LoraConfig

tokenizer = AutoTokenizer.from_pretrained(
        data_training_args.model_path,
        do_lower_case=False
    )
# DNABertForSequenceClassification is a model class for sequence classification task, which is built on top of the DNABert architecture.    
model = DNABertForSequenceClassification.from_pretrained(
        data_training_args.model_path,
        config=config
    )
    
# Define LoRA Config
LORA_R = 16
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
peft_config = LoraConfig(
                     r=LORA_R, # the dimension of the low-rank matrices
                     lora_alpha=LORA_ALPHA, #scaling factor for the weight matrices
                     lora_dropout=LORA_DROPOUT, #dropout probability of the LoRA layers
                     bias="none",
                     task_type="SEQ_CLS"
    )
model = get_peft_model(model, peft_config)

Rendimiento de la evaluación de retención

Utilizamos RMSE, MSE y MAE como métricas de evaluación y realizamos pruebas con rangos 8 y 16. Además, implementamos un método de ajuste fino simple, que consiste simplemente en agregar varias capas densas después de las incrustaciones DNABERT. La siguiente tabla resume los resultados.

Cuando rank=8, tenemos 296.450 parámetros entrenables, lo que representa aproximadamente el 33 % del total. Las métricas de rendimiento son “rmse”: 11,933, “mse”: 142,397, “mae”: 7,014.

Cuando rank=16, tenemos 591.362 parámetros entrenables, lo que representa aproximadamente el 66 % del total. Las métricas de rendimiento son “rmse”: 13,039, “mse”: 170,010, “mae”: 7,157. Es posible que haya algún problema de sobreajuste con esta configuración.

También comparamos lo que ocurre al añadir unas cuantas capas densas:

• Después de agregar una capa densa, tenemos “rmse”: 15.435, “mse”: 238.265, “mae”: 9.351
• Después de agregar tres capas densas, tenemos “rmse”: 15.435, “mse”: 238.241, “mae”: 9.505

Por último, realizamos una comparación con el método CRISPRon existente. CRISPRon es un modelo de aprendizaje profundo basado en CNN. Las métricas de rendimiento son “rmse”: 11,788, “mse”: 138,971, “mae”: 7,134.

Como era de esperar, LoRA funciona mucho mejor que simplemente agregar unas pocas capas densas. Aunque el rendimiento de LoRA es un poco peor que el de CRISPRon, con una búsqueda exhaustiva de hiperparámetros, es probable que supere a CRISPRon.

Al utilizar los cuadernos de SageMaker, tiene la flexibilidad de guardar el trabajo y los datos producidos durante el entrenamiento, apagar la instancia y volver a encenderla cuando esté listo para continuar con el trabajo, sin perder ningún artefacto. Apagar la instancia le evitará incurrir en costos de computación que no está utilizando. Le recomendamos encarecidamente que solo la encienda cuando la esté utilizando activamente.

Conclusión

En esta publicación, mostramos cómo usar los métodos PEFT para ajustar los modelos de lenguaje de ADN con SageMaker. Nos centramos en predecir la eficiencia de las secuencias de ARN CRISPR-Cas9 para determinar su impacto en las tecnologías de edición genética actuales. También proporcionamos código que puede ayudarlo a poner en marcha sus aplicaciones de biología en AWS.

Para obtener más información sobre el área de atención médica y ciencias biológicas, consulte Ejecutar AlphaFold v2.0 en amazon EC2 o ajustar e implementar el modelo ProtBERT para la clasificación de proteínas usando amazon SageMaker.

Acerca de los autores

Siddharth Varia es un científico aplicado en AWS Bedrock. Está ampliamente interesado en el procesamiento del lenguaje natural y ha contribuido a productos de AWS como amazon Comprehend. Fuera del trabajo, le gusta explorar nuevos lugares y leer. Se interesó en este proyecto después de leer el libro The Code Breaker.

Yudi Zhang es científica aplicada en marketing de AWS. Sus intereses de investigación se centran en el área de redes neuronales gráficas, procesamiento del lenguaje natural y estadísticas.

Erika Peláez Coyotl es científica aplicada sénior en amazon Bedrock, donde actualmente colabora en el desarrollo del modelo de lenguaje grande amazon Titan. Tiene experiencia en ciencias biomédicas y ha ayudado a varios clientes a desarrollar modelos de aprendizaje automático en este sector.

Zichen Wang es un científico aplicado sénior en AWS ai Research & Education. Está interesado en investigar redes neuronales gráficas y aplicar IA para acelerar el descubrimiento científico, específicamente en moléculas y simulaciones.

Rishita Anubhai Es científica aplicada principal en amazon Bedrock. Tiene una gran experiencia en procesamiento de lenguaje natural y ha contribuido a proyectos de AWS como amazon Comprehend, Machine Learning Solutions Lab y el desarrollo de modelos de amazon Titan. Está muy interesada en utilizar la investigación en aprendizaje automático, específicamente el aprendizaje profundo, para crear un impacto tangible.