LLaMo: Large Language Model-based Molecular Graph Assistant

Молекулярное машинное обучение в последние годы получило значительное внимание за счет решения разнообразных задач в химической области. Преобладающий подход для выполнения молекулярных задач — это машинное обучение на графах, которое использует структуру молекулярного графа как естественное и выразительное представление молекул. Хотя методы, основанные на графах, успешно представляют молекулы, они имеют ограниченную интерпретируемость и несовместимость для решения многомодальных молекулярных задач, которые включают взаимодействие текста и молекул.

Недавние работы в области машинного обучения на графах обучали как языковую модель, так и графовый энкодер с помощью кросс-модального контрастивного обучения. Однако такие модели не способны выполнять задачи генерации текста на основе молекул, которые более применимы на практике. Большие языковые модели (LLM) показали впечатляющий прогресс и достигли человеческого уровня в генерации открытого текста благодаря миллиардам параметров. Для использования способности LLM следовать инструкциям, многие исследования используют подходы инструкционного тюнинга для создания общих языковых моделей.

Мотивация для разработки LLaMo (Large Language Model-based Molecular graph assistant) возникла из успехов LLM и инструкционного тюнинга, а также из недостаточного изучения интеграции языковых моделей с молекулярными графами. В этой статье мы представляем LLaMo, который бесшовно интегрирует молекулярный графовый энкодер и большую языковую модель для генерации ответов, следующих инструкциям, в молекулярной области.

Основные компоненты LLaMo

Архитектура модели

Архитектура LLaMo включает три основных компонента:

1. Молекулярный графовый энкодер: Использует графовые нейронные сети (GNN) для кодирования 2D молекулярного графа в последовательность токенов.

2. Многоуровневый графовый проектор: Преобразует последовательность узловых представлений в молекулярные токены, выравнивая их с LLM. Этот проектор захватывает информацию о графе на нескольких уровнях, включая атомы, атомные группы и молекулы.

3. Большая языковая модель: Обрабатывает молекулярные и текстовые токены для генерации ответов в автопрогрессивном режиме.

Многоуровневый графовый проектор

Цель многоуровневого графового проектора — выровнять графовый энкодер с LLM, преобразуя набор узловых представлений в последовательность молекулярных графовых токенов. Это позволяет языковой модели использовать информацию о графе. В литературе проекторы обычно реализуются с помощью линейной проекции или абстракции визуальных признаков. Однако, мы наблюдаем, что высокие уровни представлений неэффективны в захвате локальной информации из-за проблемы переусреднения (over-smoothing), когда представления узлов становятся неразличимыми с увеличением количества слоев GNN.

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

Обучение LLaMo

Обучение LLaMo проходит в два этапа:

1. Предобучение для выравнивания графа и языка: На этом этапе, с замороженной LLM, обучаются многоуровневый графовый проектор и графовый энкодер на датасете пар молекул-описаний (например, PubChem).

2. Инструкционный тюнинг: Здесь мы обучаем LLM для улучшения способности следовать инструкциям и понимать молекулярные графы. Графовый энкодер замораживается, а обучаются многоуровневый графовый проектор и LLM с использованием LoRA (Low-Rank Adaptation) для адаптации LLM к данным. Для этого этапа используется набор данных инструкций, сгенерированный с помощью GPT-4.

Генерация данных инструкций с использованием GPT-4

Инструкционные данные необходимы для улучшения способности LLM следовать инструкциям. В литературе инструкционные данные для молекулярных графов изучены недостаточно из-за необходимости экспертных знаний для аннотации. Для решения этой проблемы мы используем GPT-4 для генерации данных инструкций в формате многоходовых диалогов. Эти данные включают три типа контекстов: SMILES представление молекулы, описание молекулы и IUPAC название молекулы.

Процесс генерации данных включает три шага:

1. Выбор образцов: Выбор образцов диалогов среди машинно-генерированных инструкционных данных.
2. Генерация диалогов: Создание многоходовых диалогов с помощью обучения в контексте с использованием выбранных образцов в качестве подсказок.
3. Фильтрация: Удаление неполных диалогов и диалогов с большим количеством ходов.

Эксперименты

Настройки экспериментов

Мы оценивали эффективность LLaMo на трех задачах: генерация описаний молекул, предсказание IUPAC названий и предсказание свойств молекул. Эксперименты проводились в двух режимах: общий и специализированный. В общем режиме одна модель обрабатывает все три задачи, тогда как в специализированном режиме модель обучается для каждой конкретной задачи.

Результаты экспериментов

Общие модели: LLaMo, построенный на LLaMA-7B и настроенный с помощью нашего метода инструкционного тюнинга, показал лучшие результаты во всех трех задачах по сравнению с другими моделями, включая GPT-4 с обучением в контексте.

Специализированные модели: LLaMo также показал превосходные результаты в специализированном режиме, превзойдя другие специализированные модели на различных датасетах.

Анализ

• Влияние многоуровневого графового проектора: Наш проектор показал значительное улучшение по сравнению с традиционными проекторами, захватывая информацию на различных уровнях структуры молекулы.
• Влияние данных инструкций, сгенерированных с помощью GPT: Использование данных, сгенерированных с помощью GPT-4, улучшило производительность модели по сравнению с обучением без этих данных.
• Инструкционный тюнинг против многозадачного обучения: Инструкционный тюнинг показал лучшие результаты, чем многозадачное обучение, что подчеркивает эффективность нашего подхода.

Заключение

Мы представили LLaMo, первую модель, которая интегрирует молекулярные графы и большие языковые модели для выполнения различных молекулярных задач с помощью одного универсального подхода. Многоуровневый графовый проектор и инструкционный тюнинг с использованием машинно-генерированных данных значительно улучшили способность модели понимать и генерировать текст на основе молекулярных структур. LLaMo демонстрирует потенциал для дальнейших исследований и применений в области молекулярного машинного обучения и химии.