使用 ESM-2 估计蛋白质序列嵌入的内在维度

TLDR: 在这篇文章中，我们将展示如何使用与持久同调相关的技术来估计从蛋白质语言模型（ESM-2）获得的嵌入的内在维度。然后，我们将讨论如何将其用于新蛋白质语言模型的课程学习。

理解数据集的底层结构和复杂性至关重要。衡量这一点的一种方法是估计数据的内在维度。这与神经网络的泛化能力有关，例如在内在维度、持久同调和神经网络中的泛化、深度学习中的隐式正则化可能无法用范数解释以及大型 Transformer 模型隐藏表示的几何结构中都有提及。

其中一些工作促使人们考虑诸如内在维度解释语言模型微调的有效性等论文，这为现在广泛采用的LoRA 技术提供了灵感，该技术能够以参数高效的方式微调模型。特别是，我们可以将 LoRA 视为一种重要的正则化技术，它实际上可以提高模型的泛化能力，并在未见数据上提供更好的性能。但这与我们 هنا 的目标有点离题。

在这篇文章中，我们将讨论如何使用基于蛋白质语言模型嵌入和持久同调的独特方法来估计蛋白质序列的内在维度。特别是，我们将研究 ESM-2 等蛋白质语言模型的嵌入，并讨论如何使用最小生成树和线性回归来估计嵌入的内在维度。这是对持久同调维度（一种分形维度）的估计。这将为我们提供单个蛋白质复杂性的度量，然后可用于训练新的 pLM 时的课程学习。

1. 什么是内在维度？

数据集的内在维度是衡量描述数据所需参数数量的度量，或者简单来说，是表示数据而不会造成太多损失的最小坐标数量。它让我们深入了解数据的“复杂性”。

例如，即使数据集处于高维空间中，如果数据点位于接近子空间（如平面或曲线）的位置，其内在维度可能较低。

数学上，如果我们有从嵌入在 ${\mathbb{R}}^D\backslash mathbb\{R\}^D$ 中的某个流形 $MM$ 中采样的点，则内在维度是 $MM$ 的最小维度。

2. 词元嵌入

导入

import numpy as np
from sklearn.linear_model import LinearRegression
from scipy.spatial import distance_matrix
from scipy.sparse.csgraph import minimum_spanning_tree
from transformers import AutoTokenizer, AutoModel
import torch

嵌入

要估计蛋白质序列的内在维度，第一步是计算序列中每个词元（氨基酸）的嵌入。嵌入是捕获词元上下文信息的密集向量。

def get_embeddings(text, model_name="facebook/esm2_t6_8M_UR50D"):
    """
    Compute embeddings for each token in the text using a specified model.
    
    Parameters:
    - text (str): The input text for which embeddings need to be computed.
    - model_name (str): The path to the pretrained model.
    
    Returns:
    - numpy.ndarray: A matrix where each row is the embedding of a token in the text.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)

    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=1024)
    with torch.no_grad():
        outputs = model(**inputs)

    # Return embeddings after removing <cls> and <eos> tokens and converting to numpy.
    return outputs.last_hidden_state[:, 1:-1, :].squeeze(0).numpy()

这里，我们使用一个 Transformer 模型来获取嵌入。该模型接收一个蛋白质序列，将其词元化，并返回每个词元的嵌入。我们忽略了 Transformer 架构中通常使用的特殊 `<cls>` 和 `<eos>` 词元。

3. 持久分数

一旦我们有了每个词元的嵌入，我们就会使用拓扑数据分析中的一个概念来计算这些嵌入子集的持久分数。持久分数基于从嵌入的距离矩阵形成的最小生成树（MST）的边权重之和。

def compute_persistent_score(embeddings):
    """
    Compute the persistent score for a subset of embeddings using the sum of edge weights in the MST.
    
    Parameters:
    - embeddings (numpy.ndarray): A matrix where each row is an embedding.
    
    Returns:
    - float: The persistent score for the embeddings.
    """
    dist_matrix = distance_matrix(embeddings, embeddings)
    mst = minimum_spanning_tree(dist_matrix)
    return mst.sum()

最小生成树是连接的、边加权图的边的子集，它将所有顶点连接在一起，没有循环，并且总边权重最小。这些权重之和是衡量嵌入在其空间中“散布”或“复杂性”的度量。

4. 采样和评分

给定嵌入，下一步是采样词元子集并计算其持久分数。这会针对各种样本大小重复进行。

def sample_and_score(embeddings, n, k=8, hat_n=40, J=7):
    """
    For various sample sizes, compute the median persistent score across J samples.
    
    Parameters:
    - embeddings (numpy.ndarray): A matrix where each row is an embedding.
    - n (int): Total number of embeddings.
    - k (int): Number of different sample sizes.
    - hat_n (int): A parameter for determining sample sizes.
    - J (int): Number of samples for each sample size.
    
    Returns:
    - list: List of sample sizes.
    - list: List of corresponding median persistent scores.
    """
    scores = []
    sizes = [(i - 1) * (n - hat_n) // k + hat_n for i in range(1, k + 1)]
    
    for size in sizes:
        subset_scores = [compute_persistent_score(embeddings[np.random.choice(n, size, replace=False)])
                         for _ in range(J)]
        scores.append(np.median(subset_scores))
    
    return sizes, scores

对于每个样本大小，我们随机选择一个词元子集，计算持久分数，然后重复此操作 $JJ$ 次以获得该大小的中位数分数。

5. 维度估计

该方法的关键在于样本大小及其相应的持久分数之间的关系。在对数-对数尺度上，这种关系对于许多数据集来说近似线性。最佳拟合线的斜率给我们提供了内在维度的估计。

$$\text{dimension}=\frac{1}{1-\text{slope}}\backslash text\{dimension\}\; =\; \backslash frac\{1\}\{1\; -\; \backslash text\{slope\}\}$$

具体操作如下：

def estimate_dimension(sizes, scores):
    """
    Estimate the intrinsic dimension of the data using linear regression on log-transformed sizes and scores.
    
    Parameters:
    - sizes (list): List of sample sizes.
    - scores (list): List of corresponding median persistent scores.
    
    Returns:
    - float: Estimated dimension of the data.
    """
    log_sizes = np.log(sizes).reshape(-1, 1)
    log_scores = np.log(scores)

    reg = LinearRegression().fit(log_sizes, log_scores)
    slope = reg.coef_[0]
    
    return 1 / (1 - slope)

6. 融会贯通

最后，为了获得对内在维度的可靠估计，我们重复采样、评分和维度估计几次并对结果取平均。

def estimate_sequence_dimension(text, runs=5):
    """
    Estimate the intrinsic dimension of the text by repeatedly sampling subsets of its tokens, 
    computing their persistent scores, and then using linear regression on the log-transformed values.
    
    Parameters:
    - text (str): The input text for which the dimension needs to be estimated.
    - runs (int): Number of runs with different random seeds.
    
    Returns:
    - float: Estimated dimension of the text.
    """
    embeddings = get_embeddings(text)
    n = embeddings.shape[0]
    
    slopes = []
    for _ in range(runs):
        sizes, scores = sample_and_score(embeddings, n)
        log_sizes = np.log(sizes).reshape(-1, 1)
        log_scores = np.log(scores)
        
        reg = LinearRegression().fit(log_sizes, log_scores)
        slopes.append(reg.coef_[0])
    
    kappa_F = np.mean(slopes)
    return 1 / (1 - kappa_F)

当应用于蛋白质序列时：

text = "MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE"
dimension = estimate_sequence_dimension(text)
print(f"Estimated dimension of the protein sequence: {dimension}")

将打印出：

Estimated dimension of the protein sequence: 13.063370658316673

输出提供了蛋白质序列的估计内在维度。

其他方法 (Scikit-Dimension)

存在其他估算数据内在维度的方法，例如“MLE”和“TwoNN”，根据您的用例，您可能还需要考虑这些方法。让我们看看 `scikit-dimension` 中使用的各种方法，看看是否有任何方法与我们的方法接近。首先，您需要运行：

pip install scikit-dimension

然后，务必导入它：

import skdim

接下来，我们可以使用以下函数打印由 `skdim` 中各种方法估计的内在维度：

methods = {
    "corr_int": skdim.id.CorrInt,
    "danco": skdim.id.DANCo,
    "ess": skdim.id.ESS,
    "fisher_s": skdim.id.FisherS,
    "knn": skdim.id.KNN,
    "lpca": skdim.id.lPCA,
    "mada": skdim.id.MADA,
    "mind_ml": skdim.id.MiND_ML,
    "mle": skdim.id.MLE,
    "mom": skdim.id.MOM,
    "tle": skdim.id.TLE,
    "twonn": skdim.id.TwoNN
}

def estimate_dimension(embeddings, method="twonn"):
    """
    Estimate the intrinsic dimension of embeddings using the specified method.
    
    Parameters:
    - embeddings (numpy.ndarray): A matrix where each row is an embedding.
    - method (str): The method to use for dimension estimation.
    
    Returns:
    - float: The estimated intrinsic dimension.
    """
    
    if method not in methods:
        raise ValueError(f"Unknown method: {method}")
    
    id_est = methods[method]().fit(embeddings)
    return id_est.dimension_

然后，您可以这样使用它：

# Example usage:
text = "MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE"
embeddings = get_embeddings(text)

for method in methods.keys():
    dimension_estimate = estimate_dimension(embeddings, method=method)
    print(f"Estimated Intrinsic Dimension ({method.upper()}): {dimension_estimate}")

这将打印出：

Estimated Intrinsic Dimension (CORR_INT): 7.599677519235372
Estimated Intrinsic Dimension (DANCO): 310.95695219616096
Estimated Intrinsic Dimension (ESS): 45.757353328926165
Estimated Intrinsic Dimension (FISHER_S): 11.925543780836733
Estimated Intrinsic Dimension (KNN): 6
Estimated Intrinsic Dimension (LPCA): 48
Estimated Intrinsic Dimension (MADA): 15.526202715518686
Estimated Intrinsic Dimension (MIND_ML): 10.0
Estimated Intrinsic Dimension (MLE): 11.85228294928126
Estimated Intrinsic Dimension (MOM): 4.662291966147815
Estimated Intrinsic Dimension (TLE): 11.681521116520777
Estimated Intrinsic Dimension (TWONN): 11.715313108714346

在这里，我们可以看到 `skdim` 中可用的大多数方法提供的估计值都接近我们对持久同调维度的估计值 `13.063370658316673`。我们应该在此指出，持久同调方法是随机的，这就是我们对多次迭代取平均的原因。因此，除非您使用相同的随机种子，否则您不会每次都得到完全相同的答案。

结论和一些用例

理解数据（包括蛋白质序列）的内在维度对于各种分析和计算任务都非常宝贵。它能深入了解数据的底层结构，进而可用于优化、可视化等等。这里讨论的方法尤其引人入胜，因为它结合了机器学习和拓扑数据分析（持久同调）的先进技术。

这有什么用呢？这个方法最初是在文章《内在维度估计用于鲁棒检测 AI 生成文本》中被开发出来的，试图检测 AI 生成的文本。所以，特别是，这可能被用来检测 AI 生成的蛋白质。然而，还有更有趣的应用。例如，如果我们首先用内在维度较低的蛋白质训练蛋白质语言模型，然后再用内在维度逐渐升高的蛋白质训练模型呢？这提供了一种课程学习形式，模型首先学习较简单的数据，然后逐渐学习较难的数据。

将内在维度通过课程学习整合到大型语言模型中也可能很有用，并且这很容易转移到 NLP 领域。这可能会增加 AI 生成文本的内在维度，使其与人类生成文本的区分度降低。此外，相同的想法也可能适用于蛋白质语言模型。

最后，我们还建议阅读 连接语言模型中的信息论和几何压缩、多尺度拓扑奇异点检测、大型 Transformer 模型隐藏表示的几何结构以及 蛋白质语言模型隐藏表示的几何结构。最后两篇论文似乎暗示，LoRA 的秩应该按层选择，并且应该大致与该层的内在维度匹配。此外，我最近研究表明 LoRA 和 QLoRA 可以成为有效的正则化技术，改善对未见数据的泛化能力，并显著减少过拟合。如果由于过拟合而需要正则化，选择低于该层内在维度的 LoRA 秩可能会有所帮助。