REZE: Representation Regularization for Domain-adaptive Text Embedding Pre-finetuning

1. Problem Statement

The paper addresses a critical challenge in adapting general-purpose text embedding models to specialized domains (e.g., finance, code, chemistry).

• The Context: Standard practice involves Pre-finetuning (PFT), where a base model is trained on a collection of heterogeneous, domain-specific tasks (classification, retrieval, NLI, etc.) before final fine-tuning on a specific target task.
• The Issue: Naive aggregation of these scattered, heterogeneous tasks introduces task-induced bias. While the model learns domain knowledge, it also learns conflicting task-specific artifacts.
• The Consequence: This leads to uncontrolled representation shifts that distort the pretrained embedding geometry. Instead of improving performance, naive PFT often causes negative transfer, resulting in performance degradation compared to direct fine-tuning (FT) on the target task, particularly in low-resource settings.
• Limitations of Existing Solutions: Current methods to fix representation issues (like isotropy-oriented post-processing: whitening, normalizing flows) are typically applied after training. The paper argues these are insufficient for PFT because they manipulate already-distorted spaces or rely on limited data statistics that fail to capture the true distribution in specialized domains.

2. Methodology: REZE

The authors propose REZE (Representation Regularization), a framework designed to control representation shifts during the pre-finetuning phase.

Core Insight

REZE operates on relation representations (concatenations of anchor-positive pairs) rather than individual embeddings. It posits that while generalizable domain knowledge is shared across tasks, task-specific biases manifest as outliers in the geometric structure of these relations.

Technical Workflow

1. Reference Eigenspace Construction (Offline):

   • Before pre-finetuning begins, a reference model ($f_0$) processes all source datasets to create a pooled set of relation vectors.
   • Eigenvalue Decomposition (EVD): The covariance matrix of these relations is decomposed to define an eigenspace.
   • Task-Variance Detection: For each source dataset, the mean in the eigenspace is calculated. A robust median-based dispersion metric is used to identify dimensions where datasets deviate significantly from the global consensus. These are identified as "task-variant" directions (noise).
   • Active Dimensions: Only the top eigendimensions explaining a high cumulative variance (e.g., 99%) are considered to avoid suppressing fine-grained signals.

2. Adaptive Soft-Shrinkage Mechanism:

   • REZE calculates a robust global threshold using the Median Absolute Deviation (MAD) of the task variances.
   • If a specific source dataset's mean deviates beyond a calculated band in a task-variant dimension, an adaptive soft-shrinkage factor ($\alpha$) is applied.
   • Unlike hard thresholding (which zeros out dimensions), soft-shrinkage gently pulls task-specific deviations back toward the global consensus, preserving the fundamental relational structure.
   • These shrinkage matrices are pre-computed and remain fixed during training, incurring zero inference-time overhead.

3. Debiasing During Pre-finetuning (Online):

   • During training, for every batch, the model generates a relation vector $r_i$.
   • The REZE regularization term ($L_{reze}$) is computed as the cosine dissimilarity between the current model's relation and a debiased target ($\hat{r}^{(0)}_i$).
   • The debiased target is derived by applying the pre-computed shrinkage matrices to the reference model's output.
   • Final Objective: $L = L_{main} + \alpha \cdot L_{reze}$, where $L_{main}$ is the standard InfoNCE contrastive loss.

3. Key Contributions

• Novel Framework: Introduction of REZE, the first method to explicitly control representation shifts during heterogeneous pre-finetuning using eigenspace decomposition and adaptive soft-shrinkage.
• Robust Statistics: Utilization of median-based metrics and MAD instead of mean/variance to robustly identify task-variant directions, preventing outlier datasets from skewing the regularization.
• Zero Overhead: The regularization parameters are pre-computed offline, meaning the method adds no computational cost during the actual training or inference steps.
• Theoretical Insight: Demonstrates that controlling representation shift is more critical than simply enforcing isotropy (uniformity) for robust domain adaptation.

4. Experimental Results

The authors evaluated REZE across three specialized benchmarks (FinMTEB, Code (MTEB), ChemTEB) using four diverse backbone models (E5, ModernBERT, GTE, Qwen3).

• Performance Superiority: REZE consistently outperformed standard Pre-finetuning (PFT) and direct Fine-tuning (FT) in most settings.
  • In many cases, naive PFT performed worse than FT due to negative transfer. REZE reversed this trend, absorbing domain knowledge while suppressing noise.
  • Stability: In challenging scenarios (e.g., Qwen3 on Code tasks), standard PFT variants collapsed (scores near 0), while REZE maintained robust performance.
• Comparison with Post-processing: Methods like PFT+Whitening and PFT+NormalizingFlow often degraded performance, particularly in specialized domains with limited data, because they distorted the spatial information of the already-fitted models.
• Isotropy Analysis: REZE significantly improved IsoScore (a measure of embedding uniformity) by 2x to 3x compared to standard PFT, confirming it creates a more balanced embedding space without forcing artificial isotropy.
• Ablation Studies:
  • Regularization Weight ($\alpha$): Performance was stable for $\alpha \ge 0.5$, with $\alpha=1.0$ being optimal.
  • Median vs. Mean: Using median aggregation for dispersion estimation proved superior to mean aggregation, confirming the need for robustness against outliers.

5. Significance

• Paradigm Shift: The paper challenges the notion that simply adding more domain data via multi-task learning is sufficient. It highlights that heterogeneous supervision requires explicit bias control.
• Practical Impact: REZE provides a plug-and-play solution for organizations adapting LLMs/embeddings to niche domains (Finance, Law, Medicine) where data is scarce and tasks are diverse.
• Geometric Preservation: By showing that "controlled shifts" aligned with the original manifold are better than "uncontrolled expansion" or "forced isotropy," the paper offers a new principle for designing robust embedding pipelines.

In summary, REZE solves the problem of negative transfer in domain-adaptive pre-finetuning by mathematically identifying and suppressing task-specific noise in the eigenspace, resulting in embeddings that are both domain-aware and geometrically stable.