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

эквивариантность — это когда результат функции от трансформированного входа соответствует трансформированному результату от нетрансформированного входа, или f(T(x))=T(f(x))

инвариантность — это когда результат функции от трансформированного входа такой же, как и от нетрансформированного, или f(T(x))=f(x)

Если кому интересно подробнее про это на примере свёрток, то вот неплохой разбор на кошечках: https://www.doc.ic.ac.uk/~bkainz/teaching/DL/notes/equivariance.pdf
--------------

И далее, карты внимания, полученные из 3D структур, используются для обновления фич MSA через хитрое внимание с маской (так что внимание смотрит только на близлежащие в пространстве атомы, это делается с четырьмя порогами в 8, 12, 16 и 20 ангстрем).

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

Модель получилась большая и за раз способна обрабатывать небольшие кусочки белков (кропы) по 260 аминокислот. Для получения финального предсказания авторы комбинируют и усредняют 1D фичи и 2D расстояния и ориентацию для кропов, а затем используют два подхода для генерации финальной 3D структуры. В первом подходе предсказанные расстояния между аминокислотами и ориентации отправляются в pyRosetta для генерации всех атомов; во втором 1D и 2D фичи отправляются в ещё один SE(3)-трансформер, который обучается выдавать 3D координаты. Первая сеть требует меньше памяти: 8Gb RTX2080 для инференса, но зато больше CPU для pyRosetta. Вторая требует 24Gb TITAN RTX.

Всего в модели 130М параметров, 8 2-track блоков, 5 3-track блока. Обучали 4 недели на 8 V100 32Gb. Модель называется RoseTTAFold.

Качества AlphaFold2 авторы не добились, но отчасти возможно это в силу разных мощностей — в случае AlphaFold2 предсказания для одиночных белков считались сутками на нескольких GPU, а здесь делаются за один forward-pass сетки и требуют порядка 10 минут на RTX2080 (плюс 1.5 часа на поиск шаблонов и MSA на входе).

Такая вот сложная система и чудо инженерной мысли.

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

Есть ещё экспериментальная биологическая часть, когда попытались применить RoseTTAFold к белкам неизвестной структуры и что-то про них поняли, но за этим отсылаю к статье.
👍1
Пир духа.

В Nature за 15 июля статья DeepMind про AlphaFold2:
https://www.nature.com/articles/d41586-021-01968-y
https://www.nature.com/articles/s41586-021-03819-2

Код выложен на гитхаб:
https://github.com/deepmind/alphafold


В Science за 15 июля статья про RoseTTAFold (писали про неё недавно, https://news.1rj.ru/str/gonzo_ML/642):
https://www.science.org/doi/abs/10.1126/science.abj8754

Этот код уже был на гитхабе ранее, когда мы про эту модель писали:
https://github.com/RosettaCommons/RoseTTAFold
Highly accurate protein structure prediction with AlphaFold
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli & Demis Hassabis
Статья: https://www.nature.com/articles/s41586-021-03819-2
Модель: https://github.com/deepmind/alphafold

Ну вот наконец и AlphaFold2. В статье и на гитхабе новая система называется просто AlphaFold, а под названием AlphaFold2 она фигурировала на соревновании CASP14. Новая система является совершенно другой системой по сравнению со старой AlphaFold, участвовавшей на CASP13, так что будьте внимательны и не путайте.

Итак, собственно, какова структура этой системы.

На вход системы поступает последовательность аминокислот белка. Для этой последовательности ищутся эволюционные родственники и генерируется множественное выравнивание (MSA, про него мы несколько раз писали в предыдущих постах о похожих моделях). Также в базе структур белков (с 3D координатами атомов) ищутся гомологичные структуры (если есть), и они тоже поступают на вход.

Сама сеть AlphaFold напрямую предсказывает 3D координаты всех тяжёлых атомов (углероды, азоты) по входным данным.

Сеть состоит из двух частей.

Сначала ствол сети обрабатывает входы с помощью последовательности (аж 48 штук, каждый со своими весами) новых блоков под названием Evoformer (хитрый вариант трансформера). Evoformer выдаёт на выходе тензор с обработанным MSA (размера N_seq*N_res, где N_seq — число последовательностей в MSA, а N_res — число аминокислот) и тензор (размера N_res*N_res) с репрезентацией пар аминокислотных остатков.

За стволом сети следует структурный модуль (structure module), генерирующий вращения и смещения для каждого из аминокислотных остатков (инициализируются они identity вращениями и позициями в начале координат) с помощью 8 слоёв с расшаренными весами.

И структурный модуль, и вся сеть целиком итеративно улучшает свои предсказания, подавая их себе же на вход. Этот процесс называется в работе “recycling”, он выполняется трижды, и существенно повышает качество результата.

Главная хитрость Evoformer’а в обмене информацией между репрезентациями MSA и пар аминокислот, образующих по сути отдельные стримы обработки данных.

Для обработки MSA используется вариант axial attention с чередующимся вниманием по строкам и по столбцам, причём оно gated через сигмоиду. Также есть transition layer, который по сути двуслойный MLP.

Стрим для обработки репрезентаций пар аминокислот по сути представляет собой графовый трансформер, узлами графа выступают соседние аминокислоты, а операции обновления репрезентаций работают над треугольниками аминокислот (следствие из интуиции относительно необходимости соблюдения неравенства треугольника для расстояний между аминокислотами). Две основные операции там это triangle multiplicative update и triangle self-attention.

По репрезентациям MSA считается outer product между всеми позициями попарно и добавляется к соответствующим репрезентациям пар. Так происходит регулярный обмен между стримом MSA и стримом парных фич.

Структурный модуль работает с репрезентацией бэкбона белка, используя прилетающие на вход репрезентации пар и оригинальную последовательность белка из MSA. Структура бэкбона представлена N_res независимыми вращениями и трансляциями каждой конкретной аминокислоты относительно глобального фрейма. Это называется “residue gas” (как бы независимо плавающие аминокислотные остатки, которые мы пытаемся ориентировать в пространстве).
👍1
Репрезентации этого residue gas итеративно обновляются в два этапа. Сначала специальный оператор внимания под названием Invariant Point Attention (IPA) обновляет N_res активаций сети (одиночных абстрактных репрезентаций, которые изначально прилетели из стека Evorofmer’ов) и вроде как эта операция эквивариантная (но судя по абляциям IPA, и без этого норм). Затем обновляется бэкбон белка (предсказываются вращения и трансляции). И затем мелкий резнет предсказывает все (почти) торсионные углы. Из них, так понимаю, определяются уже позиции атомов. Здесь используется урезанный вариант FAPE loss (Frame-aligned point error), штрафующий за промахи в предсказании координат C_alpha атомов.

FAPE loss также является основным лоссом для обучения сети, в полном варианте он считается по всем атомам бэкбона и боковой цепи. Кроме него в сети используются разные другие дополнительные лоссы, например, BERT-подобный лосс для замаскированных элементов MSA, кросс-энтропийный лосс на распределение расстояний, и другие. В целом самое мясо находится в Supplementary Information файле, за всеми деталями надо туда.

На выходе модель также выдаёт confidence scores для каждой отдельной аминокислоты, и это используется при инференсе — прогоняется 5 обученных моделей (с различными seed) и для каждой цели выбирается наиболее уверенная модель.

Ещё из интересного, при обучении модели используется само-дистилляция. Обученная сеть используется для предсказания структур дополнительных 350К различных последовательностей из Uniclust30, из этого датасета отфильтровываются предсказания с высокой степенью уверенности, и затем та же самая архитектура обучается с нуля на смеси оригинальных данных из PDB (25%) и этого нового датасета (75%) с различными аугментациями (сабсэмплинг MSA). Это позволяет использовать кучу неразмеченных данных и существенно повышает качество.

Обучалось всё на 128 TPUv3 с батчем 128 (то есть по одной последовательности на каждом TPU) на кропах по 256 аминокислот. Потом файнтюнили на кусочках по 384 аминокислоты. Ушла неделя на обучение и ещё 4 дня на файнтюнинг.

По скорости инференса заявлено на V100 4.8 минуты для белка длиной 256 аминокислот, 9.2 минуты для 384 и 18 часов для 2500. Это типа быстрее, чем было на CASP14, потому что заюзали XLA. После CASP14 показали также, что ансамблирование сетей не очень много добавляет, зато без него всё в 8 раз быстрее.

В общем всё интересно, достойно отдельного более глубокого рассмотрения и подробного сравнения с RoseTTAFold.