Энергоэффективные модели языка белков: Использование малых языковых моделей с LoRA для контролируемой генерации белков

Большие языковые модели (LLMs) кардинально изменили сферу обработки естественного языка (NLP), продемонстрировав исключительные способности в различных задачах, от генерации текста до перевода и анализа настроений. Модели, такие как GPT-4 и Llama 3, получили статус "фундаментальных моделей" благодаря своей способности обобщать на множество задач путем дообучения на специализированных наборах данных.

Белковые последовательности аналогичны естественному языку, так как оба состоят из набора базовых строительных блоков — аминокислот для белков и слов для языка — которые комбинируются для формирования значимых структур. В белках последовательность аминокислот определяет трехмерную структуру и функцию, тогда как в естественном языке последовательность букв и слов определяет семантическое значение и локальный контекст. Эти сходства привели к применению LLMs для различных задач, таких как предсказание структуры белка по его последовательности, de-novo генерация белковых последовательностей, также известная как неконтролируемая генерация, и генерация белков с определенными свойствами или специализированных для определенной задачи, известная как контролируемая генерация.

Дизайн белков находит свое применение в разработке лекарств и создании эффективных искусственных ферментов, которые могут разлагать промышленные отходы или пластик, способствуя достижению углеродной нейтральности.

Разработка малых моделей языка белков

В данной работе представлены две малые модели языка белков, основанные на архитектурах Llama-3-8B и Phi-3-mini. Эти модели способны как к неконтролируемой, так и к контролируемой генерации белков. Для задачи неконтролируемой генерации наша лучшая модель достигает среднего pLDDT (predicted Local Distance Difference Test) значения 69.75 ± 12.74, демонстрируя устойчивую производительность в генерации жизнеспособных белковых структур. Для задачи контролируемой генерации, где модель генерирует белки согласно свойствам, указанным в запросе, мы достигли замечательного среднего значения TM-Score (Template Modeling Score) 0.84, что указывает на высокую структурную схожесть с целевыми белками.

Мы выбрали 10 свойств, включая шесть классов ферментов, для расширения возможностей предыдущих моделей языка белков. Наш подход использует технику адаптации с низким рангом (Low-Rank Adaptor, LoRA), что позволяет сократить количество обучаемых параметров до всего лишь 4% от исходного размера модели, уменьшая вычислительные требования. Используя подмножество набора данных UniRef50 и малые модели, мы сократили общее время обучения на 70% без ущерба для производительности. В частности, модель Phi-3-mini уменьшила количество обучаемых параметров на 60%, снизив затраты на обучение на 30% по сравнению с Llama 3. В результате, Phi-3 достигла сравнимого TM-Score 0.81, демонстрируя, что меньшие модели могут соответствовать производительности более крупных, таких как Llama 3.

Энергоэффективность и развертывание

Мы также демонстрируем развертывание наших моделей на энергоэффективном чипе ET-SoC-1, значительно улучшив метрику TPS/W (Tokens Per Second per Watt) в три раза. Модели доступны на платформе Hugging Face по адресам https://huggingface.co/Esperanto/Protein-Llama-3-8Band и https://huggingface.co/Esperanto/Protein-Phi-3-mini, что способствует дальнейшим исследованиям и разработкам в области моделей языка белков.

Методы

Подготовка данных

Для первого этапа обучения, целью которого является создание модели, способной генерировать белки без условий, мы использовали 2 миллиона последовательностей из релиза 2024_03 набора данных UniRef50. Чтобы продемонстрировать достижение сравнимой производительности с меньшими ресурсами для обучения, мы не использовали полный набор данных. Для второго этапа обучения, включающего дообучение на основе инструкций, мы выбрали 10 свойств, чтобы показать эффективность нашей модели в контролируемой генерации, относящейся к этим свойствам. Мы также рассмотрели задачу генерации ферментов, данный класс фермента, что ранее не исследовалось. Шесть классов ферментов, которые мы рассмотрели, включают: оксидоредуктазы, лигазы, трансферазы, изомеразы, гидролазы и лязы.

Пайплайн обучения

Для наделения базовой LLM модели пониманием белкового языка необходимо обучить их на наборе данных белковых последовательностей. Однако полная реинициализация параметров для предобучения имеет два недостатка: (i) огромные вычислительные ресурсы, необходимые для этого, и (ii) устранение предыдущего понимания естественного языка, что важно для многофункциональных возможностей модели, понимая ввод пользователя. Чтобы избежать этого, мы использовали непрерывное обучение с LoRA, основанное на подходе, аналогичном Lv et al. Это сокращает количество обучаемых параметров и помогает сохранить предыдущее знание естественного языка, позволяя модели интерпретировать инструкции пользователя.

Мы использовали этот подход для моделей Llama-3-8B и Phi-3-mini-4k-instruct. Модели обучались для задачи причинного языкового моделирования, где функция потерь, используемая, — это кросс-энтропия между выходом модели и истинным значением, которое просто является входами, сдвинутыми на один токен влево.

Оценка структуры

В отличие от образцов естественного языка, которые могут быть оценены на основе их грамматической правильности, когерентности и релевантности запросу людьми, оценка белковых последовательностей требует специализированной модели для предсказания структурной стабильности и качества генерируемых белков. Это измеряется с помощью pLDDT (predicted Local Distance Difference Test), который является мерой уверенности модели в своем предсказании структуры и коррелирует с упорядоченностью структуры.

Контролируемая генерация

Чтобы оценить, генерирует ли модель белки в соответствии с классом, указанным в запросе, мы генерировали 100 последовательностей на класс. Эти последовательности затем передавались в модель ESMFold для получения их pLDDT оценок, и последовательности с оценкой ниже 60 отфильтровывались, чтобы гарантировать использование только последовательностей, соответствующих стабильным структурам, для оценки.

Результаты

Неконтролируемая генерация

Генерация последовательностей без предоставления класса белков в качестве контекста называется неконтролируемой генерацией. Мы генерировали 100 последовательностей и оценивали их на основе их структурной стабильности, как количественно измерялось через их pLDDT оценки, и меры их соответствия двум базам данных белковых структур - AlphaFold Protein Structure Database (AFDB) и Protein Data Bank (PDB).

Контролируемая генерация

Генерация белков в соответствии с некоторыми пользовательскими условиями важна для дизайна белков, полезных для различных приложений. Мы оценили производительность наших моделей по генерации белков, принадлежащих к 10 различным классам, как выделено в таблице 3, на основе их TM-Score и RMSD с эталонными белками.

Заключение

В данной работе мы представили две компактные модели языка белков, основанные на Llama-3-8B и Phi-3-mini, способные как к неконтролируемой, так и к контролируемой генерации белковых последовательностей. Наши модели демонстрируют хорошую производительность в генерации жизнеспособных белковых структур и синтезе белков с определенными свойствами, эффективно сокращая разрыв между большими языковыми моделями в обработке естественного языка и специализированными моделями языка белков.