多项式自编码器在 Transformer 嵌入压缩上优于 PCA

# 多项式自编码器 来源: https://ivanpleshkov.dev/blog/polynomial-autoencoder/ 压缩嵌入（除了量化之外）最直接的方法是在语料库上拟合 PCA 并保留前 d 个特征向量。这种方法有效，但 PCA 是线性投影，而球面上的神经网络嵌入在结构上是非线性的——即 Transformer 中众所周知的**锥效应**（cone effect）。部分方差存在于线性解码器无法触及的非线性尾部中。 本文介绍了一种在 PCA 之上添加**二次**解码器的闭式方法，以捕捉部分非线性尾部。编码器保持为普通的 PCA。解码器是 2 阶多项式提升加上岭回归（带 L2 正则化的普通线性回归），同样为闭式解。无需 SGD、无需训练 epoch、无需超参数搜索。只需对语料库统计数据进行一次 `np.linalg.solve` 即可。 该构造本身并非我原创。“PCA 编码器 + 二次解码器 + 最小二乘拟合”在动力系统文献中以**二次流形**（quadratic manifold）的名称出现（参见 Jain 2017 (https://arxiv.org/abs/1610.09902), Geelen-Willcox 2022 (https://arxiv.org/abs/2205.02304)/2023 (https://arxiv.org/abs/2306.13748), Schwerdtner-Peherstorfer 2024+ (https://arxiv.org/abs/2403.06732)——详见§9 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#9-where-the-method-came-from-and-where-its-already-used)）。我是在运行实验并认为该构造是新的之后，才偶然发现了这些论文。多项式提升在现代机器学习讨论中并不常见，本文是一篇笔记，记录了一个来自相邻学科且在检索领域同样有用的技巧。 具体结果如下。BEIR/FiQA，`mxbai-embed-large-v1`（1024d），每向量预算 512 字节。指标为**NDCG@10**（前 10 个结果的归一化折损累计增益，标准的检索排序质量度量；范围 [0, 1]，越高越好）： 方法 | NDCG@10 | Δ vs raw 1024d --- | --- | --- raw 1024d (4096 bytes) | 0.4525 | — PCA top-256 | 0.4168 | -3.58 p.p. **poly-AE 256d** | **0.4441** | **-0.85 p.p.** matryoshka top-256 | 0.4039 | -4.86 p.p. PCA 已经实现了 4 倍的每向量内存压缩，NDCG 下降 3.58 个百分点（p.p.）。在 PCA 之上添加二次解码器又提升了**+2.73 p.p.**——几乎完全弥补了与原始数据之间的差距，且字节预算相同。Matryoshka 作为另一个熟悉的基线列在表中（此处其下降幅度大于 PCA——这是一个已知的侧面观察结果，并非本文的核心主张；参见§3 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#3-the-headline-table--four-models)脚注和§4 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#4-where-the-quadratic-decoder-actually-helps)）。 在四个模型（nomic-v1.5, mxbai-large, bge-base, e5-base）上测量。相对于 PCA，Poly-AE 在 d=128 时提升 +1 到 +4.4 p.p.，在 d=256 时提升 +0.03 到 +2.7 p.p.。完整表格见§3 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#3-the-headline-table--four-models)。 完整实现——约 150 行 numpy 代码，MIT 许可证，仓库地址 github.com/IvanPleshkov/poly-autoencoder (https://github.com/IvanPleshkov/poly-autoencoder)。BEIR 评估脚本位于同一仓库的 `beir_eval.py`。在 M 系列 MacBook 上复现需 30-40 分钟（编码语料库需 10-15 分钟，d=256 时 Ridge 求解需约 15 分钟）。 **目录：** - §1 我们比较什么 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#1-what-were-comparing) - §2 设置 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#2-experimental-setup) - §3 主要结果表 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#3-the-headline-table--four-models) - §4 二次解码器在何处起作用 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#4-where-the-quadratic-decoder-actually-helps) - §5 为何线性投影会丢失信息 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#5-why-a-linear-projection-loses-information) - §6 方法 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#6-a-polynomial-decoder-via-a-linear-lift) - §7 小语料库注意事项 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#7-the-small-corpus-measurement-and-in-sample-magic) - §8 残差压缩 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#8-compression--what-to-do-with-the-residual) - §9 方法的来源 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#9-where-the-method-came-from-and-where-its-already-used) - §10 局限性 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#10-limits-of-the-method) - §11 接下来可以尝试什么 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#11-what-to-try-next) ## 1. 我们比较什么 每次测量中都有四条基线： - **raw**—— 完整嵌入，无压缩。质量上限，每向量字节数最多。提供公平上限的“昂贵”基线。 - **matryoshka**—— `embedding[:d]` 加上 L2 归一化。对于使用 matryoshka 损失训练的模型（我们样本中的 nomic, mxbai），这是一个有效的 matryoshka 向量。对于未经过 matryoshka 训练的模型（bge-base, e5-base），这是对天真地切片非 MRL 模型会发生什么的测试——这是 bge 系列、e5 和自定义微调嵌入用户实际面临的情况。 - **PCA**—— 语料库协方差的前 d 个特征向量。向量生活在 d 维 PCA 坐标系中。 - **poly-AE**—— 我们的方法。用 PCA 编码为 `p ∈ R^d`，用二次多项式解码回全 D 维的 `V̂`，在 `V̂` 上进行检索。 在固定的 `d` 下，所有四种方法每向量存储 `2d` 字节（fp16 坐标）。 ## 2. 实验设置 BEIR 是一组标准的检索数据集（SciFact, FiQA, NFCorpus, TREC-COVID 等）。指标为 NDCG@10。语料库和查询使用选定的模型进行编码，通过余弦相似度检索前 10 个结果，并根据标记的 qrels 计算 NDCG。 PCA 和 poly-AE 是**转导式**拟合的：统计量在我们要压缩的语料库上计算。查询从不参与拟合——它们在推理时击中固定的编码器/解码器。这符合生产部署：索引操作员在其数据上一次性计算 PCA + Ridge，然后服务查询。 对于主要运行，我们使用**FiQA**—— 5.7 万篇文档，648 个查询，1706 个 qrels。 ## 3. 主要结果表——四个模型 FiQA 上在 256 fp16（512 字节/向量）和 128 fp16（256 字节）预算下的 NDCG@10： 模型 | D | d | raw | matryoshka† | PCA | **poly-AE** | poly over matryoshka | poly vs raw --- | --- | --- | --- | --- | --- | --- | --- | --- nomic-embed-text-v1.5 | 768 | 128 | 0.3746 | 0.3190 | 0.3273 | **0.3380** | +1.90 p.p. | -3.65 p.p. nomic-embed-text-v1.5 | 768 | 256 | 0.3746 | 0.3508 | 0.3670 | **0.3673** | +1.65 p.p. | -0.73 p.p. mxbai-embed-large-v1 | 1024 | 128 | 0.4525 | 0.3503 | 0.3689 | **0.4129** | +6.26 p.p. | -3.97 p.p. mxbai-embed-large-v1 | 1024 | 256 | 0.4525 | 0.4039 | 0.4168 | **0.4441** | +4.02 p.p. | -0.85 p.p. bge-base-en-v1.5* | 768 | 128 | 0.4062 | 0.2914 | 0.3266 | **0.3654** | +7.40 p.p. | -4.09 p.p. bge-base-en-v1.5* | 768 | 256 | 0.4062 | 0.3574 | 0.3688 | **0.3958** | +3.84 p.p. | -1.05 p.p. e5-base-v2* | 768 | 128 | 0.3987 | 0.2498 | 0.3065 | **0.3317** | +8.18 p.p. | -6.70 p.p. e5-base-v2* | 768 | 256 | 0.3987 | 0.3333 | 0.3618 | **0.3852** | +5.19 p.p. | -1.35 p.p. † “matryoshka”列为 `embedding[:d]` 加上 L2 归一化。在 nomic 和 mxbai 上，它是有效的 matryoshka 向量。在标记为 `*` 的模型（bge-base, e5-base）上，模型未针对 matryoshka 进行训练，此处的切片是对天真地切片非 MRL 模型会发生什么的测试。这是 bge 系列、e5 和自定义微调嵌入用户实际面临的情况——我们在此诚实测量。 这表明： 1. **Poly-AE 在所有四个模型上始终优于 PCA**。提升：d=128 时 NDCG +1 到 +4.4 p.p.，d=256 时 +0.03 到 +2.7 p.p.。二次解码器在何处起作用以及在什么 `d` 值下起作用——在§4 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#4-where-the-quadratic-decoder-actually-helps)中讨论。 2. **在 `d=256` 时，poly-AE 相对于原始 768/1024 维的 NDCG 损失为 0.7–1.4 p.p.** 在所有四个模型上。每向量 4 倍内存压缩，损失不到 1.5 p.p.——这是本文的主要数字。 3. **在未经 matryoshka 训练的模型上，matryoshka 列的下降幅度大于 PCA**——d=128 时 NDCG 高达 -15 p.p.。这是一个侧面观察：本文比较的是 PCA 和 poly-AE，而不是 PCA 与 matryoshka。如果表中的 matryoshka 数字看起来令人惊讶，§4 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#4-where-the-quadratic-decoder-actually-helps)中有一个简短的指针。 ## 4. 二次解码器在何处起作用 PCA 是线性基线。二次解码器添加了线性解码器无法触及的非线性部分（机制见§5 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#5-why-a-linear-projection-loses-information)）。这在检索上实际有多大帮助，以及在什么 `d` 值下？ Poly-AE 相对于 PCA 的提升，按模型和 `d` 划分： 模型 | poly over PCA, d=128 | poly over PCA, d=256 --- | --- | --- nomic-v1.5 | +1.07 p.p. | +0.03 p.p. mxbai-large | +4.40 p.p. | +2.73 p.p. bge-base | +3.88 p.p. | +2.70 p.p. e5-base-v2 | +2.52 p.p. | +2.34 p.p. 情况如下： 1. **在 d=128（8 倍压缩）时，poly 始终比 PCA 领先 1–4 p.p.** 这是线性解码器开始将显著方差丢弃到非线性尾部的区域，而二次校正将其拉回。该方法的最佳甜蜜点。 2. **在 d=256（4 倍压缩）时，差距不均匀。** 在 mxbai/bge/e5 上——稳定的 +2.3–2.7 p.p.。在 nomic 上——接近零（+0.03）。可能的原因：nomic 经过多切片对比损失的仔细训练，其潜在空间更具各向同性，且在 d=256 时，线性投影已经获取了大部分信息。在非 MRL 模型上，非线性尾部更大 → poly 帮助更大。 3. **各向异性越强 → 提升越大。** 锥效应越强，PCA 无法触及但 poly 可以触及的非线性尾部中的方差就越多。这是§5 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#5-why-a-linear-projection-loses-information)中展开的几何结构。 ### 侧面：matryoshka 在表中的位置 在§3 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#3-the-headline-table--four-models)中可以看到，在非 MRL 模型（bge, e5）上，matryoshka 列的下降幅度大于 PCA——即在随机非 MRL 训练的模型上，天真切片的效果不如语料库端的线性投影。这是一个已知结果；“MRL vs PCA 在检索中的表现”问题已独立于本文被讨论——参见 Matryoshka-Adaptor 2024 (https://arxiv.org/abs/2407.20243), SMEC “Rethinking MRL” 2025 (https://arxiv.org/abs/2510.12474), CoRECT 2025 (https://arxiv.org/abs/2510.19340)，以及标题字面意思为«Is PCA enough?» (https://www.youtube.com/watch?v=lklw59jQRKE)的 YouTube 视频。本文比较的是 PCA 和 poly-AE；matryoshka 作为第三个参考点出现在§3 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#3-the-headline-table--four-models)表中。 ### poly-AE 不适用的情况 语料库端的 PCA 拟合是该方法所需的一部分。这意味着当语料库不可用时，poly-AE**不起作用**： - **多租户 SaaS**：一个模型，数千个拥有不同语料库的客户——为每个客户拟合 PCA 是运营上的痛苦； - **流式索引**：统计数据随时间漂移，PCA 需要定期重新拟合； - **边缘推理**：手机、浏览器、嵌入式设备——你不想在模型旁边分发每个客户的 PCA 矩阵。 在这些设置中，你需要一个经过 MRL 训练的模型和 `embedding[:d]`，而 poly-AE 不是替代品——它也需要语料库统计数据。 ### 实际收获 在操作员拟合设置（固定语料库，操作员一次性拟合压缩）中，你有两种工作模式： - **d=256** 提供 4 倍压缩，相对于原始 NDCG 下降 -0.7 到 -1.4 p.p.。Poly 相对于 PCA：从 +0.03（nomic）到 +2.7 p.p.（mxbai/bge/e5）。在像 nomic 这样的 MRL 训练模型上，与 PCA 的差距极小；在非 MRL 模型上，poly 明显领先。 - **d=128** 提供 8 倍压缩。Poly 相对于 PCA：在任何模型上 +1 到 +4.4 p.p.。该方法的最佳甜蜜点。 ## 5. 为何线性投影会丢失信息 PCA 是将数据投影到 d 维子空间的最佳**线性**投影。但“最佳线性”并不意味着“足够好”：如果数据具有非线性结构，线性解码器就无法触及它，句号。 Transformer 嵌入具有这种结构，且已得到充分研究——**锥效应**。点云集中在单位球面上的狭窄锥体内，并在该锥体内具有强烈的非线性结构。 左：各向同性数据。右：各向异性数据的一个示例。 拖动以旋转 PCA 仅捕获沿正交特征向量的投影——即点云实际上不适合的线性椭球。流形曲率中的任何内容对线性解码器在结构上都是不可见的。各向异性越强（锥越窄），PCA 在结构上无法触及的非线性尾部中的方差就越多。 所以我们要的很清楚。一个可以处理**坐标二次组合**的解码器——即一个捕捉流形局部曲率的解码器。然后我们就可以恢复 PCA 丢失的部分信息。 ## 6. 通过线性提升实现多项式解码器 §5 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#5-why-a-linear-projection-loses-information)明确表明我们需要非线性解码器。直接的方法是训练神经网络。但这样我们就失去了闭式流水线：SGD、学习率、批次大小、收敛、早停。我们希望一个可以通过单个公式求解的非线性解码器。 标准回归技巧：**多项式提升**。取向量 `p` 并将其提升到所有高达 2 次的单项式——偏置、线性项、平方和 pairwise 乘积。在 2D 示例中： ``` lift([p1, p2]) = [1, p1, p2, p1^2, p1·p2, p2^2] ↑ ↑ ↑ ↑ ↑ ↑ bias linear quadratic ``` 这六个的任何线性组合 = 原始 `p1, p2` 的二次函数。所以**在提升上的线性回归 = 原始空间中的二次回归**。无需非线性优化，普通的闭式岭回归 OLS 即可完成。 将此应用于我们的情况。编码器是前 d 个 PCA（`p = (V - V_bar) @ Q`, `p ∈ R^d`）。将 `p` 提升到所有高达 2 次的单项式： ```python def polynomial_lift(p, degree=2): features = [1.0] # bias features.extend(p) # degree 1 for i in range(len(p)): for j in range(i, len(p)): features.append(p[i] * p[j]) # degree 2 return np.array(features) ``` 提升维度为 `M = (d+1)(d+2)/2`。对于 `d=128`，那是 `M = 8385`，对于 `d=256`，那是 `M = 33153`。解码器是从 M 维提升回到 D 维原始空间的线性回归。尽管只是“线性回归”，但结果映射 `p → V̂` 是二次的——这正是我们需要的。 训练解码器只需一次 `np.linalg.solve`： ```python def fit_decoder(P, V, lam=1e-3): L = polynomial_lift(P, degree=2) # (N, M) G = L.T @ L + lam * np.trace(L.T @ L) / L.shape[1] * np.eye(L.shape[1]) W = np.linalg.solve(G, L.T @ V) return W ``` 多项式自编码器组装如下： 多项式自编码器架构：V → PCA → p → poly lift → L → Ridge OLS → V̂; V − V̂ → V_resid - **编码器**—— 闭式线性（PCA），无需训练， - **解码器**—— 闭式二次（多项式提升上的 Ridge OLS），通过一次 `np.linalg.solve` 训练， - **重建**—— `V̂ = polynomial_lift(p) @ W`， - **残差**—— `V_resid = V - V̂`（在§8 (https://ivanpleshkov.dev/blog/polynomial-autoencoder/#8-compression--what-to-do-with-the-residual)中用于量化）。 无需反向传播、学习率、批次大小、epoch。在 d=100 的 10 万个向量上，CPU 只需几分钟。在 d=256 时，大约需要 15 分钟（Ridge 求解一个 `M^3` 系统，与 `M` 成三次方关系）。 ### 关于归一化的一个技术说明 原始 `p` a