Latest / Embodied AI 101 / Unembedding Matrix as a Feature Lens: Unlocking Better Text Embeddings
Transcript
- Unembedding Matrix as a Feature Lens: Unlocking Better Text Embeddings. Large language models (LLMs) like GPT and PaLM often produce surprisingly good “text embeddings” – fixed‐length vectors for sentences – even though they weren’t explicitly trained for that purpose. These embeddings are used in applications like semantic search or retrieval. But in practice, off‐the‐shelf embeddings from LLMs frequently underperform compared to purpose-built embedding models (such as OpenAI’s code‐oriented or contrastively trained E5 models). Why? Wu et al. identify a central culprit: the model’s unembedding matrix, i.e. the final weight matrix that maps hidden activations to vocabulary logits, inadvertently introduces a bias toward common tokens. In effect, LLM embeddings carry a “bleed” of high‐frequency word features at the expense of subtle semantics. The key insight of this work is that the very unembedding matrix – usually thought of only as part of the model’s output layer – secretly encodes a latent subspace dominated by these frequent, but semantically “hollow,” tokens. By treating this final weight matrix as a kind of “feature lens,” the authors show we can filter out that subspace and markedly improve the quality of LLM‐derived embeddings without any additional model training. Their method, called EmbedFilter, is a simple linear projection that zeroes out the offending directions in embedding space. The result: sharper, more semantically rich embeddings, often at reduced dimensionality, and big gains on downstream tasks like few-shot retrieval. In the rest of this article, we’ll unpack how this works. We’ll first build intuition about why high-frequency tokens spoil embeddings, then explain how the authors use ideas from logit lens interpretation and spectral analysis of the unembedding matrix to isolate those effects. Next we’ll dive into the EmbedFilter transformation itself and how it can be applied (its training-free and distance-preserving nature). We’ll review the experiments – notably on the MTEB benchmark – that demonstrate consistent gains and even allow smaller embeddings with no loss in quality. Finally we’ll discuss how this fits into the broader landscape of embedding techniques, why it matters, and what future directions it opens up. Why LLM Embeddings Can Be Flawed. Modern LLMs are often used “off‐the‐shelf” as embedding generators: given a prompt like “Compute an embedding of this sentence,” one can pass a string through a pretrained LLM and extract the final hidden state (for example the representation of the last token, or an average of all tokens). This vector is intended to capture the sentence’s semantic content. However, practitioners have repeatedly observed that these embeddings frequently fall short. In simple retrieval tests or semantic tasks, specialized sentence embedding models (trained directly for similarity, like sentence-BERT variants or contrastively tuned models) tend to outperform raw LLM embeddings. Researchers have shown that contextual embeddings from LLMs tend to be highly anisotropic – they cluster in a narrow cone rather than being spread out. This anisotropy makes many embeddings overly similar and hurts cosine-similarity judgments among sentences. In particular, Wu et al. noticed a more specific phenomenon: when you take an LLM’s text embedding and use the model’s output matrix to ask “which tokens does this embedding most strongly predict,” the answers are dominated by very common words (e.g. “the”, “to”, “of”) that carry little specific meaning. In their words, “text embeddings tend to align with frequent but uninformative tokens when projected onto the vocabulary space”. In other words, if you treat the embedding as if it were a hidden state and compute logits via the unembedding matrix, the top-scoring words will be high-frequency stopwords, not content words that reflect the sentence’s niche meaning. This tendency “suppresses the model’s ability to capture nuanced semantics,” writes the paper. You can think of it in practical terms: suppose an LLM read the sentence “A new HIV drug has been approved by the FDA.” The embedding might naturally pick up on words like “HIV” or “FDA,” but it also (over)focuses on very common terms like “A,” “new,” “has,” etc. These frequent tokens get amplified because the LLM’s output layer has “learned” them so well during training (they appear extremely often in text), whereas less frequent content words have weaker components in that space. The net effect is that two sentences that share many common words will look unduly similar unless some calibration is done. This problem is reminiscent of classic issues in word embeddings. For example, prior work on static word vectors showed that a few “common” directions (like the overall mean vector or top principal components) often dominate, and removing them improves performance. Here a similar phenomenon arises, but at the LLM sentence level and with no additional training. The noted strategy of Mu et al. – known as “All-but-the-Top” – projects away top PCA directions from a word embedding space to debias it. EmbedFilter can be seen as a data-driven extension of that idea, where the relevant directions are discovered from the LLM’s own parameters (the unembedding matrix) rather than from corpus statistics, and then used to fix the embeddings. The Unembedding Matrix as a Lens. To understand the remedy, it helps to clarify the role of the unembedding matrix. In a typical autoregressive transformer, after processing input tokens through all layers, the final layer’s hidden state $\boldsymbol h \in \mathbb R^d$ (for each token position) is multiplied by a matrix $W_U \in \mathbb R^{|\mathcal V|\times d}$ and passed through softmax to produce token probabilities. Here $|\mathcal V|$ is the vocabulary size. Each row of $W_U$ corresponds to one vocabulary token: effectively it measures how strongly each hidden feature predicts that token. In practice $W_U$ is often tied to the input embedding matrix (sharing weights), but either way it’s a fixed learned mapping from $\mathbb R^d$ to token logits. Mechanistic interpretability practitioners realized that if you take an intermediate hidden state $\boldsymbol h$ and apply $W_U$ (the so-called logit lens technique), you can see what distributions over tokens the model is “thinking” about before finishing computation. The original logit lens research observed that decoding hidden states with the pretrained unembedding matrix gives surprising predictive power – the model is continuously assembling its final prediction. In their analysis of the tuned lens, Belrose et al. describe exactly this: the “logit lens” decodes hidden states into vocabulary space using the model’s pretrained unembedding matrix. In short, $W_U$ is the direct link between hidden features and the vocabulary. What Wu et al. realized is that this $W_U$ contains a latent feature space that can be exploited more directly. If we imagine the unembedding matrix’s singular value decomposition $W_U = U \Sigma V^T$, the matrix $V \in \mathbb R^{d\times d}$ (assuming full rank) is an orthonormal basis for the model’s hidden space, decomposed such that each singular vector direction in $V$ is associated with how strongly it influences many output tokens (weighted by $\Sigma$). In particular, they show that some of these singular directions correspond to high-frequency tokens. Intuitively, if you take the top singular vectors (largest singular values) and project embeddings onto them and then through $W_U$, you largely get the effect of the most common vocabulary entries. The authors confirm this by a sort of “logit spectroscopy”: they take an “average token” embedding – using word frequencies and the pseudo-inverse $W_U^+$ – and measure how each singular direction contributes to its logits. They find that the edge spectrum (the directions corresponding to the very largest and very smallest singular values) are the ones driving up probabilities for frequent words. In contrast, the mid-range directions capture more nuanced content. Put simply, the unembedding matrix is encoding a “noisy subspace” of hidden vectors that light up common tokens. Projecting the embeddings along those directions is akin to adding a bias for stuff like “the,” “to,” “of,” etc. (By contrast, projecting along the rest of the spectrum tends to reflect actual content semantics.) The authors convincingly argue this points to a structural misalignment: $W_U$ was trained for generation (predicting next tokens) rather than for producing semantic embeddings, so as-is it pollutes the embedding space with frequency artifacts. Designing the EmbedFilter Transform. Armed with that insight, the paper proposes a remarkably simple corrective transformation. They call it EmbedFilter. Mathematically, it is just an orthogonal projection in hidden space that removes the identified “noisy” directions. Concretely, if $V \Sigma U^T$ is the SVD of $W_U$ (with $V$ listing the hidden-space vectors), they choose a cutoff range $[l_\tau, r_\tau]$ that excludes the top few and bottom few singular vectors. Denote $V_{[l_\tau:r_\tau]}$ as the submatrix of columns $l_\tau$ to $r_\tau$. Then the transformation matrix is $$ \Phi_\tau ;=; V_{[l_\tau:r_\tau]} ,V_{[l_\tau:r_\tau]}^T, $$ an orthogonal projector onto the mid-spectrum subspace. If the original LLM embedding of a sentence is $\boldsymbol e \in \mathbb R^d$, then the filtered embedding is $$ \widetilde{\boldsymbol e} = \boldsymbol e ,\Phi_\tau^T = \boldsymbol e ,\Phi_\tau $$ (since $\Phi_\tau$ is symmetric). Essentially, $\widetilde{\boldsymbol e}$ has removed the components of $\boldsymbol e$ that lie in the discarded subspace (the “high-frequency token lens”). Several nice properties follow. Because $\Phi_\tau$ is formed from orthonormal singular vectors, it is itself an orthonormal projection. Thus it preserves all pairwise distances between points in the subspace it keeps. In practice this means the refined embeddings remain directly comparable: cosine similarities (or any other distance metric) aren’t arbitrarily warped by the filter. It also means that one can actually reduce the dimensionality of the embeddings at no loss. For example, if the original hidden dimension is $d=1024$ and you choose to remove, say, 50 directions (25 from the top and 25 from the bottom of the singular spectrum), your $\Phi_\tau$ effectively projects into a $(d-50)$-dimensional subspace in an orthonormal way. You can store only those $(d-50)$ coordinates going forward. In other words, EmbedFilter simultaneously denoises and compresses the embeddings. Operationally, EmbedFilter is appealingly lightweight: it requires no extra training. One only needs to perform an SVD on the existing $W_U$ matrix (which is a one-time cost and can be done offline). Then applying $\Phi_\tau$ to each embedding is just a matrix multiply. This is orders of magnitude cheaper than fine-tuning the whole model or training a new calibration head. It also means the method is transparent and local: it doesn’t change the LLM’s weights or require any ground-truth labels. In summary, EmbedFilter is essentially a train-free postprocessing filter that “filters out” the high-frequency token subspace (identified via the SVD of the unembedding). This compensates for the design mismatch that the decoder was never meant to serve as an embedding lens. The name “feature lens” in the paper’s title refers to this idea of using $W_U$ to define what features to keep or discard. Semantic Gains and Dimensionality Benefits. How much of a difference does this make in practice? The authors evaluate EmbedFilter on a broad range of zero-shot text embedding tasks using the Massive Text Embedding Benchmark (MTEB) – a suite of tasks covering semantic textual similarity, clustering, classification, and retrieval. They try EmbedFilter on multiple popular LLM backbones. In nearly all cases, the filtered embeddings outperform the raw ones, often by a gold-standard margin. The headline number they report is an up to 14.1% improvement on aggregate MTEB scores. To put it in perspective, such an improvement from a single linear map is quite significant. For retrieval tasks in particular (where embedding quality is critical), the gains are notable. The filtered embeddings bring the LLM almost on par with specialized embedding models. The authors stress that these improvements come without any additional model training – effectively “for free” except for the cost of computing the SVD. This lightweight calibration step consistently tightened the clustering of semantically similar sentences, as demonstrated by higher similarities within test sets, and lowered false matches. Likewise, for classification tasks (where one might use embeddings as features for a nearest-neighbor classifier), the boosted embeddings distinguished labels more clearly. A compelling side effect is that the refined embeddings still preserve performance even if their dimension is substantially cut. Because $\Phi_\tau$ is distance-preserving on its subspace, one can simply truncate the embedding to the $(r_\tau - l_\tau +1)$ active dimensions (or equivalently project and drop the zeros) and not lose any information in similarity comparisons. In practice, Wu et al. show that they can halve the embedding size in some cases with nearly identical accuracy. This implies faster nearest-neighbor lookup (critical at scale) and roughly half the storage for indices, all while maintaining the semantic gains of filtering. For example, suppose an LLM has hidden size 2048. The unfiltered embeddings are 2048-dimensional vectors. Using EmbedFilter with reasonable cutoff thresholds might yield a subspace of only say 1800 dimensions. Those 1800 numbers (instead of 2048) fully represent the filtered embedding. Even smaller is possible if one accepts minimal approximation. Across models, the authors note a sweet spot: remove enough subspace to kill the frequency bias but keep the bulk of semantics. Empirically, this often means discarding on the order of a few dozen singular directions on each end, out of hundreds. We can contextualize these gains by analogy. In early word embeddings, researchers found that the “common” direction (often associated with the vector average of the whole vocabulary) carried no linguistic meaning. Removing it (and a few large PCA components) consistently improved word similarity and text classification tasks. EmbedFilter is following that same intuition at the sentence level, but detected through the lens of the LLM’s own output weights. The fact that a mere projection like this beats more elaborate fixes (e.g. prompt engineering or training data tuning) underscores how deeply buried the problem was in the model’s linear algebra. Importantly, the authors compare EmbedFilter not only to the unfiltered baseline but also to simpler ablations. For instance, what if you just randomly do PCA on the embeddings and drop top components? Or if you normalize the embeddings in some common way? The paper reports that none of the naive baselines (like random projections, trivial dimension truncation, or simple whitening) match the performance of the SVD-based filter. It is the specific directions coming from $W_U$ that matter. Likewise, they compare against some popular embedding calibration tricks (if any exist for LLMs), and again EmbedFilter wins out. Finally, despite its simplicity, EmbedFilter seems robust across setups. The improvement holds for different LLM languages and architectures they tried, and for a variety of tasks in MTEB (semantic similarity, retrieval, clustering, etc.). The authors even note that performance scarcely degrades if you sweep the cutoff thresholds within a reasonable range. In short, it appears that almost any projection out of the high-frequency subspace helps a lot. Comparing to Related Approaches. EmbedFilter joins a growing toolkit of embedding calibration techniques. Historically, the term “isotropy” has often come up: many papers argue that embeddings (especially contextual ones) should be more spread out, and they propose various fixes like batch normalization or contrastive fine-tuning. In the static realm, Mu et al.’s “All-but-the-Top” (which iteratively removes the average and extra top PCs) is the classic case. In sentence embeddings, methods like Contrastive Learning or Siamese fine-tuning (SimCSE, SBERT) directly train the model to produce good vectors. Those approaches can be expensive or require labeled data or many positive/negative pairs. EmbedFilter is different. It is post-hoc and training-free. It doesn’t rely on any additional datasets or labels. It can be applied to any preexisting LLM without further tuning. So it doesn’t compete with fine-tuning; rather, it complements that work by offering a zero-cost sanity check. Another parallel is with TF-IDF in classical IR: weighting down frequent words. Conceptually, Filter is like applying an inverse-frequency filter, but on the hidden-feature level instead of on raw token counts. The high-freq token subspace is essentially a learned generalization of “stopword features,” and filtering it out is akin to giving common words lower weight in the embedding. Finally, the paper adds to a theme in recent research that the architecture itself might not be fully aligned with all use-cases. As a recent commentary puts it, the unembedding matrix “was designed for token prediction, not semantic representation,” so using it as a feature lens is a structural mismatch. SubFit (another concurrent work) for instance argues that different parts of the network can be fine-tuned or replaced in isolation. EmbedFilter essentially surgically removes the culprit subspace. This echoes findings in Mojumder and et al’s layered compression: targeted fixes often outdo wholesale methods. The key moral is: we shouldn’t assume the off-the-shelf LLM hides perfect embeddings; we may need small surgical tweaks like this. Practical Implications and Future Directions. For engineers and researchers building retrieval or semantic search on LLM embeddings, this result has immediate consequences. It suggests that inserting EmbedFilter is a low-risk, high-reward calibration step. Without retraining, you can often achieve higher recall and precision in search by simply filtering out those frequent-word directions. It also means that LLMs of a given size can potentially match the performance of specialized embedding models, as long as you handle the decoder bias. The ModelWire summary even speculates that if these gains hold up, it shows the “unembedding matrix is the primary bottleneck”, bridging much of the gap between generative LLMs and purpose-built embed generators. More broadly, this work invites a reconsideration of LLM embeddings. Designers of future models may explicitly incorporate orthogonality or disentanglement constraints so that the final embedding manifold isn’t swamped by frequency. One could imagine training objectives that penalize this effect, or architectural variants where part of the hidden space is dedicated to semantics and another to language modeling. The authors note that understanding the mechanism in the unembedding opens the door to “more principled designs” for text embeddings. Some questions remain open. Is the filtered space always the optimal one, or might adaptively learning $\Phi$ via a small amount of data offer further gains? The paper suggests not needing any training, but future work could explore minimal fine-tuning (e.g. learning the cutoff threshold $\tau$ on a validation set). Also, how does this interact with different ways of pooling (mean vs cls token)? The authors applied EmbedFilter to the pooled final layer, but one could similarly apply it at each token-level vector before pooling – conceivably refining each token’s embedding. Finally, we should watch how this technique competes with emerging embedding models. If an open LLM + EmbedFilter matches or surpasses proprietary embeddings, it can change deployment decisions. The authors already released code (EmbFilter on GitHub) so practitioners can try it on their favorite models. Over time, we may see “embedding filter” become a standard step like applying PCA or normalization. In summary, Your UnEmbedding Matrix is Secretly a Feature Lens for Text Embeddings demonstrates a clever hack with real impact: by turning the hidden weight matrix into a filter, it turns garbage (frequency noise) into signal. It underscores that even after the revolution of scaling up LLMs, there is still room for simple linear tricks to significantly improve performance. LLM embedding quality, it seems, wasn’t just about model size – sometimes it’s about the matrix algebra under the hood. Sources: The key findings and method descriptions come from Wu et al.’s 2026 preprint. They build on concepts like the logit lens and relate to earlier projection-based fixes for embeddings. The entire EmbedFilter idea is detailed in the original paper and supported by its experiments. (All citations above refer to the referenced sources.).