Collapse-Free Prototype Readout Layer for Transformer Encoders

TL;DR

DDCL-Attention introduces a prototype-based readout layer for transformer encoders that ensures stable, collapse-free token summarization with broad applicability across NLP, vision, and scientific data.

Contribution

It proposes a novel, stable prototype-based readout mechanism that avoids collapse and supports multiple uses, outperforming standard quantization methods.

Findings

Prototypes remain distinct due to exact loss decomposition.

Stable joint training with the encoder is achievable under practical conditions.

The codebook achieves full utilization and outperforms standard quantization.

Abstract

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.