Collapse-Free Prototype Readout Layer for Transformer Encoders

📄 arXiv: 2604.03850 📥 PDF

作者: Giansalvo Cirrincione, Rahul Ranjeev Kumar

分类: cs.LG, cs.NE

发布日期: 2026-04-07

💡 一句话要点

提出DDCL-Attention，一种无崩溃原型读取层，用于Transformer编码器。

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: Transformer编码器 原型学习 注意力机制 向量量化 无崩溃训练

📋 核心要点

  1. 传统Transformer编码器的读取层依赖池化或类别token，存在信息损失和原型崩溃问题。
  2. DDCL-Attention通过学习一组原型向量，使用软概率匹配将token分配给原型，实现高效压缩。
  3. 实验证明，该方法避免了原型崩溃，提高了码本利用率，并在多个任务上优于现有方法。

📝 摘要（中文）

DDCL-Attention是一种基于原型的Transformer编码器读取层，它使用学习到的压缩机制取代了简单的池化方法（如平均池化或类别token）。该方法使用一小组全局原型向量，并通过软概率匹配将token分配给它们，从而以序列长度的线性复杂度生成紧凑的token摘要。该方法具有三个主要优点。首先，它通过将训练损失精确分解为重构项和多样性项来避免原型崩溃，确保原型保持不同。其次，使用Tikhonov奇异摄动理论和显式学习率约束，证明了其与编码器的联合训练在实际时间尺度条件下是稳定的。第三，同一框架支持三种用途：最终读取层、扩展VQ-VAE的可微码本和分层文档压缩器。在四个数据集上的实验证实了理论预测：损失分解完全成立，当满足稳定性条件时，原型分离如预期增长，并且码本达到完全利用率，优于标准硬向量量化。一项关于轨道碎片分类的额外研究表明，该方法也适用于标准NLP和视觉任务之外的科学表格数据。

🔬 方法详解

问题定义：Transformer编码器通常使用平均池化或类别token作为读取层，将变长序列压缩为固定长度的向量表示。这些方法存在信息损失，并且在原型学习中容易出现原型崩溃问题，即所有原型向量收敛到相同的值，导致表示能力下降。

核心思路：论文的核心思路是使用一组可学习的原型向量来表示输入序列，并通过软概率匹配的方式将每个token分配给不同的原型。通过优化重构损失和原型多样性损失，避免原型崩溃，并学习到更具区分性的原型表示。

技术框架：DDCL-Attention包含以下主要模块：1) Transformer编码器：用于提取输入序列的token表示。2) 原型向量：一组可学习的全局原型向量。3) 软分配模块：计算每个token与每个原型之间的相似度，并使用softmax函数生成软分配概率。4) 加权平均：使用软分配概率对token表示进行加权平均，得到原型表示。5) 重构损失：衡量原始token表示与原型表示之间的重构误差。6) 多样性损失：鼓励原型向量之间的差异性，避免原型崩溃。

关键创新：该方法最重要的创新点在于提出了一个精确的损失分解，将训练损失分解为重构项和多样性项。这种分解保证了原型向量在训练过程中保持多样性，避免了原型崩溃问题。此外，论文还分析了联合训练的稳定性条件，并给出了显式的学习率约束。

关键设计：关键设计包括：1) 使用余弦相似度计算token和原型之间的相似度。2) 使用softmax函数生成软分配概率。3) 使用Tikhonov奇异摄动理论分析联合训练的稳定性。4) 设计了多样性损失函数，鼓励原型向量之间的差异性。5) 通过实验验证了损失分解的有效性和稳定性条件的正确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DDCL-Attention在四个数据集上都取得了良好的性能。损失分解能够精确成立，原型分离随着训练的进行而增长，码本利用率达到完全。在轨道碎片分类任务中，该方法也表现出良好的泛化能力，证明了其在非标准NLP和视觉任务中的适用性。该方法优于标准硬向量量化方法。

🎯 应用场景

DDCL-Attention可应用于各种需要将变长序列压缩为固定长度表示的任务，例如文本分类、图像分类、文档压缩和科学数据分析。该方法能够学习到更具区分性的表示，提高下游任务的性能。此外，该方法还可以作为可微码本扩展VQ-VAE，用于生成模型。

📄 摘要（原文）

DDCL-Attention is a prototype-based readout layer for transformer encoders that replaces simple pooling methods, such as mean pooling or class tokens, with a learned compression mechanism. It uses a small set of global prototype vectors and assigns tokens to them through soft probabilistic matching, producing compact token summaries at linear complexity in sequence length.The method offers three main advantages. First, it avoids prototype collapse through an exact decomposition of the training loss into a reconstruction term and a diversity term, ensuring that prototypes remain distinct. Second, its joint training with the encoder is shown to be stable under a practical timescale condition, using Tikhonov's singular perturbation theory and explicit learning-rate constraints. Third, the same framework supports three uses: a final readout layer, a differentiable codebook extending VQ-VAE, and a hierarchical document compressor.Experiments on four datasets confirm the theoretical predictions: the loss decomposition holds exactly, prototype separation grows as expected when the stability condition is met, and the codebook reaches full utilization, outperforming standard hard vector quantization. An additional study on orbital debris classification shows that the method also applies beyond standard NLP and vision tasks, including scientific tabular data.