Whitening Sentence Representations

• Whitening sentence representations is a geometric transformation that eliminates statistical dependencies and scales embedding dimensions to achieve an isotropic distribution.
• Techniques such as ZCA, group whitening, and kernel whitening decorrelate embeddings, addressing issues like high cosine similarities and hubness.
• Applied as a post-processing step or integrated in training, whitening boosts semantic similarity and retrieval accuracy while sometimes reducing classification performance.

Whitening sentence representations is a geometric transformation applied to sentence embeddings that aims to eliminate statistical dependencies and scale disparities among embedding dimensions, thereby inducing an isotropic distribution. This transformation is motivated by the empirical observation that pretrained representations from models such as BERT exhibit severe anisotropy—most embeddings occupy a narrow cone or subspace, with disproportionately high pairwise cosine similarity and a few principal axes capturing most variance. Whitening counteracts this phenomenon by enforcing zero mean and unit covariance, enhancing the utility of embeddings in tasks dependent on similarity or uniform coverage of the semantic space.

1. The Anisotropy Problem in Sentence Embeddings

Pretrained LLMs, including BERT and its derivatives, generate dense sentence vectors with highly non-uniform spatial geometry. The bulk of their variance is concentrated along a minimal subset of dimensions, resulting in embeddings whose average pairwise cosine similarity may exceed 0.8–0.99 (Su et al., 2021, Zhang et al., 2024). Such anisotropy diminishes the discriminative power of cosine similarity, causing most distinct sentences to cluster along a subspace and reducing retrieval and ranking quality. High anisotropy also amplifies spurious correlations learned from dataset biases, undermining generalization in out-of-distribution settings (Gao et al., 2022). Hubness—the phenomenon whereby certain points become nearest neighbors ("hubs") for a disproportionately high fraction of other points—further disrupts the geometry, yielding asymmetric neighborhood structure that complicates kNN classification and retrieval (Nielsen et al., 2023).

2. Whitening Transformations: Linear, Grouped, and Nonlinear Variants

Linear Whitening

The classical whitening transformation involves centering embeddings and applying a decorrelating linear map:

• Given embeddings $X \in \mathbb{R}^{N\times d}$, compute mean $\mu$ and covariance $\Sigma$.
• Eigendecompose $\Sigma = U\Lambda U^\top$.
• Apply the ZCA transformation $W = U\Lambda^{-1/2}U^\top$, yielding whitened vectors $x'_i = W(x_i - \mu)$, with $Cov(\{x'_i\}) = I$ (Huang et al., 2021, Zhang et al., 2024, Su et al., 2021).

Group and Rank-based Whitening

To balance isotropy with local semantic structure, group whitening partitions the embedding into $G$ groups (of dimension $d/G$), independently decorrelating and standardizing each group before reassembling the full vector. PyTorch-style pseudocode and detailed steps are given in (Zhang et al., 2024). A related marginal approach (f-norm) applies rank-based normalization per-dimension and projects to the unit sphere, ensuring each coordinate marginal is $\mathcal{N}(0,1)$ but without decorrelating across dimensions (Nielsen et al., 2023).

Kernel Whitening

Linear whitening fails to address nonlinear dependencies. Kernel-Whitening maps embeddings into a reproducing kernel Hilbert space (RKHS) using a positive-definite kernel $\mu$0, applies the whitening transform via a low-rank Nyström approximation, and uses the resulting kernel features directly (Gao et al., 2022).

Comparison Table

Whitening Type	Decorrelates All Dimensions?	Nonlinear Dependencies?
ZCA/PCA Whitening	Yes	No
Group Whitening	Partly (within group)	No
Marginal (f-norm)	No (per-dim only)	No
Kernel Whitening	Yes (in RKHS)	Yes

3. Integration with Learning Objectives and Models

As a Post-processing Step

Whitening is typically applied post hoc to off-the-shelf embeddings (e.g. average pooled token vectors from BERT or Sentence-BERT), yielding immediate gains in isotropy, as shown in WhiteningBERT (Huang et al., 2021) and related work (Su et al., 2021). For instance, classic whitening increases average Spearman's $\mu$1 correlation on seven STS tasks from 62.97 to 67.76 using BERT-base (Huang et al., 2021).

Joint Architectures

Incorporating whitening within training frameworks—such as with contrastive losses—can yield further improvements. Shuffled Group Whitening (SGW) introduces random group assignments of features across batch samples, computes per-group whitening, and uses multiple whitened “views” per sentence as positive samples, benefiting both alignment and uniformity in contrastive learning (Zhuo et al., 2023). This results in state-of-the-art unsupervised sentence embeddings, e.g., raising SimCSE-BERT performance from 76.25% to 78.78% average Spearman’s correlation on STS tasks. Kernel-Whitening integrates feature whitening into fine-tuning using kernel PCA in RKHS, significantly enhancing out-of-distribution robustness without harming in-distribution task performance (Gao et al., 2022).

Hubness Mitigation and Retrieval

Marginal whitening combined with Mutual Proximity rescaling specifically targets hubness, transforming both global distribution and neighborhood structure to dramatically reduce hubness metrics (e.g., $\mu$2-skewness reduced by 69–83%) and improve 10-NN classification error on standard text benchmarks (Nielsen et al., 2023).

4. Empirical Impact: Semantic Similarity, Retrieval, Classification

Semantic Textual Similarity

Whitening consistently improves performance of unrefined BERT-family embeddings on semantic similarity tasks (STS), yielding average absolute gain of +4 to +9 Spearman $\mu$3 points across STS12-16, STS-B, and SICK (Huang et al., 2021, Su et al., 2021). Removing dominant axes, the isotropic embedding distribution provides finer semantic gradation and higher correlation with human similarity judgments.

Retrieval and Storage/Latency

Dimension reduction via whitening (retaining only the top $\mu$4 principal components post-decorrelation) yields significant compression: reducing dimensionality from 768 to 256 maintains nearly all semantic similarity performance with >66% storage and retrieval speedup (Su et al., 2021). For large repositories and online services, this enables scalable approximate nearest-neighbor retrieval.

Out-of-Distribution Robustness and Recommendation

Kernel-Whitening specifically targets elimination of spurious, dataset-specific correlations. It achieves >9% absolute accuracy gains on out-of-distribution generalization tests while holding in-domain accuracy constant (Gao et al., 2022). In sequential recommendation, whitening reduces the average cosine similarity of item representations from ~0.85 to near zero, enabling item-only models to achieve competitive Recall@20 and cold-start performance relative to ID-augmented models (Zhang et al., 2024).

Classification

In contrast, whitening is detrimental for direct classification (especially on LLM embeddings). Across a suite of text classification tasks (MR, CR, SUBJ, etc.), ZCA, PCA, and Cholesky whitening all yield notable accuracy losses, with magnitude increasing with embedding dimensionality. The isotropy enforced by whitening destroys anisotropic, discriminative directions required for class separation. For example, accuracy on LLaMA2 embeddings falls from 81.39% to 66.40% after whitening (Forooghi et al., 2024). A plausible implication is that whitening should be strictly avoided for classification-centric workflows unless embeddings are explicitly optimized for such downstream adaptation.

5. Design Trade-offs, Limitations, and Best Practices

Balancing Isotropy and Semantic Preservation

Full whitening maximizes global uniformity but can distort local manifold structure and semantic relations, especially when used as a drop-in for downstream tasks other than similarity search or retrieval (Zhang et al., 2024). Group whitening and ensemble approaches (WhitenRec+) capture a trade-off: group-wise decorrelation maintains some original feature relationships, and ensembling (summing projected full/group whitened vectors) preserves more information for workflows such as recommendation (Zhang et al., 2024). Theoretical analysis shows that WhitenRec+ retains $\mu$5 more "degrees of freedom" than full whitening alone.

Nonlinear Whitening and Bias Mitigation

Linear whitening cannot remove higher-order statistical dependencies. Kernel-Whitening addresses this limitation by operating in an RKHS, achieving deeper debiasing while retaining tractable training and inference cost via the Nyström method (Gao et al., 2022). It is effective for mitigating nonlinear spurious correlations and confers strong out-of-distribution robustness.

Hyperparameter Sensitivity

Group size, number of positive augmentations, whitening output dimension, and batch size govern the isotropy/semantic preservation trade-off. Empirically, group sizes of 256–384 (for $\mu$6) and $\mu$7 positive augmentations maximize alignment and uniformity benefits in contrastive whitening (Zhuo et al., 2023). Full-rank whitening remains optimal for similarity and retrieval, but group or kernel whitening may be preferred for debiased modeling.

Task and Model Dependency

Whitening is advantageous only for unsupervised similarity and retrieval tasks and can consistently harm classification accuracy, particularly for embeddings from large LLMs, where class-separating structure is encoded in the original anisotropic feature geometry (Forooghi et al., 2024). Whitening may be beneficial for base or non-semantic-fine-tuned models but less so for models such as SBERT or SimCSE already specialized for semantic similarity (Forooghi et al., 2024).

6. Summary of State-of-the-Art and Future Directions

Whitening, in its linear, grouped, rank-based, and kernel forms, provides a principled framework for correcting geometric pathologies of pretrained sentence embeddings, mainly anisotropy and hubness. Its greatest utility is observed in tasks reliant on nearest-neighbor or similarity computations, where it improves isotropy, retrieval accuracy, and semantic gradation. Recent advances—including SGW-augmented contrastive learning (Zhuo et al., 2023), probabilistic reweighting for hubness (Nielsen et al., 2023), and kernelized whitening for debiasing (Gao et al., 2022)—extend these gains to broader modeling contexts. Limitations remain where downstream applications depend on anisotropic, discriminative feature directions (notably classification), requiring selective, task-aware application of whitening and motivating hybrid or ensemble approaches (e.g., WhitenRec+) (Zhang et al., 2024). Ongoing research explores integration of whitening within end-to-end learning and data-dependent schemes for dynamic trade-off adjustment, as well as extending these methods to very large-scale, heterogeneous, or multilingual embedding scenarios.