AlphaFold2 gibi modellerin ortaya çıkışı, moleküler yapay zeka modellerinin daha hızlı çıkarım ve eğitimine olan talebi artırdı. Hız ihtiyacı, benzersiz hesaplama zorluklarıyla birlikte gelir; bu zorluklar arasında algoritmik karmaşıklık, bellek verimliliği ve sıkı doğruluk gereksinimleri bulunur. Bu sorunları çözmek için NVIDIA, ortaklarıyla işbirliği yaparak, daha hızlı eşitlikçi işlemler ve daha hızlı MSA üretimi gibi hızlandırılmış çözümler sunmuştur.
Bugün, moleküler yapay zeka modellerinin eğitimini ve çıkarımını hızlandırmak için cuEquivariance kütüphanesinde yeni çekirdekler ve bir NVIDIA NIM mikroservisi sunduk. Bu hızlandırmalar, MIT ve Recursion tarafından geliştirilen açık kaynak bir temel model olan Boltz-2 gibi moleküler yapay zeka sistemlerinin daha sofistike bir şekilde geliştirilmesine ve moleküler yapıların ölçeklenebilir şekilde daha hızlı elde edilmesine olanak tanır.
NVIDIA cuEquivariance ile Gelecek Nesil Protein Yapı Modellerini Hızlandırma
NVIDIA cuEquivariance, MACE, Allegro, NequIP ve DiffDock gibi geometriye duyarlı sinir ağlarının yoğun hesaplamalarını hızlandırmak için tasarlanmış bir CUDA-X kütüphanesidir. Segmented Tensor Products gibi temel eşitlikçi işlemleri önemli ölçüde hızlandıran yüksek doğrulukta CUDA çekirdekleri ve kapsamlı API’ler sunar.
cuEquivariance v0.5 sürümüyle birlikte kütüphanede, Nobel ödüllü protein yapı tahmin modelleri gibi AlphaFold2 için kritik olan hızlandırılmış Üçgen Dikkat ve Üçgen Çarpma çekirdekleri bulunmaktadır. Hızlandırılmış üçgen işlemleri eklenerek, cuEquivariance’ın etkisi protein katlanması, RNA/DNA bağlanması, blind docking, protein kompleks tahmini ve afinitet puanlama gibi uygulamalara genişlemektedir.
Bir proteinin üç boyutlu yapısını anlamak kritik öneme sahiptir çünkü bu yapı, proteinlerin nasıl çalıştığını ortaya koyar. Ancak, hücrelerin ve yaşamın karmaşıklığı, biyolojik kompleksler içindeki dinamik etkileşimlerden kaynaklanmaktadır. Bu kompleksler sadece proteinlerden oluşmaz; çeşitli küçük moleküller de dahil olmak üzere, proteinler, nükleik asitler (DNA ve RNA gibi), lipidler ve karbonhidratların karmaşık birleşimleriyle çalışmaktadırlar.
Üçgen İşlemlerin Hesaplama Zorluğu
Bu bireysel moleküllerin ve komplekslerin yapısını ve dinamik davranışlarını tahmin etmek, moleküler yapay zekanın bir sonraki sınırını temsil etmektedir. Bu bilimsel atılım, hücresel yolları aydınlatabilir, hastalık mekanizmalarını tanımlayabilir ve belirli moleküler etkileşimleri hedefleyen ilaçların tasarımını sağlayabilir.
Proteinler, RNA ve DNA, tekrar eden birimlerden oluşan uzun moleküllerdir: DNA ve RNA için nükleotid, proteinler için ise amino asitler. Bu diziler hücrelerde üretildiğinde, yapı taşlarının belirli dizilimi, karmaşık üç boyutlu yapılar haline katlanmalarını sağlar. Bu 3D şekiller, moleküllerin işlevini ve diğer hücresel bileşenlerle olan etkileşimlerini belirlemektedir.
State-of-the-art geometriye duyarlı sinir ağlarında (AlphaFold3, Proteina, Chai-1, Neo-1 ve Boltz-2 gibi) iki temel, ancak hesaplama açısından yoğun işlemler vardır: üçgen çarpma ve üçgen dikkati. Bu işlemler genellikle bu modellerde en fazla zaman alan bileşenler arasında yer almaktadır.
Pearwise dikkate alma mekanizmaları, Transformers ile popüler hale gelmiştir. Bir tokenın (Transformer kelime hazinesindeki bir yapı taşı) her diğer token ile olan ilişkisini hesaplayarak, modelin bir kelimenin, tüm diğer kelimelerle bağlam dışındaki önemini anlamasını sağlamakta kullanılmaktadır.
Moleküler yapay zeka modelleri 2D temsili temel alarak 3D yapı tahmini yapmaya çalıştığı için, pearwise ilişkiler tüm bağlamsal bilgiyi sağlamamaktadır. İşte burada “Üçgen İlişkiler”, mekansal ilişkileri yakalamak için güçlü bir proxy işlevi görebilmektedir. Örneğin, bir yapı taşı ik‘ya yakınsa ve kj‘ye yakınsa, i ve j‘nin mekansal olarak ilişkili olma olasılığı yüksektir, dolayısıyla direkt bir pearwise sinyal olmadan bile bu ilişki göz önüne alınabilir.
Bir molekül N yapı taşı içeriyorsa, bu işlemler naif bir yaklaşımla O(N3) karmaşıklığı sergilemektedir. Bu hesaplama yoğunluğu, büyük moleküller ve karmaşık çok moleküllü yapılar için önemli bir zorluk teşkil eder; bu da önemli hesaplama maliyetlerine ve bir AI modelinin ölçeklenme potansiyeli üzerinde sert kısıtlamalara yol açar.
cuEquivariance ile Üçgen İşlemlerin Hızlandırılması
Burada, cuEquivariance ön uç Üçgen Dikkat modülünün performansını, basit bir PyTorch uygulaması ile karşılaştırıyoruz. Bu ölçüm, tam uçtan uca hızlanma yerine yalnızca modül çalışma süresini tartışır. İlerleyen bölümlerde uçtan uca performans ölçümlerinden bahsedeceğiz.
cuEquivariance Üçgen Dikkat çekirdekleri BF16 hassasiyeti ile, çıkarım için, PyTorch uygulamasına göre 5 kat çekirdek düzeyi hız artışı ve bellek kullanımını O(N3‘ten O(N2)’ye düşürdüğü gözlemlenmiştir. Ayrıca, cuEquivariance Üçgen Dikkat Ön Çekirdeklerini Trifast ile karşılaştırdığımızda 1.5x ile 2x arasında çekirdek düzeyinde bir hız artışı gözlemlenmiştir.
Benzer şekilde, cuEquivariance Üçgen Çarpma çekirdekleri BF16 hassasiyeti ile, herhangi bir doğruluk kaybı olmadan 5x modül düzeyi hız kazancı sağlar.
“Bu çekirdekler, beklenen zaman diliminin ötesinde ve Boltz ailesinin bir parçası olacak, hız ve bellek tüketimindeki darboğazları gidermeye yardımcı olacaktır” dedi MIT araştırmacısı Gabriele Corso.
Boltz-1x, Boltz-1‘in bir sonraki nesil versiyonu üzerinde, PyTorch, Trifast ve cuEquivariance tabanlı uygulamaların çeşitli hassasiyetleri (TF32, FP32 ve BF16) ile karşılaştırmalar yaptık ve uçtan uca çıkarım sürelerini değerlendirdik. Bu çalışmalarda, Boltz-1x yazarları tarafından yayımlanan varsayılan test veri seti kullanılmıştır. Aynı hassasiyetlerle PyTorch BF16’dan cuEquivariance BF16’a geçtiğimizde 1.75 kat performans artışı gözlemlendi. PyTorch FP32’dan cuEquivariance BF16’a geçişte ise 2.5 kat performans artışı sağlandı.
Boltz-1x ile PyTorch FP32’den cuEquivariance BF16 konfigürasyonları kullanarak, uçtan uca eğitimde 1.35 kat hız artışları görüyoruz. Uçtan uca hız artışları, modelin mimarisine bağlı olarak değişiklik gösterebilir.
“cuEquivariance bu genişleme, bizim için son derece değerlidir, eğitim sürelerinde 2 kat, çıkarım sürelerinde ise 3 kat hız artışı gördük; bu da model tekrarlama döngülerini önemli ölçüde azaltıyor ve daha büyük moleküller üzerinde çıkarım yapmamıza olanak tanıyor” dedi VantAI CTO’su Luca Naef.
cuEquivariance tarafından sağlanan bu hızlandırmalar, MIT, VantAI, Molecular Glue Labs (MGL), Dyno, Peptone, Genesis ve Xaira gibi birçok ortak tarafından beğenilmiştir. Bu ortaklar, erken versiyonları test etmiş ve geri bildirim sağlamışlardır. Topluluğun bu hızlandırmalardan yararlanmasını ve işimizi geliştirmemize yardımcı olacak kritik geri bildirimleri sağlamasını dört gözle bekliyoruz; bu sayede bilimsel yeniliklerin sınırlarını daha da ileriye taşıyabileceğiz.
Boltz-2 NIM ile Dijital Biyoloji için Kurumsal Düzeyde İşbirliği
Boltz-1 gibi modellerin başarısından yola çıkarak, MIT’deki Boltz ekibi ile Recursion işbirliği ile geliştirilen bir sonraki nesil Boltz-2 modeli, önemli bir adım atmaktadır. Boltz-2, daha büyük ve daha yetenekli bir model olarak tasarlanmıştır; çıkarım zamanı optimizasyonları Boltz-1x ile birleştirilmiş ve benzersiz, en son teknolojiye sahip afinitet tahmin yetenekleri içermektedir. Bu devrim niteliğindeki modelin en erişilebilir versiyonunu sunmak için NVIDIA, Boltz-2’yi bir NVIDIA NIM olarak paketlemektedir.
NIM’ler, optimize edilmiş, üretime hazır çıkarım sunan kullanımı kolay, önceden oluşturulmuş konteynerlerdir. Boltz-2 NIM, araştırmacılara ve geliştiricilere güçlü yeteneklerine kolay erişim sunarak, gerçek zamanlı tahminler yapmalarını ve zorlu ilaç keşif iş akışları için test zamanı ölçeklenmesini sağlar. Bu yaklaşım, öncü moleküler yapay zekaya erişimi demokratikleştirerek, daha geniş bir kullanıcı yelpazesinin Boltz-2’nin tahmin gücünden yararlanmasını mümkün kılar.
Moleküler Yapay Zeka için Hızlandırılmış Hesaplama
cuEquivariance ile sağlanan gelişmiş hesaplama verimliliği oldukça önemlidir. Eğitim için, bu hızlandırılmış çekirdekler, araştırmacıların önceden eğitim ölçekleme yasalarından daha büyük temel modeller oluşturmalarına yardımcı olur; burada artan hesaplama verimliliği genelde model performansında iyileşme ile ilişkilidir. Ayrıca, hesaplama süresi ve maliyetindeki elde edilen verimlilikler, daha fazla model geliştirme döngüsü için kaynakları serbest bırakmaktadır ve böylece bir sonraki nesil kapasitelerin sınırlarını daha da ileriye taşır. Test aşamasında, sağlanan hızlandırmalar, daha geniş in silico deneylerin gerçekleştirilmesini kolaylaştırmakta; bu, sanal tarama kampanyalarının yüz binlerce veya hatta milyonlarca çıkarımla ölçeklenmesine olanak tanımaktadır.
“NVIDIA’nın cuEquivariance kütüphanesi, Boltz-2 gibi yapısal olarak duyarlı biyomoleküler modeller için temel hızlandırmalar sağlıyor” dedi Recursion CTO’su Ben Mabey. “Hesaplama darboğazlarını aşarak, ilaç keşfi için bu güçlü modellerin endüstride daha hızlı Ar-Ge döngülerinin sağlanmasına yardımcı olacaktır.”
Kütüphane düzeyindeki bu hızlandırmalara ek olarak, NVIDIA, NVIDIA NIM mikroservisleri gibi ürünlerle en son teknolojilerdeki modellere erişimi artırmaktadır. Örneğin, Boltz-2’yi bir NIM olarak paketlemek, araştırmacılara ve geliştiricilere bu güçlü yeteneklerin uygulanması için hatasız, üretim için hazır bir çözüm sunmaktadır. Zorlu ilaç keşif iş akışlarının verimli bir şekilde ölçeklenmesine olanak tanımaktadır.
Geniş bilim topluluğu ile işbirliği yaparak, NVIDIA, cuEquivariance kütüphanesi gibi temel yazılımları geliştirip rafine etmekte ve NIM aracılığıyla optimize edilmiş model erişimi sunmaktadır. Bu teklifler, NVIDIA’nın hesaplama platformu ile desteklenmekte ve bilim insanlarını araştırmaların sınırlarını zorlamaya, hesaplama içgörülerinden gerçek dünya etkisine geçişi hızlandırmaya olanak tanımaktadır.
cuEquivariance’ı Bugün Deneyin
Bu hızlandırmalar, Apache 2.0 lisansı altında PyTorch API ön ucu ile günümüzde mevcut. Hızlandırmalar, desteklenen hassasiyetler ve örnekler hakkında daha fazla bilgi için buraya tıklayın.
