LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures

Motivation

Autoregressive LLMs are strong generators but their representations may be suboptimal because next-token prediction is a local objective. Joint Embedding Predictive Architectures (JEPA), inspired by LeCun’s energy-based framework, train representations by predicting embeddings rather than tokens — potentially capturing higher-level semantic structure.

Approach

LLM-JEPA combines two training objectives:

1. Standard next-token prediction (preserving generation capability)
2. JEPA-style embedding loss: given paired inputs (e.g., natural language + regular expression for the same concept), predict the embedding of one view from the other

A [PRED] token implements the predictor using the LLM’s own weights, keeping architecture unchanged. Separate forward passes for each view prevent attention from mixing information between views.

Results

Gains across model families (Llama 3, Gemma 2, OpenELM, OLMo) on:

• Natural language → regex conversion
• Mathematical reasoning (GSM8K)
• SQL generation (Spider)
• Sentiment analysis in pre-training experiments

Trade-offs

• Training cost triples: each view pair requires additional forward passes
• Requires datasets with naturally aligned multiple views — not universally applicable