基于太初AI加速卡，自研 DNA 語言模型 MxDNAPro

MxDNAPro 模型是一個(gè)專門閱讀理解DNA語言的人工智能大模型，能夠深度解讀 DNA：快速精準(zhǔn)地找出 DNA 最可能導(dǎo)致疾病的遺傳變異、預(yù)測出哪些DNA區(qū)域或分子可能成為治療疾病的新靶點(diǎn)。

該自研 MxDNAPro 模型在多項(xiàng)基因功能預(yù)測任務(wù)中表現(xiàn)優(yōu)異，超越主流模型，并能夠顯著提升計(jì)算效率，為復(fù)雜基因序列精準(zhǔn)建模和DNA語言潛在規(guī)律探索開辟了新范式。

MxDNAPro 模型首創(chuàng)自主學(xué)習(xí)劃分token策略，通過可變形卷積與稀疏專家模塊，讓模型自行識別生物學(xué)意義片段；并引入交叉注意力（Cross-Attention）機(jī)制，實(shí)現(xiàn)任意輸入/輸出長度的精準(zhǔn)對齊；同時(shí)基于八張?zhí)魽I加速卡的多卡并行訓(xùn)練和 FlashAttention 算子庫，提升訓(xùn)練吞吐量、縮短訓(xùn)練時(shí)間。

MxDNAPro 模型的平均性能，相較 DNABERT-2 模型，AUC 提高約 1.9%；組蛋白標(biāo)記預(yù)測，相較 DNABERT-2 模型，AUC 提高約 3.4%；并形成一套可復(fù)用、可加速、可解釋的 DNA 語言模型訓(xùn)練框架。

SNP 是基因組中最普遍且與復(fù)雜遺傳疾病關(guān)聯(lián)最緊密的變異類型。SNPBERT 模型開創(chuàng)性地借鑒自然語言處理思想，將 DNA 序列視作“語言”，利用 Transformer 架構(gòu)精準(zhǔn)捕捉單堿基的復(fù)雜上下文關(guān)系。

該自研 SNPBERT 模型實(shí)現(xiàn)了高精度的基因型填充（imputation），大幅提升了下游疾病關(guān)聯(lián)研究與功能預(yù)測的準(zhǔn)確性，為解鎖疾病機(jī)制和推動(dòng)個(gè)性化健康管理提供關(guān)鍵洞見。

SNPBERT 模型通過 Byte Pair Encoding 的 token 詞匯本進(jìn)行編碼，并在此基礎(chǔ)上采用 RoPE（Rotary Position Embedding）對 token 位置進(jìn)行編碼；基于 Transformer 框架，創(chuàng)新性地引入父系/母系雙掩碼自監(jiān)督預(yù)訓(xùn)練，實(shí)現(xiàn)等位基因級的并行堿基預(yù)測；并分別編碼父系與母系兩條單倍型序列，在同一位置同步掩碼并預(yù)測堿基，解碼后形成精準(zhǔn)的基因型推斷；同時(shí)基于八張?zhí)魽I加速卡的多卡并行訓(xùn)練和 FlashAttention 算子庫，提升模型訓(xùn)練速度。

SNPBERT 模型訓(xùn)練時(shí)間縮短 50%；整體吞吐量提升 80%；并借助太初元碁的 FP16/FP8 混合精度訓(xùn)練，GPU 顯存占用降低約 30%；batch size 從 128 增大至 256。

本次合作，湖南大學(xué)的合作者均表示：

● 國產(chǎn)卡體驗(yàn)好，算子深度適配：在高性能計(jì)算方面的知識得到了顯著積累，之前一直是使用 Nvidia 的計(jì)算卡，現(xiàn)在體驗(yàn)到國產(chǎn)卡也非常不錯(cuò)。太初AI加速卡同樣支持 flash-attention 的加速訓(xùn)練算子，這點(diǎn)非常給力。

● 性能強(qiáng)勁，遷移成本低：太初AI加速卡性能強(qiáng)勁，多卡并行訓(xùn)練 Translatomer 模型，極大地縮短了訓(xùn)練時(shí)長。同時(shí)核糖體圖譜預(yù)測模型中主要使用的 PyTorch Lightning 庫，太初AI加速卡也均支持，復(fù)現(xiàn)效率高。

● 教程詳細(xì)、易于上手：太初元碁提供的教程詳盡且易于上手。這些教程不僅系統(tǒng)地介紹了 SDAA 的基本概念，還通過實(shí)例引導(dǎo)我快速實(shí)踐，極大地降低了學(xué)習(xí)門檻，使我在短時(shí)間內(nèi)能掌握核心知識。

● 工程師專業(yè)素養(yǎng)過硬，響應(yīng)迅速：代碼開發(fā)過程中遇到的問題，太初元碁的工程師們能夠在生態(tài)支持群中快速響應(yīng)，第一時(shí)間給出解決方案，非常高效。