Collapse-Free Prototype Readout Layer for Transformer Encoders

Giansalvo Cirrincione; Rahul Ranjeev Kumar

cs.LGApr 4, 2026

Collapse-Free Prototype Readout Layer for Transformer Encoders

Authors: Giansalvo Cirrincione, Rahul Ranjeev Kumar

Organizations: Laboratory LTI, Université de Picardie Jules Verne, Amiens, France · Institute of Energy and Resources, Charles Darwin University, Darwin, NT, Australia

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.

Figures & tables

Property	Self-Attn	Slot Attn	Perceiver	DDCL-Attn
Key type	dynamic	dynamic (iter.)	dynamic	global static
Similarity	dot product	dot product	dot product	$-\\|\cdot\\|^{2}$
Complexity	$O(T^{2})$	$O(TKI)$	$O(TK)$	$\mathbf{O(TK)}$
Inference iter.	none	$I$ GRU steps	none	none
Anti-collapse	none	none	none	$\nabla_{P}V$
Stability proof	n/a	none	none	Thm. 1

Table 1: Structural comparison of prototype based attention mechanisms.

Pos.	Paradigm	Theoretical motivation	Val.
1	Final readout	Soft centroid compresses sequence into prototype vocabulary; $L_{\mathrm{OLS}}$ is the quantisation cost	$\checkmark$
2	Enc.–dec. bottleneck	Differentiable VQ bottleneck with anti-collapse guarantee	—
3	FFN replacement	Prototypes as learnable feed-forward basis, linear in $T$	—
4	Soft vector quantization	$\mathcal{L}_{q}=L_{\mathrm{OLS}}+V$ generalises VQ-VAE; $V$ prevents dead codes	$\checkmark$
5	Hierarchical compression	Decomposition holds level by level; anti-collapse at every level	$\checkmark$

Table 2: Five application paradigms for DDCL-Attention.

\checkmark

= validated in this paper.

Method	ACC	NMI	ARI
DDCL-Attention ( $\mathcal{L}_{q}$ , best epoch)	0.772	0.752	0.669
$k$ -means (raw, $d=7$ )	0.756	0.751	0.667
$k$ -means + PCA ( $m=5$ )	0.756	0.751	0.667

Table 3: Space debris clustering results (

N=1600

K=4

d=7

m=5

, best epoch ACC). Bold: best result.

Figure 1: 2D PCA projection of space debris features (

K=4

, after 500 epochs). Left: coloured by true orbital regime. Right: coloured by DDCL-Attention prototype assignment. Stars mark prototype positions

\mathbf{p}_{k}

. The LEO/HEO overlap in the lower left region reflects the genuine orbital ambiguity between low-altitude circular and Molniya-type eccentric orbits in 2D; the full

m=5

dimensional space resolves this via the eccentricity feature.

Dataset	Method	ACC	NMI	ARI
SST-2 ( $K=2$ )	CLS + logistic regression	0.861	—	—
	$k$ -means on CLS	0.519	0.003	0.001
	DDCL-Attention	0.867	0.435	0.538
IMDB ( $K=4$ )	$k$ -means on CLS	—	—	—
IMDB ( $K=4$ )	DDCL-Attention	0.913	0.472	0.540
20NG ( $K=20$ )	$k$ -means on CLS	0.196	0.189	0.065

Table 4: Paradigm 1 results: text readout with frozen BERT (

\varepsilon=0.05

). ACC = clustering accuracy (Hungarian); NMI = normalised mutual information; ARI = adjusted Rand index; best epoch reported.

Panel of Figure 2 — Figure 2: Exp 1: training dynamics for SST-2 (top), IMDB (middle), and 20 Newsgroups (bottom). Each row shows the loss decomposition $\mathcal{L}_{q}=L_{\mathrm{OLS}}+V$ (top left), variance term $V\geq 0$ (top centre), prototype separation $\mathcal{S}(P)$ (top right), assignment entropy $H(Q)$ (bottom left), clustering quality ACC/NMI/ARI (bottom centre), and phase portrait $(V/N,\mathcal{S}(P))$ (bottom right).

Figure 3: PCA projection of the 20 learned prototypes (

K=20

) on 20 Newsgroups after 15 epochs.

Method	$K$	Util. ep. 1	Epochs to 100%	$V\geq 0$
DDCL-Attention ( $\varepsilon=0.05$ , $\lambda=0.5$ )	16	100%	1	$\checkmark$
DDCL-Attention ( $\varepsilon=0.05$ , $\lambda=0.5$ )	64	100%	1	$\checkmark$
VQ-VAE (hard, straight-through)	16	81.2%	6	—
VQ-VAE (hard, straight-through)	64	18.8%	44	—

Table 5: Paradigm 4 results: soft VQ-VAE on CIFAR-10 (50 epochs). Codebook utilisation = fraction of prototypes with mean assignment

>0.01

Panel of Figure 4 — Figure 4: Exp 2: DDCL-Attention soft VQ on CIFAR-10 ( $K=64$ , 50 epochs). Left: training dynamics (loss decomposition, $V\geq 0$ , codebook utilisation). Right: codebook utilisation over epochs for DDCL-Attention vs. hard VQ-VAE; DDCL-Attention achieves 100% from epoch 1 while VQ-VAE requires 26 epochs.

Method	ACC	NMI	ARI
$k$ -means on [CLS]	0.200	0.211	0.060
$k$ -means on mean pooling	0.351	0.377	0.178
Hier. DDCL L1+L2 ( $\varepsilon=0.1$ , $\lambda=0.5$ )	0.112	0.075	0.009
Hier. DDCL L1+L2 ( $\varepsilon=0.05$ , $\lambda=1.5$ )	0.133	0.093	0.016

Table 6: Paradigm 5 results: hierarchical compression on 20 Newsgroups (

K_{1}=32

K_{2}=20

, frozen BERT, 15 epochs,

\varepsilon=0.05

\lambda=1.5

Figure 5: Exp 3: hierarchical DDCL-Attention on 20 Newsgroups (

K_{1}=32

K_{2}=20

, frozen BERT, 15 epochs). Left: training dynamics at both levels;

V^{(1)}\geq 0

and

V^{(2)}\geq 0

confirmed simultaneously at every epoch (Proposition 5 ). Right: t-SNE projection of the level-2 document representations coloured by DDCL-Attention assignment.

Figure 6: Stability ablation on MNIST Digits (

K=10

m=32

, 300 epochs). Left: best clustering ACC (bars) and final prototype separation

\mathcal{S}(P)

(line with markers) as a function of learning rate ratio

\varepsilon=\eta_{\theta}/\eta_{P}

. Green shading: stable regime (

\varepsilon\leq 0.1

, condition (iv) of Theorem 1 ). Red shading: boundary/unstable regime (

\varepsilon\geq 0.5

). Right:

(V/N,\mathcal{S}(P))

phase portrait coloured by epoch for three ratios; lower ratios converge to higher separation fixed points.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Symbol	Dimension	Meaning
$T$	scalar	sequence length
$d$	scalar	input embedding dimension
$m$	scalar	prototype (latent) dimension
$K$	scalar	number of prototypes
$N$	scalar	dataset size (number of sequences)
$H$	scalar	number of heads

Appendix

Table 7: Summary of notation.

Token Compression

See all 26 papers

Contrastive Projection: Reading Transformer Internals by Differencing Logit Lenses
Sep 9, 2026Olli TuomiTransformer Interpretability Language Model Introspection
Evident Solutions Oy
Decoding Looped Transformers Better for (Almost) Free
Oct 1, 2026Weihao Liu, Huangjie Zheng, Tianrong Chen +5Efficient Transformer Inference Language Model Decoding
Apple
DDCL-INCRT: A Self-Organising Transformer with Hierarchical Prototype Structure (Theoretical Foundations)
Apr 2, 2026Giansalvo CirrincioneTransformer Hierarchical Representation Learning
Laboratory LTI, Université de Picardie Jules Verne, Chemin du Thil, Amiens, 80025, France

Sep 9, 2026J/K move · Enter open · S save

Contrastive Projection: Reading Transformer Internals by Differencing Logit Lenses

Olli Tuomi

Evident Solutions Oy

Transformer Interpretability Language Model Introspection Transformer Attention

Open paper AskPDF arXiv

Reading a transformer's internal states in token space is easy to do and hard to trust: a logit lens on a single hidden state is dominated, at intermediate layers, by the generic tokens the model would predict for almost any input. We read the difference instead. Subtracting two closely matched prompts' hidden states and projecting through the unembedding cancels the shared component and surfaces what separates them, an operation equivalent to reading a RepE/ActAdd steering vector through a logit lens. Built into a training-free tracer that reads at every position, sub-layer, and head and averages over designed baselines, it traces a compound- noun MLP->attention chain in Phi-2, confirmed there by activation patching, with the same distinction recovered across three architectures by readout and probe rather than by patching; it reads what retrieval surfaces for real versus fictional entities, and reads metaphor as a set of domain-to-domain mappings rather than a single figurativity feature. A cross-seed control marks the boundary: across five networks differing only in initialization, the same distinction surfaces as almost entirely different tokens (top-10 overlap 0.08). What a computation looks like in token space is network-specific; the distinction it draws is not

Abstract

Token Compression

Explore similar work

Contrastive Projection: Reading Transformer Internals by Differencing Logit Lenses

Decoding Looped Transformers Better for (Almost) Free

DDCL-INCRT: A Self-Organising Transformer with Hierarchical Prototype Structure (Theoretical Foundations)