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.
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.
bioRxiv
Transformer protein language models are unsupervised structure learners
Unsupervised contact prediction is central to uncovering physical, structural, and functional constraints for protein structure determination and design. For decades, the predominant approach has been to infer evolutionary constraints from a set of related…
Ещё пробуют ансамблирование по нескольким последовательностям из MSA, качество улучшается.
Снова (как и в предыдущей работе) показали, что перплексия хорошо коррелирует с точностью предсказания контактов (и снова видно, что есть куда скейлить).
Логрегрессия выделила 102 из 660 голов, которые помогают предсказаниям. Две головы оказались в top-10 по предсказаниям на всех масштабах. Одна голова оказалась особенно хороша для предсказания локальных контактов и мешала предсказывать дальние. И в целом видна специализация по типам контактов.
Ну и наконец показывают, что вероятности, которые выдаёт логрегрессия, близки к реальным вероятностям контактов, модель неплохо откалибрована и не overconfident. А в кейсах, где она ошибочно предсказывает контакт, он часто оказывается в пределах нескольких аминокислот от предсказанной точки, либо оказывается контактом между двумя цепочками гомодимера. К вставкам модель устойчива, можно вставить в начало/конец 256 аланинов (или 64 в середину) и модель не особенно деградирует. Более раннему трансформеру TAPE хватало 8 аланинов в начале, чтобы качество сильно ухудшилось (возможно, потому что тот обучался на белковых доменах).
Работа взята постером на ICLR 2021 (https://openreview.net/forum?id=fylclEqgvgd).
Снова (как и в предыдущей работе) показали, что перплексия хорошо коррелирует с точностью предсказания контактов (и снова видно, что есть куда скейлить).
Логрегрессия выделила 102 из 660 голов, которые помогают предсказаниям. Две головы оказались в top-10 по предсказаниям на всех масштабах. Одна голова оказалась особенно хороша для предсказания локальных контактов и мешала предсказывать дальние. И в целом видна специализация по типам контактов.
Ну и наконец показывают, что вероятности, которые выдаёт логрегрессия, близки к реальным вероятностям контактов, модель неплохо откалибрована и не overconfident. А в кейсах, где она ошибочно предсказывает контакт, он часто оказывается в пределах нескольких аминокислот от предсказанной точки, либо оказывается контактом между двумя цепочками гомодимера. К вставкам модель устойчива, можно вставить в начало/конец 256 аланинов (или 64 в середину) и модель не особенно деградирует. Более раннему трансформеру TAPE хватало 8 аланинов в начале, чтобы качество сильно ухудшилось (возможно, потому что тот обучался на белковых доменах).
Работа взята постером на ICLR 2021 (https://openreview.net/forum?id=fylclEqgvgd).
OpenReview
Transformer protein language models are unsupervised structure...
Unsupervised contact prediction is central to uncovering physical, structural, and functional constraints for protein structure determination and design. For decades, the predominant approach has...
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.
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.
bioRxiv
MSA Transformer
Unsupervised protein language models trained across millions of diverse sequences learn structure and function of proteins. Protein language models studied to date have been trained to perform inference from individual sequences. The longstanding approach…
На supervised contact prediction (на фичах MSA transformer обучают resnet) побили в том числе более качественную trRosetta_full.
Проверились на задаче предсказания вторичной структуры на CB513, обошли профили HMM, ESM-1b и Netsurf.
Поэкспериментировали с разнообразием последовательностей в MSA. Стратегии, максимизирующие разнообразие, дают качество выше ESM-1b с всего 16 последовательностями на входе. И добавление всего одной последовательности в высоким diversity перекрывает 31 последовательность с низким.
Головы внимания к строкам любят столбцы с высокой изменчивостью. А внимание по столбцам больше смотрит на более информативные последовательности. Также логрегрессия нашла 55 (из 144) голов, предсказывающих контакты.
Такие дела. Скейлить по идее есть куда, ждём продолжения.
Проверились на задаче предсказания вторичной структуры на CB513, обошли профили HMM, ESM-1b и Netsurf.
Поэкспериментировали с разнообразием последовательностей в MSA. Стратегии, максимизирующие разнообразие, дают качество выше ESM-1b с всего 16 последовательностями на входе. И добавление всего одной последовательности в высоким diversity перекрывает 31 последовательность с низким.
Головы внимания к строкам любят столбцы с высокой изменчивостью. А внимание по столбцам больше смотрит на более информативные последовательности. Также логрегрессия нашла 55 (из 144) голов, предсказывающих контакты.
Такие дела. Скейлить по идее есть куда, ждём продолжения.