Tag: ml

JEPA is an example of a predictive coding architecture. Predictive coding is the idea that the brain works by predicting the next sensory input and learns from the error between the prediction and the actual input. Hierarchies of this kind of prediction allow higher level elements to predict the outputs of lower-level elements, building up deeper and more complex semantics.

The core idea in JEPA is to take two related things (say consecutive frames of a video) x and y, encode each of them into a latent space (an embedding), and then predict the embedding for y s(y) based on s(x). The encoders can be Transformer-based models — in practice models like I-JEPA have trained the x encoder and updated the weights via a moving average for the y (target) encoder.

The learning is not based on how well the end-result predicts the target e.g. how close the pixels of the next frame are predicted. Instead, it’s based on how well the latent representation of the next frame is predicted.

The advantage of working in the latent space for the prediction is the model can choose what level of detail it wants to capture, discarding some aspects and focusing on more foundational concepts. This helps build a more robust world-model, with the hope being that training in this way will then allow easier generalization to more tasks, with less data required

Similarities

This is somewhat similar to autoencoders. Autoencoders take an input, compress it in a latent space, then reconstruct the original from the latent space and propagate back the error. JEPA does a similar process across two different items with separate encoders, and only cares about error within the latent space.

Contrastive models embed two different items into the same space and try to increase similarity between the embeddings for things known to be similar and make them dissimilar to other items. This is used in CLIP and other multimodal text-image encoders, where the text and the image embed to the same space so that a text caption and a matching image are close in embedding space. This requires a lot of pairwise comparisons, while JEPA is a more straightforward s(x)->s(y) prediction in training.

Challenges

Because JEPA models leave you with a latent they need to be paired with a generator for getting an observable/human viewable output, which is a per-domain challenge. This makes it harder to evaluate how well the model is learning, beyond measuring loss.

Training stability can also be tricky — it is possible for the model to collapse and learn trivial representations to minimize prediction error. Even without complete collapse it can require some experimentation to ensure the model is learning a deep enough conceptual level. For example, I-JEPA, which worked in image space, found that using large enough masked patches was important to ensure the model captured sufficient detail.

https://ai.pydantic.dev/evals/

Ed Yang was recently recommending keeping your own benchmark of LLM evals, so you can test newer models on problems that they have struggled with in the past. I have recommended similar things to people, but there is some barrier to entry into knowing how to start. Ed references (and forks) Nicolas Carlini’s personal benchmark repo, but its nice to have some light(ish) weight options too.

Pydantic Evals is a powerful evaluation framework designed to help you systematically test and evaluate the performance and accuracy of the systems you build, especially when working with LLMs.

You can install the library with uv or pip:

uv add pydantic-evals

I tried it out with a strawberry test, calling openrouter with different models. I needed a custom eval as the default Contains is a bit rigid, but the approach seems nice!

import os
import asyncio
from dataclasses import dataclass
from pydantic_evals import Case, Dataset
from pydantic_evals.evaluators import Evaluator, EvaluationReason, EvaluatorContext
from pydantic_evals.evaluators.common import _truncated_repr
from openai import OpenAI
from typing import Any, Optional, cast
@dataclass
class FlexibleContains(Evaluator[object, object, object]):
    """
    Check if the output contains any one of the expected options.
    """
    value: Any
    case_sensitive: bool = False
    def evaluate(
        self, ctx: EvaluatorContext[object, object, object]
    ) -> EvaluationReason:
        failure_reason: Optional[str] = None
        # Normalize value into a list of options if it isn't already a list or tuple.
        options = self.value if isinstance(self.value, (list, tuple)) else [self.value]
        output_str = str(ctx.output)
        if not self.case_sensitive:
            output_str = output_str.lower()
        match_found = False
        for opt in options:
            opt_str = str(opt)
            if not self.case_sensitive:
                opt_str = opt_str.lower()
            if opt_str in output_str:
                match_found = True
                break
        if not match_found:
            failure_reason = (
                f"Output string {_truncated_repr(output_str, max_length=100)} does not contain "
                f"any of expected strings: {[str(opt) for opt in options]}"
            )
        return EvaluationReason(value=match_found, reason=failure_reason)
strawberry = Case(
    name="strawberry",
    inputs="How many rs are in strawberry?",
    evaluators=[FlexibleContains(value=["3", "three"])],
    metadata={"difficulty": "easy"},
)
dataset = Dataset(cases=[strawberry])
MODELS = [
    "anthropic/claude-3.5-sonnet",
    "openai/gpt-4o",
    "meta-llama/llama-4-maverick:free",
    "meta-llama/llama-4-scout:free",
    "openrouter/optimus-alpha",  # secret model!
]
def generate_completion(inputs: str, model: str) -> str:
    """Generate a completion using OpenRouter with specified model"""
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=os.getenv("OPENROUTER_API_KEY"),
    )
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": inputs},
        ],
        max_tokens=50,
        temperature=0.7,
    )
    return response.choices[0].message.content.strip()
def evaluate_models():
    """Run evaluations across multiple models"""
    for model in MODELS:
        print(f"\nResults for model: {model}")
        print("=" * 50)
        # Wrap the synchronous generate_completion in an async function:
        async def model_specific_generate(inputs: str) -> str:
            loop = asyncio.get_running_loop()
            return await loop.run_in_executor(None, generate_completion, inputs, model)
        # Run evaluation for this model
        report = dataset.evaluate_sync(model_specific_generate)
        # Print results for this model
        report.print(include_input=True, include_output=True, include_durations=False)
def main():
    evaluate_models()
if __name__ == "__main__":
    main()

To give a trimmed output:

Results for model: openrouter/optimus-alpha
==================================================
                                      Evaluation Summary: model_specific_generate
┏━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Case ID    ┃ Inputs                         ┃ Outputs                                                   ┃ Assertions ┃
┡━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ strawberry │ How many rs are in strawberry? │ The word **"strawberry"** contains **three** letter "r"s. │ ✔          │
├────────────┼────────────────────────────────┼───────────────────────────────────────────────────────────┼────────────┤
│ Averages   │                                │                                                           │ 100.0% ✔   │
└────────────┴────────────────────────────────┴───────────────────────────────────────────────────────────┴────────────┘