Transformer protein language models are unsupervised structure learners
Roshan Rao, Joshua Meier, Tom Sercu, Sergey Ovchinnikov, Alexander Rives
Статья: https://www.biorxiv.org/content/10.1101/2020.12.15.422761v1
Модель: https://github.com/facebookresearch/esm
Продолжение темы про большую языковую модель для белков ESM-1b частично от тех же авторов. Название перекликается с известными статьями из NLP/CV, например, про GPT-2 (“Language Models are Unsupervised Multitask Learners”), GPT-3 (“Language Models are Few-Shot Learners”) или SimCLRv2 (“Big Self-Supervised Models are Strong Semi-Supervised Learners”).
В данной работе фокусируются на задаче предсказания контактов в белке с целью заменить текущие многоступенчатые пайплайны на один forward pass уже предобученной сетки (та самая ESM-1b на 650М параметров). И показывают, что все необходимые для этого фичи уже содержатся в картах внимания определённых голов self-attention’а.
Предсказание контактов (contact prediction) бывает как водится supervised и unsupervised.
В supervised части уже было много интересных результатов, в том числе с обучением глубоких residual networks на структурах белков.
Unsupervised предсказание контактов обучается на последовательностях без какой-либо информации о структуре белка. Широко используемый для этого метод — это обучение Potts model (специальный вариант Markov Random Field) на множественном выравнивании (Multiple sequence alignment, MSA — это когда множество эволюционно связанных последовательностей выравнивается между собой по позициям, делая вставки, делеции и замены где надо, так чтобы минимизировать метрику типа edit distance по всем последовательностям).
Фичи из предсказания контактов могут потом использоваться для предсказания структуры (AlphaFold использовал фичи из Potts model).
Языковые модели уже использовались для предсказания контактов. Это было и в оригинальной работе про ESM-1b, и в других работах на трансформерах и LSTM, в том числе с supervised обучением предобученной модели. В данной работе показывают, что SoTA unsupervised предсказатель контактов можно напрямую извлечь из карт внимания предобученного трансформера (того же ESM-1b).
В работе сравниваются с тулом Gremlin, который принимает на вход MSA (и от качества этого выравнивания сильно зависит и качество результата). Gremlin берётся в двух вариантах: 1) на всех MSA из обучающей выборки trRosetta (которые сгенерены по Uniref100 и куда добавлены метагеномные последовательности, если глубина выравнивания по Uniref100 оказалась мала), и 2) на MSA полученных через Jackhmmer на том же датасете, на котором обучался ESM-1b (подмножество Uniref50). MSA из trRosetta называются ключевым ингредиентом state-of-the-art пайплайнов.
Что собственно делают? Последовательность белка (без какого либо MSA) прогоняют через ESM-1b, получают в трансформере карты внимания (attention maps, то есть скоры, которые получаются внутри трансформера при перемножении матриц Q и K). Эти карты отправляются независимо для каждой пары аминокислот (i,j) в L1 логрегрессию, которую обучают на каком-то небольшом числе последовательностей (например, 20). Это обучение нужно лишь для того, чтобы найти веса логрегрессии и понять, какие головы каких слоёв трансформера важны, сам трансформер никак не обновляется. Также пробуют обучать отдельные модели для предсказания близких (6-12 аминокислот), средних (12-24) и далёких контактов (24+).
Результат? На 20 обучающих последовательностях ESM-1b даёт на всех дистанциях более высокую точность, чем Gremlin (а также предыдущие языковые модели типа TAPE или ProtBERT-BFD).
Интересные абляции.
Всего с одной отобранной головой (из 660), качество ESM-1b сравнимо с Gremlin, а при усреднении 10 голов приближается к Gremlin’у на оптимизированном MSA.
Одного обучающего примера достаточно, чтобы получить статистически неотличимый от Gremlin’а результат. 10 примеров достаточно, чтобы его обойти.
Также вместо структур пробуют обучать по MSA (обучают Gremlin, берут оттуда top-L couplings и считают их настоящими контактами), получается похуже, но сравнимо с Gremlin.
