单细胞组学大模型（2）--- scBERT，已开源详细代码，且有预训练模型权重，可自行DIY！

最新推荐文章于 2024-09-02 10:06:40 发布

组学之心

最新推荐文章于 2024-09-02 10:06:40 发布

阅读量1.9k

点赞数 26

分类专栏：单细胞转录组大模型文章标签： scBERT 单细胞转录组测序

本文链接：https://blog.csdn.net/weixin_56751316/article/details/141607873

版权

单细胞转录组大模型专栏收录该内容

12 篇文章

订阅专栏

–https://doi.org/10.1038/s42256-022-00534-z

scBERT as a large-scale pretrained deep language model for cell type annotation of single-cell RNA-seq data

scBERT模型是2022年腾讯的Healthcare AI实验室研发的，它和基础BERT模型架构相似。

和iSEEEK模型相比，它的策略是用上基因表达的数据，并把transformer块换成performer块。接下来仔细的看看它的设计策略，和关键部分的代码怎么实现的。

留意更多内容，欢迎关注微信公众号：组学之心

数据和处理简单描述：

数据：1M / cross-tissue, human；来自PanglaoDB数据库，文章列出了具体的使用数据

数据处理：统一基因symbol，去除不匹配的基因和重复的基因。然后使用scanpy中的sc.pp.normalize_total和sc.pp.log1p方法对数据进行归一化（所以输入的是归一化后的）。

代码地址：https://github.com/TencentAILabHealthcare/scBERT

1.scBERT算法

1.1 模型架构概览

A：（上半部分）对未标记的数据进行自监督学习，并对特定于任务的数据进行微调。将掩码表达嵌入和基因嵌入组合后input到 Performer encoder中（⊕ 表示加法）。reconstructor 用于生成掩码基因的outputs，用于计算重建损失。

（下半部分）在监督微调阶段，将特定于任务的 scRNA-seq 数据的掩码表达嵌入和基因嵌入组合后，输入到预训练encoder中。输出表示传递到一维卷积层和分类器，来生成细胞类型预测。

Performer encoder是在预训练和微调阶段使用的模型之间共享的组件。在预训练和微调过程中，reconstructor和分类器独立且单独地用于模型。

B：scBERT 的嵌入图示。首先将预处理后的 scRNA-seq 数据转换为离散化表达，然后随机掩码。比如，将EG1 基因嵌入（来自 gene2vec 的基因同一性落入第一个 bin）和EB2 表达嵌入（落入第二个 bin 并转化为与 EG1 相同的维度）相加并传入到 Performer中以生成基因的表示。

1.2 Embedding表示

单细胞数据的掩码表达embeddings：

原始BERT 的嵌入包括 token 和 position，其中token embeddings是一个离散变量（代表一个单词，例如iSEEEK模型），而scBERT模型的原始表达输入是代表单个细胞中基因表达的连续变量。

scBERT利用 NLP 字段中的bag-of-words技术来对基因的表达进行分箱（binning），从而将它们转换为离散值（先解肢再装袋），并将它们转换为 200 维向量。

等大小分箱：整个表达值的范围被平均划分为若干个等宽的区间，然后再用整数编号来表示区间。
bag-of-words：将文本表示为一组词汇及其出现频率，而不考虑词汇之间的顺序或语法结构。

bag-of-words具有以下特征：

无序性：忽略了词汇在文本中出现的顺序，只关注哪些词汇出现以及它们出现的频率。
词汇表：文本中的所有独特词汇组成了一个词汇表，即每个文档中的所有不同的词汇。
词频：记录每个词汇在文档中出现的次数。

所以，依照Bag-of-Words的思想，细胞相当于Bag-of-Words中的文档，基因相当于词汇，而分箱后的表达值相当于词频，可以认为是每个细胞中的基因转录频率。

预训练模型Gene2vec的基因embeddings：

改变输入的序列是不会改变其含义的，因此绝对位置对基因没有意义。而从 gene2vec 获得基因嵌入来表示基因同一性（每个基因都有一个独特的 gene2vec 嵌入），可以将其视为相对嵌入，从一般共表达层面来捕获语义相似性。

Gene2vec模型是2019年的一项研究，发表在《BMC Genomics》。https://doi.org/10.1186/s12864-018-5370-x

Gene2vec是从GEO 数据库的中获取984 个人类芯片测序数据，用PCC挑选出每个数据集中基因表达相关性>0.9的基因对来表示基因共表达，结合基因功能注释数据库，得到具有功能注释的270,704 对，涉及 5369 个基因；负数据集作为基因对在功能上无关联的集合，涉及 12,521 个（19,307 个中的 64.85%）人类基因，有 40,879,714 个基因对。

用word2vec来生成基因嵌入。在基因嵌入中，提取所有共同表达的基因对，并最大化每对基因中一个基因在另一个基因存在下的概率，相当于 skip-gram 模型。

该模型训练了一个 200 维的向量表示，用于表示所有人类基因。

作者直接从该模型中拿到每个基因的训练好的基因向量表示作为嵌入，与单细胞数据掩码嵌入相加。

1.3 Performer特征学习

具有大感受野的 Transformer 可以有效地利用 scRNA-seq 数据中的全局信息，并通过无偏倚地捕获长程基因-基因相互作用来学习每个细胞的全面全局表示。由于计算复杂性，Transformer 的输入序列长度限制为 512 个，而大多数 scRNA-seq 数据包含超过 10,000 个基因。

scBERT用 Performer 替换了 BERT 中使用的 Transformer 编码器，允许超过 16,000 个基因输入。Performer保留了完整的基因水平解释，放弃了 HVG 和降维的使用，能够学习判别性基因和有用的相互作用。因此，scBERT 允许以无偏倚的数据驱动方式发现基因表达模式和细胞类型注释的长距离依赖性。scBERT 不是严重依赖超参数选择。

Performer则不依赖于注意力的稀疏性和低秩矩阵，也能达到线性的效果。它使用了FAVOR+的方法。

attention的计算公式其实就是 $softmax(QKT/sqrt(d_k))*V$ ，而在Performer中，采用了另一种等价的表达形式，公式在下面，其中diag是将一个输入向量变成对角矩阵， $1_L$ 则是一个长度为L的全为1的向量。这里的不同点就是把 $so f t ma x$ 进行了拆分，其中， $diag(A{1_L})$ 是为了求和， $D^{-1}$ 则是将这些和变成倒数，所以 $D^{-1} A$ 跟 $softmax(QKT/sqrt(d_k))$ 是等价的。

自注意力机制的核心是计算注意力矩阵，而这个矩阵的计算通常涉及输入序列的所有元素之间的点积操作，导致计算复杂度较高。为了降低这种计算复杂度，在Performer模型中采用了一种基于核函数的方法近似注意力矩阵A。

$K (x, y)$ ：这个函数是输入向量 𝑥 和 𝑦 的核函数，实际上它代表了在特征空间中 𝑥 和 𝑦 的相似度。
𝜙(𝑥)：这是一个映射函数，它将输入向量 𝑥 映射到一个高维空间，通常称为特征空间。
𝐸[⋅]：这个符号表示对随机变量取期望值。在Performer中，使用随机特征映射的技术来估计点积核的期望值，这样可以通过样本估计的方式来近似原始的核函数。

所以整理一下，Performer注意力矩阵计算如下：

$Q' = \emptyset （ Q ）$ ， $K' = \emptyset （ K ）$ ， $\emptyset （ x ）$ 定义为： $c$ 是常数， $ω$ 是随机特征矩阵， $m$ 是矩阵维度。

Performer能够将原本计算复杂度为 𝑂(L2d) 的全局自注意力机制降低到 𝑂(Lrd)，极大地提高了计算效率，同时又不会显著降低模型的表现。

1.4 在没有标签数据上的自监督学习

scBERT 遵循 BERT 模型在 NLP 任务中的常规自学习策略，随机掩码输入数据值并根据剩余输入进行预测。

考虑到 dropout zeroes 现象，在随机掩码非零基因表达后，使用剩余基因通过模型预测重建原始输入。并利用交叉熵损失作为reconstructor（全连接层+softmax）的损失：

$L_{Rec}$ 中 $M$ 是细胞数， $N$ 是被屏蔽的基因表达值数； $y_{i,j}$ 和 $p_{i,j}$ 分别是基因 $j$ 在细胞 $i$ 中的真实表达和预测表达，对于每个被掩码的基因表达值，交叉熵损失会计算它们之间的差异，并最小化它。 $L o g$ 部分是模型预测基因表达的对数概率。损失的每一项 $y_{i,j}log(p_{i,j})$ 反映了真实值和预测值之间的不一致性。因为是最小化损失，所以加了个负号。

1.5 特定任务的监督学习

scBERT 的输出是对应于每个基因的 200 维特征，并对每个基因特征应用1D卷积进行信息提取。然后将三层神经网络作为分类头，并将基因特征转化为每种细胞类型的概率。交叉熵损失也被用作细胞类型标签预测损失:

公式中 $z_{i}$ 和 $q_i$ 分别表示细胞 $i$ 的真实细胞类型标签和预测标签。

2.下游任务

2.1 评估数据集内的细胞类型注释的稳健性

在多个数据集用多种聚类降维的方法来测评细胞分类的准确性和F1值

2.2 跨队列和器官的细胞注释

2.3 发现新细胞类型

通过在 scBERT 训练过程中去除 α-β T 细胞、γ-δ T 细胞、成熟 B 细胞和浆细胞群，scBERT 在人类肝组织的 MacParland 数据集上的性能。新细胞类型和已知细胞类型的准确度和 F1 分数都显示在箱形图中。

左图是scBERT 为 MacParland 细胞类型提供的置信度分数;对于所有已知细胞类型，模型预测概率低（< 0.5）的细胞被指定为潜在的新细胞类型。桑基图将已知和新细胞类型的 scBERT 预测与 MacParland 数据集的原始细胞类型注释进行比较，其中浆细胞被标记为新细胞类型，因为它们在 scBERT 训练过程中看不到。

3.模型代码

3.1 预训练

3.1.1 随即掩码实现

# 随机生成和t形状相同的随机掩码矩阵
def prob_mask_like(t, prob):
    return torch.zeros_like(t).float().uniform_(0, 1) < prob

# 检查t的元素是或否和token_ids中的值匹配，匹配则为True，生成布尔掩码矩阵，用于标记哪里不需要被掩码
# 通常是特殊符号的位置
def mask_with_tokens(t, token_ids):
    init_no_mask = torch.full_like(t, False, dtype=torch.bool)
    mask = reduce(lambda acc, el: acc | (t == el), token_ids, init_no_mask)
    return mask

def get_mask_subset_with_prob(mask, prob):
    batch, seq_len, device = *mask.shape, mask.device
    max_masked = math.ceil(prob * seq_len) # 计算每个序列中要遮掩的最大标记数     
    num_tokens = mask.sum(dim=-1, keepdim=True) # 计算每个序列中要遮掩的最大标记数
    
    # 1.计算超过允许遮掩的部分
    mask_excess = torch.cat((torch.zeros(0), torch.arange(mask.size(-1)).repeat(mask.size(0)))).reshape(mask.size(0),mask.size(-1)).to(device)
    mask_excess = (mask_excess >= (num_tokens * prob).ceil())        # 只有15%的token需要被mask
    mask_excess = mask_excess[:, :max_masked]       #只保留最大遮掩数量 max_masked 的部分
    
    # 2.生成随机数矩阵并处理特殊标记
    # 生成一个随机数矩阵，并将 mask 中无效的位置（即 False 的位置）填充为 -1e9，确保这些位置不会被选择进行遮掩
    rand = torch.rand((batch, seq_len), device=device).masked_fill(~mask, -1e9)     
    
    # 3.选择概率最大的标记进行遮掩
    #  找出每个序列中随机数最大的前 max_masked 个位置的索引
    _, sampled_indices = rand.topk(max_masked, dim=-1) 
    # 使用 mask_excess 掩盖掉多余的遮掩部分
    sampled_indices = (sampled_indices + 1).masked_fill_(mask_excess, 0)        
    
    # 4.创建新掩码矩阵
    new_mask = torch.zeros((batch, seq_len + 1), device=device)     # 创建一个全零矩阵
    # 根据 sampled_indices 的索引将对应位置设置为 1。根据 sampled_indices 的索引将对应位置设置为 1。
    new_mask.scatter_(-1, sampled_indices, 1)
    
    # 5.返回最终的布尔掩码矩阵，去掉多加的第一个维度，并将矩阵转换为布尔类型，True是要mask的。
    return new_mask[:, 1:].bool()

把上面三个自定义函数调用，获得掩码之后的矩阵输入和pad_token_id填充的标签：

def data_mask(data,
    mask_prob = MASK_PROB,
    replace_prob = REPLACE_PROB,
    num_tokens = None,
    random_token_prob = RANDOM_TOKEN_PROB,
    mask_token_id = MASK_TOKEN_ID,
    pad_token_id = PAD_TOKEN_ID,
    
    # 需要忽略掩码操作的标记 ID 集合，例如 [CLS]、[SEP] 等特殊标记
    mask_ignore_token_ids = MASK_IGNORE_TOKEN_IDS
):
    
    # 1. 设置不被掩盖的标记
    mask_ignore_token_ids = set([*mask_ignore_token_ids, pad_token_id])
    no_mask = mask_with_tokens(data, mask_ignore_token_ids)   # 其中 True 表示这些位置不应被遮掩
    
    # 2. 生成掩码矩阵
    mask = get_mask_subset_with_prob(~no_mask, mask_prob) 
    
    # 3. 创建掩码后的输入
    mask_indices = torch.nonzero(mask, as_tuple=True)
    masked_input = data.clone().detach()
    
    # 4. 随机标记替换
    # 如果 random_token_prob 大于 0，则以一定的概率将部分标记替换为随机标记。
    if random_token_prob > 0:
        assert num_tokens is not None, 'num_tokens keyword must be supplied when instantiating MLM if using random token replacement'
        random_token_prob = prob_mask_like(data, random_token_prob)       # get the mask matrix of random token replace
        random_tokens = torch.randint(0, num_tokens, data.shape, device=data.device)     # generate random token matrix with the same shape as input
        random_no_mask = mask_with_tokens(random_tokens, mask_ignore_token_ids)        # not masked matrix for the random token matrix
        random_token_prob &= ~random_no_mask        # get the pure mask matrix of random token replace
        random_indices = torch.nonzero(random_token_prob, as_tuple=True)        # index of random token replace
        masked_input[random_indices] = random_tokens[random_indices]        # replace some tokens by random token
    
    # 5. 应用 [MASK] 标记
    replace_prob = prob_mask_like(data, replace_prob)   
    masked_input = masked_input.masked_fill(mask * replace_prob, mask_token_id)
    
    # 6. 生成标签
    labels = data.masked_fill(~mask, pad_token_id)     
    # 生成标签张量 labels，其中只有被遮掩的标记会被保留，其余位置被替换为 pad_token_id，表示这些位置不参与训练。
    
    return masked_input, labels

随机抽取数据样本并对其进行一些预处理，包括数值截断和转换为张量，最终将处理后的样本返回给模型训练时使用：

class SCDataset(Dataset):
    def __init__(self, data):
        super().__init__()
        self.data = data

    # 获取数据样本
    def __getitem__(self, index):
        rand_start = random.randint(0, self.data.shape[0]-1)
        full_seq = self.data[rand_start].toarray()[0]
        full_seq[full_seq > (CLASS - 2)] = CLASS - 2
        full_seq = torch.from_numpy(full_seq).long()
        full_seq = torch.cat((full_seq, torch.tensor([0]))).to(device)
        return full_seq

    def __len__(self):
        return self.data.shape[0]

3.1.2 gene2vec的embedding

作者的这部分专门有一个代码文件叫performer_pytorch.py很长，这里展示关键部分：

class Gene2VecPositionalEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len):
        super().__init__()
        gene2vec_weight = np.load('../data/gene2vec_16906.npy')
        gene2vec_weight = np.concatenate((gene2vec_weight, np.zeros((1, gene2vec_weight.shape[1]))), axis=0)
        gene2vec_weight = torch.from_numpy(gene2vec_weight)
        self.emb = nn.Embedding.from_pretrained(gene2vec_weight)

    def forward(self, x):
        t = torch.arange(x.shape[1], device=x.device)
        return self.emb(t)

加载预训练的 gene2vec 权重矩阵，该矩阵存储在.npy 文件中，每一行对应一个基因的嵌入向量。
np.concatenate部分: 在 gene2vec 权重矩阵的最后添加一行全零向量。这通常用于处理特殊标记（如 [PAD]），模型需要一个用于填充的嵌入向量。然后使用加载的 gene2vec 权重矩阵初始化一个嵌入层 self.emb。
前向传播部分：生成一个表示序列位置的整数张量 t，其长度等于输入 x 的序列长度（即 x.shape[1]）。

3.1.3 Performer的自注意力机制

类初始化：

class FastAttention(nn.Module):
    def __init__(self, dim_heads, nb_features = None, ortho_scaling = 0, causal = False, generalized_attention = False, kernel_fn = nn.ReLU(), no_projection = False):
        super().__init__()
        nb_features = default(nb_features, int(dim_heads * math.log(dim_heads)))

        self.dim_heads = dim_heads
        self.nb_features = nb_features
        self.ortho_scaling = ortho_scaling

        self.create_projection = partial(gaussian_orthogonal_random_matrix, nb_rows = self.nb_features, nb_columns = dim_heads, scaling = ortho_scaling)
        projection_matrix = self.create_projection()
        self.register_buffer('projection_matrix', projection_matrix)

        self.generalized_attention = generalized_attention
        self.kernel_fn = kernel_fn

        # if this is turned on, no projection will be used
        # queries and keys will be softmax-ed as in the original efficient attention paper
        self.no_projection = no_projection

        self.causal = causal
        if causal:
            try:
                import fast_transformers.causal_product.causal_product_cuda
                self.causal_linear_fn = partial(causal_linear_attention)
            except ImportError:
                print('unable to import cuda code for auto-regressive Performer. will default to the memory inefficient non-cuda version')
                self.causal_linear_fn = causal_linear_attention_noncuda

dim_heads: 每个注意力头的维度。
nb_features: 投影后的特征数，默认值是 dim_heads * log(dim_heads)。
ortho_scaling: 正交矩阵的缩放系数。
causal: 如果为 True，则使用因果注意力，这在自回归模型中常用。
generalized_attention: 如果为 True，使用广义注意力，通过可选的核函数来生成查询和键的核表示。
kernel_fn: 当使用广义注意力时，应用的核函数（默认为 ReLU）。
no_projection: 如果为 True，不使用任何投影。
重要属性：
self.create_projection: 生成投影矩阵的部分函数，投影矩阵用于将查询和键映射到较低维度空间。
self.projection_matrix: 通过 register_buffer 将投影矩阵注册为模型的一部分，它不会在反向传播中更新，但会随模型保存和加载。

重绘投影矩阵：

@torch.no_grad()
    def redraw_projection_matrix(self, device):
        projections = self.create_projection(device = device)
        self.projection_matrix.copy_(projections)
        del projections

在不进行梯度计算的情况下，重新生成投影矩阵，这可以用于在训练过程中定期更新投影矩阵以增加模型的鲁棒性。

前向传播：

def forward(self, q, k, v, output_attentions = False):
        device = q.device
        # inds = [8060, 8064, 6243, 8575, 10342, 10913, 9366, 993, 7796, 5210, 5212, 5504, 6851, 6559, 5508, 13107, 13820]
        if self.no_projection:
            q = q.softmax(dim = -1)
            k = torch.exp(k) if self.causal else k.softmax(dim = -2)

        elif self.generalized_attention:
            create_kernel = partial(generalized_kernel, kernel_fn = self.kernel_fn, projection_matrix = self.projection_matrix, device = device)
            q, k = map(create_kernel, (q, k))

        else:
            create_kernel = partial(softmax_kernel, projection_matrix = self.projection_matrix, device = device)
            q = create_kernel(q, is_query = True)
            k = create_kernel(k, is_query = False)

        attn_fn = linear_attention if not self.causal else self.causal_linear_fn
        out = attn_fn(q, k, v)
        if output_attentions:
            v_diag = torch.eye(v.shape[-2]).to(device)
            v_diag = v_diag.unsqueeze(0).unsqueeze(0).repeat(v.shape[0],v.shape[1],1,1)
            # attn_weights = torch.zeros(1, 1, len(inds), len(inds)).to(device).to(torch.float16)
            # attn_weights = torch.zeros(1, q.shape[1], len(inds), len(inds)).to(device).to(torch.float16)
            attn_weights = torch.zeros(1, 1, q.shape[2], q.shape[2]).to(device).to(torch.float16)
            for head_dim in range(q.shape[1]):
                # attn_weights[0, head_dim] = torch.abs(attn_fn(q[:,head_dim].to(torch.float16), k[:,head_dim].to(torch.float16), v_diag[:,head_dim].to(torch.float16)))[0, inds][:, inds]
                attn_weights += torch.abs(attn_fn(q[:,head_dim].to(torch.float16), k[:,head_dim].to(torch.float16), v_diag[:,head_dim].to(torch.float16)))
                # attn_weights += norm_tensor(torch.abs(attn_fn(q[:,head_dim].to(torch.float16), k[:,head_dim].to(torch.float16), v_diag[:,head_dim].to(torch.float16))), dim=-1)
            attn_weights /= q.shape[1]
            return out, attn_weights
        else:
            return out

no_projection：如果 no_projection 为 True，直接对 q 进行 softmax 处理，对 k 进行 softmax 或指数变换（取决于是否是因果注意力）。这种情况下，没有进行降维投影。
generalized_attention：如果 generalized_attention 为 True，则使用 generalized_kernel 对 q 和 k 进行变换。这个核函数可能对输入进行某种非线性映射。
标准注意力：如果上述两个选项都未启用，使用 softmax_kernel 进行降维投影。
注意力计算：根据 causal 参数，选择不同的注意力函数 (linear_attention 或 causal_linear_fn) 计算注意力输出。
输出注意力权重：如果 output_attentions 为 True，则额外计算和返回注意力权重。通过对值矩阵 v 进行对角矩阵的注意力操作，计算每个头的注意力贡献，并对所有头的结果进行平均。

3.2 微调模型

模型权重下载：https://drive.weixin.qq.com/s?k=AJEAIQdfAAoUxhXE7r

微调：

python -m torch.distributed.launch finetune.py --data_path "fine-tune_data_path" --model_path "pretrained_model_path"

细胞类型信息会从label和label_dict文件中读取

微调后的模型拿来预测细胞类型：

python predict.py --data_path "test_data_path" --model_path "finetuned_model_path"

检测新的细胞亚群：

python predict.py --data_path "test_data_path" --model_path "finetuned_model_path" --novel_type True --unassign_thres "custom_threshold"

在这里插入图片描述