Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

BERT-JEPA: Reorganizing CLS Embeddings for Language-Invariant Semantics

Taj Gillin, Adam Lalani, Kenneth Zhang, Marcel Mateos Salles

TL;DR

This work tackles CLS collapse and language-specific entanglements in BERT-style models by introducing BEPA, a training paradigm that adds a JEPA-style CLS alignment objective to MLM. By packing two sentences (monolingual or bilingual) and enforcing cross-language CLS alignment with an InfoNCE-based loss $L_{Alignment}$ alongside the standard MLM objective, BEPA reshapes the CLS embedding space into a language-invariant semantic space. The approach yields a fuller-rank variance distribution, reduced CLS collapse, and improved cross-lingual transfer on benchmarks such as XNLI and MLQA, while preserving English GLUE performance. Practically, BEPA demonstrates that joint MLM and embedding-space alignment can enhance multilingual understanding without heavy English-task degradation, and the authors provide code and resources for replication.

Abstract

Joint Embedding Predictive Architectures (JEPA) are a novel self supervised training technique that have shown recent promise across domains. We introduce BERT-JEPA (BEPA), a training paradigm that adds a JEPA training objective to BERT-style models, working to combat a collapsed [CLS] embedding space and turning it into a language-agnostic space. This new structure leads to increased performance across multilingual benchmarks.

BERT-JEPA: Reorganizing CLS Embeddings for Language-Invariant Semantics

TL;DR

This work tackles CLS collapse and language-specific entanglements in BERT-style models by introducing BEPA, a training paradigm that adds a JEPA-style CLS alignment objective to MLM. By packing two sentences (monolingual or bilingual) and enforcing cross-language CLS alignment with an InfoNCE-based loss alongside the standard MLM objective, BEPA reshapes the CLS embedding space into a language-invariant semantic space. The approach yields a fuller-rank variance distribution, reduced CLS collapse, and improved cross-lingual transfer on benchmarks such as XNLI and MLQA, while preserving English GLUE performance. Practically, BEPA demonstrates that joint MLM and embedding-space alignment can enhance multilingual understanding without heavy English-task degradation, and the authors provide code and resources for replication.

Abstract

Joint Embedding Predictive Architectures (JEPA) are a novel self supervised training technique that have shown recent promise across domains. We introduce BERT-JEPA (BEPA), a training paradigm that adds a JEPA training objective to BERT-style models, working to combat a collapsed [CLS] embedding space and turning it into a language-agnostic space. This new structure leads to increased performance across multilingual benchmarks.
Paper Structure (20 sections, 2 equations, 10 figures, 10 tables)

This paper contains 20 sections, 2 equations, 10 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Method: A Novel JEPA-style BERT Finetuning Framework
  4. BEPA Framework
  5. BEPA Loss
  6. BEPA Training Semantically Reorganizes the CLS Embedding Space
  7. BEPA Training Semantically Reorganizes the CLS Embedding Space
  8. Semantically Organized Token Spaces Improve Cross-Lingual Transfer
  9. Conclusion
  10. Appendix
  11. Contributions
  12. GenAI Use
  13. Resources Used
  14. BEPA Repository
  15. The Full BEPA Framework
  16. ...and 5 more sections

Figures (10)

  • Figure 1: PCA plot showing the variance per component of the [CLS] embedding space. BEPA finetuning increases (↑) the number of high variance components. Meanwhile, standard RoBERTa and XLM-RoBERTa display the majority of variance, 47% and 78% respectively, being displayed by a single component. Compared to 34% in the first component for both BEPA models.
  • Figure 2: (Left): t-SNE plot for RoBERTa, samples are tightly distributed by language, with English distinctly unique in its large span. (Middle): t-SNE plot for XLM-RoBERTa, some languages are isolated, similar to RoBERTa, though less tightly packed. Many languages, especially English, span a much larger space and share a space with multiple languages. (Right): t-SNE plot for BEPA Bilingual, languages are evenly distributed in overlapping. BEPA creates a shared and aligned space between languages. Results for our finetuned BEPA Monolingual model can be found in \ref{['tsne_plot_mono']}.
  • Figure 3: Performance G-XLT matrix on MLQA benchmark. F1 scores across 49 language pairs (7 context languages × 7 question languages). Diagonal entries represent monolingual performance, while off-diagonal entries show cross-lingual transfer. (Top): XLM-RoBERTa Base shows strong performance on monolingual tasks but struggles significantly with cross-lingual scenarios, particularly those not involving English. (Bottom): BEPA Bilingual demonstrate improved cross-lingual transfer, with more uniform performance across language pairs, while maintaining strong monolingual performance. This indicates BEPA's learns token-level semantic alignment in the representation space.
  • Figure 4: Input packing strategy for BEPA. Two sequences are concatenated using two [CLS] tokens and a [SEP] boundary. $x^{i}_1$ represents tokens corresponding to the first sequence while $x^{i}_2$ represents tokens corresponding to the second sequence. Segment A marks the first sentence and Segment B the second. The entire sequence is encoded jointly.
  • Figure 5: Masking strategy for CLS-token JEPA alignment stage of the BEPA framework. (Left): Masking out tokens corresponding to the second sentence to allow the for the model to build a [CLS] token representation of the first sequence. (Right): Masking out tokens corresponding to the first sentence to allow for the model to build a [CLS] token representation of the second sequence.
  • ...and 5 more figures