Большая фишка для большой науки: наблюдение за вирусом COVID-19 в действии

Трудно представить лучший пример ИИ во благо, чем выяснить, как работает вирус, вызывающий COVID-19. Если мы узнаем, как это работает, мы сможем разработать способы предотвращения репликации вируса и положить конец глобальному бедствию. Мне посчастливилось стать соавтором вместе с моими коллегами Джессикой Лю и Танвиром Раза масштабного научного исследования, чтобы выяснить, как это сделать. Исследование было номинировано на Специальную премию Гордона Белла и было представлено на прошлогодней суперкомпьютерной конференции SC21.

Это большая наука. Исследователи из 12 национальных лабораторий, университетов и таких компаний, как Cerebras Systems, разработали множество новых вычислительных методов для моделирования механизма репликации вируса, работающего на 4 суперкомпьютерных площадках!

Идея заключалась в том, чтобы создать полнофункциональную модель «механизма репликации-транскрипции» вируса SARS-CoV-2. Уместно слово «машина»: это сложный биологический механизм, состоящий из миллионов атомов, движущихся в трех измерениях, когда он захватывает собственный механизм репликации хозяина, чтобы создавать свои копии.

Процесс начинается с трехмерных изображений вируса, полученных с помощью криоэлектронной микроскопии. Этот метод позволяет достичь разрешения, близкого к атомарному, но изображения все еще недостаточно хороши или недостаточно динамичны, чтобы показать нам, как на самом деле работает механизм. Чтобы восполнить недостающие данные, исследовательская группа наложила сверху два совершенно разных, но дополняющих друг друга метода, работающих в разных масштабах. Во-первых, мы можем относиться к биомолекулам так же, как к любой проблеме с материалами. Мы можем использовать тип инструментов анализа конечных элементов, которые мы обычно используем для проектирования объектов непрерывного масштаба, таких как детали двигателя. И, во-вторых, мы можем моделировать молекулы атом за атомом, как гораздо более сложную версию моделей шарика и стержня, которые мы все помним из уроков химии.

Собрать все это воедино — колоссальная задача. Инновации требовались, так сказать, на всех уровнях, от новой архитектуры рабочего процесса, которая позволяет широко распределенным вычислительным ресурсам беспрепятственно и автоматически объединяться, до повышения вычислительной эффективности отдельных моделей.

В этой последней части — повышении эффективности вычислений — вступает в игру Cerebras. В прошлом на создание этих симуляций уходило так много времени, что за один раз можно было изучить только несколько десятков наносекунд движения. Однако, чтобы достичь более широкого понимания, им нужно было изучить микросекунды — в 50 раз больший период времени. Команда поняла, что этапы машинного обучения были узким местом для достижения 50-кратного ускорения, необходимого при интеграции моделирования с ИИ.

Вы спросите, какую именно роль здесь играет ML? Каждый эксперимент по моделированию надолго связывает суперкомпьютер с тысячами вычислительных узлов. Чтобы избежать напрасной траты времени, жизненно важно «направлять» эти эксперименты, распознавая и останавливая симуляции, которые заходят в тупик, и, так сказать, поощряя те, которые могут оказаться плодотворными. Это легче сказать, чем сделать. Заранее указать характеристики «плохой» симуляции очень сложно. Но легко постфактум признать, что случилось плохое. Это классическая возможность машинного обучения: вы знаете, как выглядит ответ, но не знаете, как определить правила для его описания.

Мы решаем эту проблему с помощью модели машинного обучения, называемой «сверточный вариационный автоэнкодер» или CVAE. Упрощая, CVAE берет сложные «многомерные» входные данные и преобразует или «кодирует» их в меньшую форму. Вы можете думать об этом как о своего рода фигуре достоинства. Мы обучаем модель, позволяя ей наблюдать за моментальными снимками симуляций. Затем мы запускаем обратное преобразование — или декодируем его. Если декодированная версия хорошо соответствует оригиналу, мы знаем, что CVAE работает. Затем эту обученную модель можно использовать во время «настоящих» экспериментов с другим алгоритмом, который выполняет реальное управление. Однако, как отмечается в документе: «CVAE квадратичен по временной и пространственной сложности и может быть непомерно сложным для обучения».

Cerebras вступает в игру здесь, потому что эта часть проблемы исследовалась в Ок-Риджской национальной лаборатории на суперкомпьютере Саммит и на Аргоннском испытательном стенде ИИ в Аргоннской национальной лаборатории, где как раз есть ускоритель Cerebras. . Исследователи ANL сравнили обучение своей модели CVAE на 256 узлах Summit, в общей сложности на 1536 графических процессорах и на одной системе Cerebras CS-2.

И как мы поступили? С точки зрения чистой производительности, довольно хорошо. Снова цитируя документ: «CS-2 обеспечивает производительность 24 000 выборок в секунду, что эквивалентно 110–120 графическим процессорам».

Каким бы впечатляющим ни было это число, возможно, еще более впечатляющим является комментарий из коробки. Распределить многообещающий алгоритм по большому кластеру вычислительных узлов сложно и отнимает много времени даже у специалистов в этой области. Система CS-2, напротив, намеренно спроектирована как единый сверхмощный узел с производительностью кластерного масштаба. Наше программное обеспечение позволяет легко запустить нейронную сеть, изменив всего пару строк кода.

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

Еще раз процитируем статью: «Поскольку один CS-2 обеспечивает производительность более 100 графических процессоров, это практическая альтернатива для организаций, заинтересованных в этом рабочем процессе, которые не имеют очень больших кластеров графических процессоров». Мы не могли не согласиться.

Наконец, важно иметь в виду, что, хотя это исследование имеет прямые преимущества в лечении COVID-19, новые инструменты и рабочий процесс могут в конечном итоге оказаться гораздо более значимыми. Эта методология может быть применена к любому типу молекулярного оборудования, прокладывая путь к более быстрому и лучшему пониманию молекулярных взаимодействий в широком диапазоне вариантов использования, включая разработку методов лечения ряда заболеваний. Очень приятно осознавать, что я смог сыграть свою роль.

Чтобы узнать больше об исследовании, прочитайте статью Интеллектуальное разрешение: интеграция крио-ЭМ с моделированием с несколькими разрешениями на основе ИИ для наблюдения за механизмом репликации-транскрипции SARS-CoV-2 в действии», которая будет опубликована в International Journal. приложений для высокопроизводительных вычислений, 2021 г.

Изображение баннера Аргоннской национальной лаборатории/Университета Иллинойса в Урбане-Шампейне.

(Репост из Блога Церебра)