[Перевод] Как работает DeepMind AlphaFold2?08.01.2022 22:46

DeepMind AlphaFold2 и Github Copilot являются одними из наиболее ярких успехов науки и техники в 2021 году. Спустя два года после исходного прорыва AlphaFold, специалисты из DeepMind выпустили вторую версию их революционной системы предсказания 3D-структуры белков по последовательности. На сей раз им, фактически, удалось решить задачу фолдинга белка, которая оставалась неразрешенной более 50 лет. В этом посте я постараюсь кратко описать принцип действия этой системы, как я его вижу.

Содержание

• Постановка задачи и принцип действия

• Поток данных в системе в целом

  • Последовательности: базы данных Uniprot, Mgnify; пакет программ HMMER

  • Структуры: база данных PDB; пакеты программ HHpred, OpenMM

  • Evoformer: эмбеддинги множественных выравниваний и парные представления, а также их итеративное улучшение

  • Структурный модуль: 3D структура из эмбеддинга последовательности

• Механизм внимания и трансформерная архитектура

  • Механизм внимания на масштабированном скалярном произведении: Запросы (Queries), Ключи (Keys) и Значения (Values), нестрогий поиск в базе данных ключ-значения

  • Многоголовое внимание на векторах эмбеддингов аминокислотных остатков

• Evoformer: уточнение эмбеддингов множественных выравниваний и парных представлений

  • Внимание по осям (Criss-cross (axial) attention)

  • Внимание по колонкам выравнивания, дополненное использованием информации из парных представлений

  • Внимание по последовательностям выравнивания

  • Уточнение парного представления из уточненного эмбеддинга выравнивания через внешнее произведение

  • Уточнение парного представления через мультипликативный модуль неравенства треугольника

  • Уточнение парного представления через механизм внимания, реализующий неравенство треугольника

• Структурный модуль

  • Внимание с неподвижной точкой (Invariant Point Attention, IPA)

  • Уточнение остова

  • Релаксация структуры физическими методами (OpenMM)

• Функции потерь: FAPE и дополнительные функции потерь

• Переработка (refinement/recycling)

• Самодистилляция (self-distillation)

• Протокол обучения и исследование абляции

Постановка задачи

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

Прочтение последовательности белка дешево и может быть произведено напрямую методами масс-спектрометрии (порядка $100 за образец) или секвенирования кодирующей данный белок ДНК ($100-$1000 за экзом/геном). Экспериментальное определение трехмерной структуры белка — гораздо более дорогая и трудоемкая задача ($100k-$1M за структуру), которая обычно решается методами рентгеновской кристаллографии или ЯМР.

Стоимость и вероятность успеха эксперимента по рентгеновской кристаллографии (данные 2003 года).

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

Конвейер эксперимента по рентгеновской кристаллографии. 1. Белок кристаллизуется. 2. Делается рентгеновская дифракция на кристалле. 3–4. Решается фазовая проблема. 5. Карты электронной плотности интерпретируется человеком, и по ним строится модель пространственной структуры белка.

Результат такого эксперимента нельзя считать абсолютной истиной, так как многие шаги в нем являются в какой-то мере искусством. Для получения структуры высокого разрешения требуется очень качественный кристалл (который может и не свариться, особенно для плохо растворимых мембранных белков). Различных методов решения фазовой проблемы немало. По одной и той же карте электронной плотности (особенно, низкого разрешения) два разных интерпретатора могут предложить структуры, существенно (процентов на 10–20%) различающиеся между собой. Наконец, белки в растворе живут и «дышат», и их рентгенные структуры похожи на «живой» белок примерно в той же мере, в какой фотографии замороженных трупов, сложенных штабелями в морге, похожи на живых людей. Уровень ошибки предсказания, достигнутый AlphaFold2 по сравнению с экспериментом, мало отличается от уровня разногласия между результатами двух разных групп, делающих расшифровку структуры одного и того же белка, поэтому я считаю, что AF2 скорее решил задачу фолдинга белка, чем не решил (пусть и с некоторыми оговорками).

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

Белковые последовательности известных белков содержатся в банке данных Uniprot (где их более 200 миллионов) и базе данных MGnify (о ней — позже).

Структуры белков доступны в базе данных PDB (содержащей менее 200 тысяч структур). В ней находятся файлы .pdb со структурами белков, каждый файл содержит координаты всех атомов или только тяжелых атомов (то есть не атомов водорода) белка.

Принцип действия AlphaFold2

Методы предсказания трехмерной структуры белка in silico, основанные на чистой физике, в общем случае не способны достигать результатов, сопоставимых с экспериментом. Статистический подход оказался более действенным.

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

Основная масса событий в эволюция носят нейтральный характер (т.е. большинство мутаций не влияют на функцию белка, поэтому в последовательностях белков со временем происходит множество замен). При этом структура белков значительно консервативнее последовательности. Типичным является случай, когда последовательность белка между не очень близкими видами измениалсь на 70%, а 3D-структура белка более-менее не изменилась. Сопоставление гомологов из разных видов позволяет получить важную дополнительную информацию о белках, и обычно делается посредством записи последовательностей белков в двумерную табличку, которая называется выравниванием (multiple sequence alignment, MSA). Последовательности белков из разных видов записываются в ряды, так что соответствующие друг другу аминокислотные остатки находятся в одних колонках.

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

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

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

Хочу также коротко отметить роль метагеномики в недавних успехах предсказания 3D-структур. Обычно белки в базах данных вроде Uniprot происходят из организмов, которые удается культивировать в лабораторных условиях и секвенировать. Однако, огромное множество микроорганизмов в лабораторных условиях выращивать не удается, а разнообразие белков в них огромно. Поэтому в 2000-ые годы некоторые исследователи пошли по другому пути: они стали секвенировать не отдельные виды, а просто зачерпывать образцы из среды, и секвенировать всю ДНК из этого бульона. Оказалось, что такие «метагеномные» образцы содержат массу новых гомологов известных белков, и использование данных коэволюции из них повысило качество предсказания 3D-структур белков в середине 2010-ых. AlphaFold2 использует такую метагеномную информацию из базы данных MGnify, которая используется как дополнение к Uniprot.

Конвейер AlphaFold2

Я использую термин «AlphaFold2» в двух контекстах: в широком смысле это вся система предсказания трехмерной структуры, использующая множество сторонних программ и баз данных. В узком смысле это нейронная сеть, лежащая в сердце этого конвейера.

Конвейер работы AlphaFold2 с высоты птичьего полета.

Система AlphaFold2 получает на вход последовательность белка. На первом шаге она использует программы из пакета HMMER для поиска гомологов входной последовательности в базах данных Uniprot и MGnify. Пакет HMMER был создан в начале 2000-ых, и основан на скрытых Марковских моделях/Марковских цепях. Тогда этот подход доминировал в распознавании/синтезе речи, и биоинформатики позаимствовали его для своих нужд. HMMER находит гомологи данного белка в базах и возвращает выравнивание всех гомологов, которые он нашел.

AF2 также использует пакет HH-suite для поиска 3D-структур по банку данных PDB, которые могли бы соответствовать гомологам поданной на вход последовательности. Если нам повезло, и для какого-то из гомологов последовательность белка уже экспериментально расшифрована, задача кратно упрощается — фактически, остается только провести гомологичное моделирование, то есть, грубо говоря, натянуть последовательность нашего белка на эту структуру. Однако, менее чем для 0.1% белков имеется расшифрованная структура, поэтому в общем случае нужно ожидать, что такой шаблон структуры будет недоступен.

Если 3D-структура все-таки нашлась, то на основе ее AF2 строит так называемую таблицу парного представления остатков — квадратную таблицу, где каждый элемент описывает взаимодействие i-ой аминокислоты с j-ой. Если 3D-структуры не нашлось, то такая таблица инициализируется разумными значениями по умолчанию.

Далее AF2 генерирует векторные эмбеддинги из аминокислотных остатков в выравнивании и из пар остатков в парном представлении. Я не буду углубляться в то, как это делается, но нетрудно придумать несколько способов. Только замечу, что для тонких выравниваний, в которых менее 30 последовательностей, AlphaFold2 работает не очень хорошо — недостаточно данных. Если же число последовательностей в выравнивании сильно больше сотен, то это тоже плохо, потому что сложность алгоритма растет квадратично от размера входа. В таком случае делается кластеризация наиболее близких последовательностей.

Далее следует ядро AlphaFold2: end-to-end нейронная сеть, основанная на архитектуре трансформер. Нейросеть получает на вход эмбеддинги выравнивания и парного представления, итеративно улучшает их на все более глубоких слоях, и возвращает 3D-структуру, предсказанную на их основе.

Нейросеть состоит из 2 ключевых суб-модулей:

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

• Структурный модуль нейросети, состоящий из 8 однотипных блоков, предсказывает 3D-структуру белка на основе эмбеддинга нашей последовательности (он вырезается из эмбеддинга полного выравнивания), и эмбеддинга парного представления.

После того, как некая 3D-стукрутра была получена на выходе из нейросети, дополнительно используется пакет OpenMM, который на основании физических методов релаксирует полученную структуру.

Данная процедура повторяется трижды в ходе процесса под называнием «переработка» (recycling).

Механизм внимания и трансформерная архитектура

Как было сказано выше, нейросеть, лежащая в основе AlphaFold2, основана на трансформерной архитектуре и механизме внимания. Я коротко объясню их далее. Смело пропускайте этот раздел, если уже знакомы с этими конструкциями.

Механизм внимания получил широкую популярность в районе 2014–2015 в задачах машинного перевода (к примеру, с английского на французский). На тот момент лучшие архитектуры были основаны на рекуррентных нейросетях (RNN).

Сначала обнаружилось, что добавление к ним механизма внимания улучшает качество перевода. Позднее выяснилось, что от рекуррентной нейросети можно избавиться вовсе, и наслоив несколько однотипных блоков, основанных на механизме внимания, можно добиться примерно тех же результатов, что и с рекуррентными нейросетями, используя в 100 раз меньше вычислительных ресурсов.

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

Архитектуры, основанные только на механизме внимания, называются трансформерами. Блок Evoformer AlphaFold2, как подсказывает его название, является частным случаем трансформера (да и структурный модуль, на самом деле, тоже).

Механизм внимания

Механизм внимания формулируется в терминах нестрогого поиска в базе данных ключ-значение.

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

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

# Suppose that we have a key-value database and implement a fuzzy search in it database = { "Ivanov": 25, "Petrov": 100, "Sidorov": 5 } query = "Pvanov" # note the typo in query def hamming_distance(query: str, key: str) -> int: distance = 0 for i, _ in enumerate(query): if query[i] != key[i]: distance += 1 return distance # construct a vector of similarities between query and keys # similarity=1 if query=key; # similarity=0 if there is nothing in common between query and key similarities = {} for key in database: similarities[key] = (1 - hamming_distance(query, key) / len(query)) print(f"Similarities before softmax: {similarities}") output = 0 for key, value in database.items(): output += similarities[key] * value output

Similarities before softmax: {'Ivanov': 0.8333333333333334, 'Petrov': 0.5, 'Sidorov': 0.0}
70.83333333333334

В данном примере кода содержится одна неточность: вместо вероятности, что под запросом пользователь понимал ключ, мы измеряли просто сходство между запросом и ключом. Таким образом сумма этих мер сходства больше 1, и их нельзя интерпретировать как вероятности.

Поэтому требуется нормализовать эту меру сходства, так, чтобы сумма равнялась 1. Для этого используется функция Softmax:

# However, note that the sum of similarities in previous solution was more than 1.
# We probably want to preserve the sum of similarities equal to 1.
# Hence, apply Softmax to similarities vector to normalize it.
import math


def softmax(weights: dict) -> dict:
    """Transform a dict of weights into a dict with sum of weights normalized to 1."""
    # Calculate statistical sum (softmax denominator)
    zustandssumme = 0
    for key, weight in weights.items():
        zustandssumme += math.exp(weight)
    
    return { key: math.exp(weight)/zustandssumme for key, weight in weights.items() }


similarities = softmax(similarities)
print(f"Similarities after softmax: {similarities}")

output = 0
for key, value in database.items():
    output += similarities[key] * value
    
output

Similarities after softmax: {'Ivanov': 0.4648720549505913, 'Petrov': 0.33309538278287776, 'Sidorov': 0.202032562266531}
45.94150246338521

Softmax — это по сути мультиномиальный аналог сигмоиды. Подобно тому, как функция сигмоиды используется в логистической регрессии чтобы конвертировать значения из произвольного диапазона в диапазон [0,1], так что их можно интерпретировать как вероятности, Softmax делает то же самое в мультиномиальном случае. Другой способ думать о Softmax — это рассмотреть его как обобщение распределения Больцмана из статфизики.

Здесь следует сделать два замечания.

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

Во-вторых, обычно запросы, ключи и значения — это векторы. В качестве меры сходства между векторами можно использовать разные функции, например, косинумное расстояние, но обычно в механизмах внимания используется скалярное произведение (обычно поделенное на нормировочный множитель).

Вот изображение механизма внимания из классической статьи «Attention is all you need paper» (игнорируйте блок «Mask» на этой картинке, он не актуален в данной задаче):

Механизм внимания (Scaled Dot-Product Attention). На вход принимаются матрицы запросов (Q), ключей (K) и значений (V), каждым столбцом которых является один вектор запроса, вектор ключа и вектор значения. Далее блок MatMul вычисляет скалярные произведения между каждой парой (запрос, ключ). Блок Scale нормирует каждое скалярное произведение из технических соображений. Блок Mask опционален и не исопльзуется в AlphaFold2. Блок Softmax применяет Softmax к каждому вектору нормированных скалярных произведений, соответствующих одному запросу и разным ключам, переводя их в такое пространство, что их сумма равна 1. Блок MatMul для каждого вектора-запроса в Q возвращает вектор-взвешенную сумму векторов значений (V) с коэффициентами, равными выходу блока Softmax для данного вектора-запроса.

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

Вот реализация механизма внимания на PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F


def attention(query, key, value, mask=None, dropout=None):
    """Compute 'Scaled Dot Product Attention'.
    
    Mostly stolen from: http://nlp.seas.harvard.edu/2018/04/03/attention.html.
    """

    # MatMul and Scale
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Mask (optional)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Softmax
    p_attn = F.softmax(scores, dim = -1)
    
    # MatMul
    if dropout is not None:
        p_attn = dropout(p_attn)
    
    return torch.matmul(p_attn, value), p_attn

Внимание со множеством голов

Предположим, что ваши запросы, ключи и значения являются идентичными векторами (например, эмбеддингами аминокислот):

Предположим, что первая координата x_{i,1}​ отражает растворимость аминокислоты, вторая координата,  x_{i,2} — ее размер, третья координата x_{i,3}​ — ее положительный заряд, четвертая координата x_{i,4} — отрицательный заряд.

Нам может потребоваться установить различные отношения между аминокислотами. Например, насколько взаимозаменяемы аминокислоты S и D в разных качествах?

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

Давайте попробуем спроецировать данную аминокислоту на признаковое пространство «заряженная большая аминокислота». Наше признаковое пространство будет иметь размерность 2, где первая координата будет соответствовать абсолютной величине заряда, а вторая координата — размеру:

Теперь применим механизм внимания в разрезе этого отношения к 2 аминокислотам. Сходство между ними будет записываться как:

Многогололвое внимание. На вход подаются матрицы запросов (Q), ключей (K) и значений (V). В случае самовнимания все они — одна и та же матрица X. К ним применяется h проекций на признаковые пространства h голов с помощью матриц линейного преобразования W_h (блок Linear), далее каждая голова применяет механизм внимания (блок Scaled Dot-Product Attention), результаты работы каждой головы конкатенируются в единый массив (блок Concat), и он проецируется на исходное пространство еще одним блоком Linear.

После применения механизма внимания, выходы со всех голов конкатенируются в единый вектор (см. блок «Concat» на картинке).

Затем конкатенированный вектор этих выходов проецируются обратно на пространство размерности исходного эмбеддинга (блок «Linear» на картинке).

Получившийся механизм называется блоком многоголового внимания (multi-head attention). Подобные блоки являются базовым элементом нейросети-трансформера. Чем больше подобных блоков мы наслоим, тем более высокоуровневые отношения сможет идентифицировать наша нейросеть. Код на PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadedAttention(nn.Module):
    """Mostly stolen from: http://nlp.seas.harvard.edu/2018/04/03/attention.html."""
    
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        
        assert d_model % h == 0
        
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        """Implements multihead attention"""
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)

        nbatches = query.size(0)
        
        # 1) Do all the linear projections in batch from d_model => h x d_k 
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]
        
        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        
        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)

        return self.linears[-1](x)

Здесь я опускаю некоторые важные технические инженерные моменты, такие как применение LayerNorm и Residual connections между блоками multi-head attention, нормировку Scaled Dot-Product Attention, Positional Encoding и т.д.

Модуль Evoformer

Модуль Evoformer состоит из 48 идентичных блоков, которые принимают эмбеддинги множественного выравнивания и парного представления на вход и возвращают их усовершенствованную версию на выходе.

Модуль Evoformer. (a) общая архитектура; (b) матрица/тензор эмбеддинга парного представления; © неравенства треугольника.

Эвоформер состоит из 3 основных шагов:

1. Эвоформер улучшает эмбеддинг выравнивания вниманием по осям (axial/criss-cross), используя дополнительную информацию, содержащуюся в парном представлении.

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

3. Эвоформер применяет блоки неравенств треугольника к парному представлению, обновленному на предыдущем шаге, чтобы обеспечить согласованность данных.

Разберем каждую из частей эвоформера по отдельности.

Evoformer: улучшение множественного выравнивания

Эвоформер: красной рамкой выделен набор модулей, отвечающих за уточнение эмбеддинга выравнивания.МЕХАНИЗМ ВНИМАНИЯ ПО ОСЯМ В ЭВОФОРМЕРЕ

MSA update part of Evoformer uses an approach, called axial or criss-cross attention. It was suggested in visual transformers circa 2020. In visual transformers attention needs to be applied to each pixel of an image, and using the full 2D image as keys would be computationally inefficient.

A frugal alternative to full 2D attention, is to first attend to all pixels in the same row as the query pixel and then attends all pixels in the same column. Same approach was employed here by evoformer. It first attends to other aminoacid residues in the same sequence (which is called row-wise gated self-attention), and then residues from other sequences in the same column (column-wise gated self-attention).

САМОВНИМАНИЕ С ЗАТВОРОМ ПО ПОЗИЦИЯМ ВЫРАВНИВАНИЯМеханизм самовнимания с затвором по колонкам (позициям) выравнивания в Эвоформере

Считаю нужным обсудить 2 аспекта самовнимания по колонкам выравнивания с механизмом запирания/затвора.

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

Во-вторых, верхняя стрелка на этой схеме соответствует механизму запирания на этой схеме, о котором я еще не говорил. Механизм запирания похож на затвор транзистора — если в векторе затвора «закрыт» некий элемент, это практически обнуляет соответствующий элемент выходного вектора механизма внимания. Затвор будет закрыт, если проекция эмбеддинга соответствующей аминокислоты имеет малую «длину» (длину в кавычках, потому что «вектор» затвора все-таки имеет c каналов, и запирание происходит для каждого канала в отдельности). Технически для запирания используется все та же сигмоидная функция от проекции эмбеддинга позиции выравнивания, которая будет близка к 1, если значение координаты проекции достаточно велико или близка 0, если оно достаточно мало. Соответственно, в первом случае затвор будет открыт, а во втором — открыт. Однако, назначение этого механизма для меня неочевидно, и может быть просто инженерным приемом, для которого известно, что так сеть лучше работает (см. ссылки на литературу в конце статьи).

САМОВНИМАНИЕ С ЗАТВОРОМ ПО ПОСЛЕДОВАТЕЛЬНОСТЯМ ВЫРАВНИВАНИЯМеханизм самовнимания по последовательностям выравнивания с затвором в Эвоформере.

Самовнимание по последовательностям выравнивания обеспечивает обмен информацией между последовательностями выравнивания внутри каждой колонки выравнивания. Этот шаг помогает AlphaFold2 обнаружить, что между последовательностями в данной позиции есть коэволюция, или что позиция консервативна. Кроме того, поскольку именно для первой последовательности выравнивания делается предсказание 3D-структуры белка, через нее обновленная информация о 3D-структуре распространяется на другие последовательности.

В дополнительных материалах к статьей DeepMind приводит визуализацию карт внимания на этом блоке (см. изображение ниже). Ряд (a) на картинке показывает, что для каждой позиции выравнивания имеется какое-то подмножество последовательностей, которое несет основную информацию. Ряд (b) показывает, что после нескольких итераций переработки все последовательности обращают внимание на первую последовательность (для которой предсказывается 3D-структура). Ряд © показывает, что промежуточные слои AF2 постепенно выучивают, по сути, филогенетическое дерево, описывающее эволюционное родство последовательностей (см. правую нижнюю картинку «Hamming distance», треугольную матрицу над главной диагональю которой можно рассматривать как филогенетическое дерево).

Карты внимания механизма самовнимания по последовательностям. Интерпретацию рядов (a)-© см. в тексте статьи. Картинка (d) содержит предсказанные расстояния между парами C-альфа атомов аминокислотных остатков белка. Картинка (e) содержит матрицу расстояний Хэмминга между последовательностями выравнивания, фактически, эквивалентную филогенетическому дереву, описыающему эволюционное родство данных последовательностей.

Приведу интерпретацию этих карт внимания из дополнительных материалов к статье:

In Suppl. Fig. 13 we show a visualization of the attention pattern in the MSA along the columns a h_{sti}hsti​ (line 4 of Algorithm 8). We slice along the last axis ii and display the stst array as heat map.

The original MSA subset shown to the main part of the model is randomly sampled. As such the order is random except for the first row. The first row is special because it contains the target sequence and is recycled in consecutive iterations of the model. Due to the shallow MSA of this protein, the random subset leads to a random permutation of the sequences. In order to facilitate easier interpretation of the attention patterns here we reorder the attention tensor by using a more suitable order for the MSA. We perform a hierarchical clustering using the Ward method with simple Hamming distance as a metric and use the output to re-index the sequence dimension in the MSA attention tensor. We resort the indices from the hierarchical clustering manually to keep the target sequence in the first row. This manual sorting is done in such a way as to keep the tree structure valid. The Hamming distances between the reordered sequences (see Suppl. Fig. 13e) show a block-like structure quite clearly after the reordering.

The attention pattern in the first layer of the network in the first recycling iteration (e.g. layer 0, head 7 in Suppl. Fig. 13a) is not very informative and is largely averaging as can be seen by looking at the range of the attention weights.

In the same head, but in the last recycling iteration (Suppl. Fig. 13b) we see that all sequences at all positions attend to the first row. Therefore this head behaves differently upon recycling and is presumably important for distribution the information in the recycled first row to the rest of the MSA representation. Layer 6, head 6 (Suppl. Fig. 13c) shows a pattern that is fairly common in the column-wise MSA attention, here the pattern only varies lightly as one goes along the sequence and there is a clear structure in blocks of sequences that attend to each other. We note that these seem somewhat similar to the blocks derived from hierarchical clustering using Hamming distance. Whether attention patterns provide a good explanation for the behaviour of a given model or are predictive of interventions into the model is a topic of debate in the community, see [127], [128], [129]. A detailed analysis of the generality and predictivity of these attention patterns is beyond the scope of this paper.

ПЕРЕХОДНЫЙ БЛОК ВЫРАВНИВАНИЯ

Последний элемент подблока уточнения эмбеддинга выравнивания в Эвоформере — это переходный блок. Я не вполне понимаю его назначение, и оно не обсуждается подробно в статье и дополнительных материалах. Полагаю, его функция — чисто инженерная.

Переходный блок в Эвоформере.

Эвоформер: уточнение парного представления после улучшения эмбеддинга выравнивания

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

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

Однако, надо помнить, что мы должны взять среднее по всем последовательностям в данной колонке. Поэтому делается что-то подобное, просто математическим представлением этой операции является среднее по последовательностям внешних произведений эмбеддингов колонок. Напоминаю, что операция внешнего произведения (outer product) в линейной алгебре — это частный случай Кронекеровского произведения или тензорного произведения.

Уточнение парного представления после улучшения выравнивания в Эвоформере через внешнее произведение.

Эвоформер: неравенства треугольника как естественные условия для парн

Habrahabr.ru прочитано 29291 раз