gonzo-обзоры ML статей – Telegram
gonzo-обзоры ML статей
24.1K subscribers
2.72K photos
2 videos
3 files
1.34K links
Авторы:
Гриша Сапунов, ранее руководитель разработки Яндекс-Новостей, ныне CTO Intento. Области интересов: AI/ML/DL, биоинформатика.
Лёша Тихонов, ранее аналитик в Яндексе, автор Автопоэта, Нейронной Обороны... Области интересов: discrete domain, NLP, RL.
Download Telegram
Ещё пробуют ансамблирование по нескольким последовательностям из MSA, качество улучшается.

Снова (как и в предыдущей работе) показали, что перплексия хорошо коррелирует с точностью предсказания контактов (и снова видно, что есть куда скейлить).

Логрегрессия выделила 102 из 660 голов, которые помогают предсказаниям. Две головы оказались в top-10 по предсказаниям на всех масштабах. Одна голова оказалась особенно хороша для предсказания локальных контактов и мешала предсказывать дальние. И в целом видна специализация по типам контактов.

Ну и наконец показывают, что вероятности, которые выдаёт логрегрессия, близки к реальным вероятностям контактов, модель неплохо откалибрована и не overconfident. А в кейсах, где она ошибочно предсказывает контакт, он часто оказывается в пределах нескольких аминокислот от предсказанной точки, либо оказывается контактом между двумя цепочками гомодимера. К вставкам модель устойчива, можно вставить в начало/конец 256 аланинов (или 64 в середину) и модель не особенно деградирует. Более раннему трансформеру TAPE хватало 8 аланинов в начале, чтобы качество сильно ухудшилось (возможно, потому что тот обучался на белковых доменах).

Работа взята постером на ICLR 2021 (https://openreview.net/forum?id=fylclEqgvgd).
MSA Transformer
Roshan Rao, Jason Liu, Robert Verkuil, Joshua Meier, John F. Canny, Pieter Abbeel, Tom Sercu, Alexander Rives
Статья: https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1
Модель: https://github.com/facebookresearch/esm

Продолжим тему про DL в биологии. Этот пост является логическим продолжением и развитием работ из https://news.1rj.ru/str/gonzo_ML/609 и https://news.1rj.ru/str/gonzo_ML/621 от той же команды.

В предыдущих работах на последовательностях одиночных белков (но в гигантском количестве) обучили языковую модель в стиле BERT (на Masked Language Modeling задаче, когда маскируются отдельные токены входа и модель должна восстановить их на выходе) и показали, что модель выучивает интересные и полезные для других задач фичи. Всё как в NLP в общем.

Но в биологии есть более богатый подход (что, кстати, могло бы быть его аналогом в NLP?) — работать не на одиночных последовательностях, а на множественных выравниваниях (Multiple sequence alignment, MSA), когда к последовательности добавляется множество её эволюционных родственников, которые все выровнены по позициям. На таком выравнивании сразу видно, где произошли замены аминокислот в белке и какие позиции в белке более консервативны (а значит и более важны), и где были вставки и делеции. Посмотреть примеры и подробнее почитать про MSA можно в Википедии https://en.wikipedia.org/wiki/Multiple_sequence_alignment.

При обучении на большом датасете одиночных последовательностей сеть неявно всё равно вытягивает это знание, но оно оказывается зашито внутрь сети. Альтернативно можно сделать этот процесс более явным, подавая выравнивание на вход, тогда и сети вроде как не надо всё это выучивать, а можно извлекать из входных данных, что по идее проще и потребует меньше параметров (spoiler: в реальности обученная в работе сеть использует микс обоих вариантов).

Встречайте MSA Transformer — трансформер, работающий на множественных выравниваниях.

Вход теперь, в отличие от обычного трансформера, получается двумерным и в модели чередуются блоки внимания по строкам и по столбцам. Это по сути практически архитектура Axial Transformer, мы про неё подробно не писали, но упоминали в разных местах (в посте про картиночный трансформер https://news.1rj.ru/str/gonzo_ML/434 и в слайдах про разные трансформеры https://youtu.be/KZ9NXYcXVBY?t=3082 и https://youtu.be/7e4LxIVENZA?t=976).

Если делать это на обычном трансформере, растягивая MSA в одномерную последовательность, то будет слишком дорого, получится сложность O(M^2*L^2), где M — число последовательностей, а L — длина выравнивания. А так выходит O(ML^2) для блоков внимания по строкам и O(LM^2) для внимания по столбцам. Этот вариант трансформера ещё и шарит карту внимания между всеми строками, так что для строк получается даже O(L^2).

Есть ещё мелкие изменения по части нормализации и расположения полносвязных слоёв.

Задача обучения всё та же — MLM. Модель пытается восстановить повреждённую MSA. То есть всё тот же self-supervised learning.

Обучают трансформер на 100М параметров (скромно на фоне ESM-1b, который был на 650М и содержал 34 слоя; здесь всего 12 слоёв) на большом (4.3 Тб) датасете в 26М множественных выравниваний. Отдельное выравнивание генерится для для каждой последовательности из UniRef50 поиском в UniClust30 с помощью HHblits. Средняя глубина MSA получается 1192.

При инференсе пробовали считать не на полном MSA, а на подвыборке с разными стратегиями сабсемплинга.

Результаты?

На задаче unsupervised contact prediction сравниваются с ESM-1b (650M параметров) и ProTrans-T5 (3B, видимо ProtT5-XL, про него как-нибудь отдельно), а также с Potts model. Используют методику из предыдущей статьи, где с помощью разреженной логрегрессии определяли головы внимания, хорошо справляющиеся с данной задачей. При инференсе делали подвыборку в 256 последовательностей с помощью hhfilter.

MSA трансформер побил все бейзлайны (и ESM-1b побил ProTrans-T5). Также сравнились на CAMEO hard targets и CASP13-FM, на последнем сравнимо с supervised trRosetta_base.
На supervised contact prediction (на фичах MSA transformer обучают resnet) побили в том числе более качественную trRosetta_full.

Проверились на задаче предсказания вторичной структуры на CB513, обошли профили HMM, ESM-1b и Netsurf.

Поэкспериментировали с разнообразием последовательностей в MSA. Стратегии, максимизирующие разнообразие, дают качество выше ESM-1b с всего 16 последовательностями на входе. И добавление всего одной последовательности в высоким diversity перекрывает 31 последовательность с низким.

Головы внимания к строкам любят столбцы с высокой изменчивостью. А внимание по столбцам больше смотрит на более информативные последовательности. Также логрегрессия нашла 55 (из 144) голов, предсказывающих контакты.

Такие дела. Скейлить по идее есть куда, ждём продолжения.
Accurate prediction of protein structures and interactions using a 3-track network
Minkyung Baek, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N. Kinch, R. Dustin Schaeffer, Claudia Millán, Hahnbeom Park, Carson Adams, Caleb R. Glassman, Andy DeGiovanni, Jose H. Pereira, Andria V. Rodrigues, Alberdina A. van Dijk, Ana C. Ebrecht, Diederik J. Opperman, Theo Sagmeister, Christoph Buhlheller, Tea Pavkov-Keller, Manoj K Rathinaswamy, Udit Dalwadi, Calvin K Yip, John E Burke, K. Christopher Garcia, Nick V. Grishin, Paul D. Adams, Randy J. Read, David Baker
Статья: https://www.biorxiv.org/content/10.1101/2021.06.14.448402v1
Модель: https://github.com/RosettaCommons/RoseTTAFold

С того момента, как AlphaFold2 всех побил на соревновании по предсказанию белковых структур CASP14 (https://www.nature.com/articles/d41586-020-03348-4) было много обсуждений и домыслов, что у этой штуки внутре. Некоторые команды пытаются собрать свою версию AlphaFold2 (https://github.com/lucidrains/alphafold2) на основе доступной информации. DeepMind вот уже обещает детальную статью (в стадии ревью) и код в опенсорсе (https://twitter.com/demishassabis/status/1405922961710854144), но пока не выдал. В это время другие исследователи предлагают свои альтернативные подходы к таким задачам.

Вот недавно большая команда исследователей предложила свою сеть, работающую с тремя типами данных (они это называют 3-track network, что, наверное, можно назвать трёхмодальной сетью): 1D информация о последовательностях, 2D карта расстояний и 3D уровень координат атомов. AlphaFold2 вроде как 2-track network (1D и 2D, а часть с 3D и итеративным предсказанием структуры там вроде на выходе уже после того как полностью отработали первые две, но это не точно).

Цель сети — предсказать трёхмерные координаты атомов N, Cɑ, C каждой аминокислоты.

Независимо от результата интересен подход. В модели происходит довольно активный обмен данными между разными треками. И в работе собраны несколько интересных вещей, редко встречающихся в других областях.

Пайплайн работает так:

Для целевого белка через HHblits ищутся гомологи в UniRef30 и BFD, строится множественное выравнивание (MSA).

По MSA делается поиск белков-шаблонов (или как по-русски называются эти templates?) в PDB100. Далее из известных структур шаблонов берутся попарные расстояния и ориентации для выравненных позиций, плюс собираются дополнительные фичи (скоры выравнивания, вероятности HHsearch и т.д.).

MSA обрабатывается аналогом MSA трансформера с axial attention. Для колоночного внимания используется архитектура Performer (https://news.1rj.ru/str/gonzo_ML/404), уменьшающая требования к памяти с O(LM^2) до O(LM), где L — длина выравнивания, а M — число последовательностей. Для внимания по строкам берут модифицированный вариант зашаренного внимания, уменьшающий вклад невыравненных участков с помощью обучаемого зависящего от позиции веса (это называется soft-tied attention).

Парные фичи (2D) обновляются по сигналу из фич MSA с использованием axial attention также с Performer’ом. MSA фичи в свою очередь обновляются на основе информации, закодированной в парных фичах (и эта идея взята из AlphaFold2).

Далее задействован графовый трансформер архитектуры UniMP (https://arxiv.org/abs/2009.03509), он до сих пор довольно высоко висит в лидерборде Open Graph Benchmark (https://ogb.stanford.edu/). На вход ему прилетает полносвязный граф, где узлы это аминокислоты белка, MSA фичи дают эмбеддинги узлов, а парные фичи — эмбеддинги рёбер. Этот трансформер выдаёт (через 4 слоя графовых слоёв и один слой линейной трансформации) на выходе начальные координаты бэкбона белка для 3D трека.

Затем подключается хитрый SE(3)-трансформер (https://arxiv.org/abs/2006.10503), который эквивариантен относительно трёхмерных вращений и перемещений . Он уточняет xyz координаты на основе обновлённых MSA и парных фич.