# JEPA as Categorical World Model FunctorFlow.jl - [Introduction](#introduction) - [Part 1: Coalgebra — World Models as State Transitions](#part-1-coalgebra--world-models-as-state-transitions) - [Part 2: JEPA Block — Prediction in Embedding Space](#part-2-jepa-block--prediction-in-embedding-space) - [Part 3: JEPA Execution — The Obstruction Loss](#part-3-jepa-execution--the-obstruction-loss) - [Part 4: KAN-JEPA — Merging KET and JEPA](#part-4-kan-jepa--merging-ket-and-jepa) - [Part 5: Hierarchical JEPA (H-JEPA)](#part-5-hierarchical-jepa-h-jepa) - [Part 6: Energy-Based Cost Module](#part-6-energy-based-cost-module) - [Energy Function Implementations](#energy-function-implementations) - [VICReg Regularization](#vicreg-regularization) - [Part 7: Lean 4 Proof Certificates](#part-7-lean-4-proof-certificates) - [Part 8: EMA Update (Collapse Prevention)](#part-8-ema-update-collapse-prevention) - [Summary](#summary) ## Introduction **JEPA** (Joint Embedding Predictive Architecture) is LeCun’s proposed framework for building world models that predict in *representation space* rather than *observation space*. The key insight is that predicting pixel-by-pixel is intractable, but predicting *embeddings* of future states is feasible. FunctorFlow.jl reveals that JEPA is fundamentally a **coalgebraic construction**: | JEPA Component | Categorical Analog | |------------------------|----------------------------------| | Encoder | Coalgebra morphism | | Predictor | Endofunctor dynamics | | Target encoder (EMA) | Frozen reference coalgebra | | Prediction loss | Obstruction to commutativity | | World model | F-coalgebra (X → F(X)) | | Optimal representation | Final coalgebra (Lambek’s lemma) | ``` julia using FunctorFlow ``` ## Part 1: Coalgebra — World Models as State Transitions A **coalgebra** is a pair (X, α : X → F(X)): a state space equipped with a transition structure. This is the categorical formalization of a world model. ``` julia # Build a simple world model diagram D = world_model_block(; name=:SimpleWorldModel, observation_object=:Pixels, latent_object=:Embedding, encoder_name=:encode, dynamics_name=:predict_next, decoder_name=:decode, ) println("Objects: ", join(keys(D.objects), ", ")) println("Operations: ", join(keys(D.operations), ", ")) println() # The coalgebra structure coalgebras = get_coalgebras(D) for (name, c) in coalgebras println(c) end ``` Objects: Pixels, Embedding Operations: encode, predict_next, decode, encode_then_predict, autoencoder Coalgebra(:coalgebra, Embedding →_{identity} Embedding) The coalgebra declares that `predict_next` is the world model dynamics: given a latent state, it produces the next latent state. ``` julia # Bind concrete implementations bind_morphism!(D, :encode, x -> x ./ 255.0) # normalize bind_morphism!(D, :predict_next, x -> x .+ 0.01) # simple drift bind_morphism!(D, :decode, x -> x .* 255.0) # denormalize compiled = compile_to_callable(D) result = FunctorFlow.run(compiled, Dict( :Pixels => [128.0, 64.0, 200.0], )) println("Encoded: ", round.(result.values[:encode]; digits=4)) println("Predicted:", round.(result.values[:predict_next]; digits=4)) println("Decoded: ", round.(result.values[:autoencoder]; digits=4)) ``` Encoded: [0.502, 0.251, 0.7843] Predicted:[0.512, 0.261, 0.7943] Decoded: [128.0, 64.0, 200.0] ## Part 2: JEPA Block — Prediction in Embedding Space The JEPA block encodes the fundamental prediction-in-embedding pattern as a FunctorFlow diagram: Observation ──encoder_ctx──→ ContextRepr ──predictor──→ PredictedRepr Target ────── encoder_tgt──→ TargetRepr ↕ prediction_loss (obstruction) ``` julia # Build a JEPA block D_jepa = jepa_block(; name=:ImageJEPA, observation_object=:Context, target_object=:MaskedRegion, context_repr=:CtxEmb, target_repr=:TgtEmb, comparator=:l2, ) println("=== JEPA Diagram: $(D_jepa.name) ===") println("\nObjects:") for (name, obj) in D_jepa.objects println(" $(name) :: $(obj.kind)") end println("\nMorphisms:") for (name, op) in D_jepa.operations if op isa Morphism role = get(op.metadata, :role, :unknown) net = get(op.metadata, :network, :none) detach = get(op.metadata, :detach, false) println(" $(name): $(op.source) -> $(op.target) [role=$(role), net=$(net), detach=$(detach)]") elseif op isa Composition println(" $(name): $(op.source) → $(op.target) [chain=$(op.chain)]") end end println("\nLosses:") for (name, loss) in D_jepa.losses println(" $(name): paths=$(loss.paths), comparator=$(loss.comparator)") end println("\nCoalgebra:") for (name, c) in get_coalgebras(D_jepa) println(" ", c) end ``` === JEPA Diagram: ImageJEPA === Objects: Context :: observation MaskedRegion :: observation CtxEmb :: representation TgtEmb :: representation Morphisms: context_encoder: Context -> CtxEmb [role=encoder, net=online, detach=false] target_encoder: MaskedRegion -> TgtEmb [role=encoder, net=target, detach=true] predictor: CtxEmb -> TgtEmb [role=predictor, net=none, detach=false] prediction_path: Context → TgtEmb [chain=[:context_encoder, :predictor]] Losses: prediction_loss: paths=[(:prediction_path, :target_encoder)], comparator=l2 Coalgebra: Coalgebra(:jepa_dynamics, CtxEmb →_{identity} CtxEmb) ## Part 3: JEPA Execution — The Obstruction Loss The prediction loss IS an obstruction loss: it measures how far the encoder-predictor square is from commuting. ``` julia # Bind implementations: encoders as linear scaling bind_morphism!(D_jepa, :context_encoder, x -> x .* 0.5) bind_morphism!(D_jepa, :target_encoder, x -> x .* 0.5) # same encoder (EMA would make it close) bind_morphism!(D_jepa, :predictor, identity) # trivial predictor compiled_jepa = compile_to_callable(D_jepa) # Case 1: Context and target are the same → perfect prediction result1 = FunctorFlow.run(compiled_jepa, Dict( :Context => [1.0, 2.0, 3.0], :MaskedRegion => [1.0, 2.0, 3.0], # same input )) println("Case 1 (same input): loss = ", round(result1.losses[:prediction_loss]; digits=6)) println(" Predicted: ", result1.values[:prediction_path]) println(" Target: ", result1.values[:target_encoder]) # Case 2: Different context and target → non-zero loss result2 = FunctorFlow.run(compiled_jepa, Dict( :Context => [1.0, 2.0, 3.0], :MaskedRegion => [4.0, 5.0, 6.0], # different! )) println("\nCase 2 (different input): loss = ", round(result2.losses[:prediction_loss]; digits=6)) println(" Predicted: ", result2.values[:prediction_path]) println(" Target: ", result2.values[:target_encoder]) ``` Case 1 (same input): loss = 0.0 Predicted: [0.5, 1.0, 1.5] Target: [0.5, 1.0, 1.5] Case 2 (different input): loss = 2.598076 Predicted: [0.5, 1.0, 1.5] Target: [2.0, 2.5, 3.0] When the loss is zero, the JEPA square commutes — the encoder is an **exact coalgebra morphism**. ## Part 4: KAN-JEPA — Merging KET and JEPA KAN-JEPA uses a left Kan extension (Σ) as the predictor, merging the KET attention pattern with JEPA’s embedding-space prediction: ``` julia D_kan = kan_jepa_block(; name=:KanJEPA_Demo, observation_object=:Tokens, target_object=:MaskedTokens, context_repr=:TokenEmb, target_repr=:MaskedEmb, relation_object=:Neighborhood, reducer=:mean, ) println("=== KAN-JEPA Diagram ===") for (name, op) in D_kan.operations if op isa KanExtension println(" Σ($(op.source); along=$(op.along)) → $(op.target) [$(op.reducer)]") elseif op isa Morphism role = get(op.metadata, :role, :unknown) println(" $(name): $(op.source) → $(op.target) [$(role)]") end end println("\nThis is the natural fusion:") println(" KET's attention = left-Kan aggregation (Σ)") println(" JEPA's loss = obstruction to commutativity") ``` === KAN-JEPA Diagram === context_encoder: Tokens → TokenEmb [encoder] target_encoder: MaskedTokens → MaskedEmb [encoder] Σ(TokenEmb; along=Neighborhood) → MaskedEmb [mean] This is the natural fusion: KET's attention = left-Kan aggregation (Σ) JEPA's loss = obstruction to commutativity ## Part 5: Hierarchical JEPA (H-JEPA) H-JEPA uses multiple abstraction levels for multi-scale prediction. Fine levels handle short-range details; coarse levels handle long-range abstract planning. ``` julia D_hjepa = hjepa_block(; name=:MultiScaleJEPA, levels=[:fine, :medium, :coarse], ) println("=== H-JEPA: 3 Abstraction Levels ===") println("\nObjects per level:") for level in [:fine, :medium, :coarse] obs = Symbol("$(level)__Obs_$(level)") repr = Symbol("$(level)__CtxRepr_$(level)") println(" $(level): $(obs) → $(repr)") end println("\nAbstraction morphisms:") for (name, op) in D_hjepa.operations if op isa Morphism && get(op.metadata, :role, nothing) == :abstraction from = op.metadata[:from_level] to = op.metadata[:to_level] println(" $(name): $(from) → $(to)") end end println("\nPorts:") for (name, p) in D_hjepa.ports println(" $(name) ($(p.direction)): $(p.ref)") end ``` === H-JEPA: 3 Abstraction Levels === Objects per level: fine: fine__Obs_fine → fine__CtxRepr_fine medium: medium__Obs_medium → medium__CtxRepr_medium coarse: coarse__Obs_coarse → coarse__CtxRepr_coarse Abstraction morphisms: abstract_fine_to_medium: fine → medium abstract_medium_to_coarse: medium → coarse Ports: input (INPUT): fine__Obs_fine fine_repr (OUTPUT): fine__CtxRepr_fine coarse_repr (OUTPUT): coarse__CtxRepr_coarse ## Part 6: Energy-Based Cost Module The energy function measures compatibility in representation space. The cost module decomposes into intrinsic (immutable) and trainable components: $$C(s) = \sum_i u_i \cdot \text{IC}_i(s) + \sum_j v_j \cdot \text{TC}_j(s)$$ ``` julia # Build an energy block with VICReg-style regularization D_energy = energy_block(; name=:JEPAEnergy, energy_type=:l2, variance_weight=0.5, covariance_weight=0.1, collapse_strategy=VICREG, ) cost_mods = get_cost_modules(D_energy) cm = cost_mods[:cost] println("=== Energy-Based Cost Module ===") println("Intrinsic costs:") for ic in cm.intrinsic_costs println(" ", ic) end println("\nCollapse prevention: VICReg (variance + covariance regularization)") ``` === Energy-Based Cost Module === Intrinsic costs: IntrinsicCost(:prediction_cost, type=prediction, weight=1.0) IntrinsicCost(:variance_cost, type=variance, weight=0.5) IntrinsicCost(:covariance_cost, type=covariance, weight=0.1) Collapse prevention: VICReg (variance + covariance regularization) ### Energy Function Implementations ``` julia x = [1.0, 0.5, -0.3, 0.8] y = [0.9, 0.6, -0.2, 0.7] println("L2 energy: ", round(energy_l2(x, y); digits=6)) println("Cosine energy: ", round(energy_cosine(x, y); digits=6)) println("Smooth L1: ", round(energy_smooth_l1(x, y); digits=6)) # Self-similarity should give zero energy println("\nSelf-similarity:") println(" L2(x, x) = ", energy_l2(x, x)) println(" Cosine(x, x) = ", round(energy_cosine(x, x); digits=10)) ``` L2 energy: 0.04 Cosine energy: 0.007994 Smooth L1: 0.02 Self-similarity: L2(x, x) = 0.0 Cosine(x, x) = 5.1e-9 ### VICReg Regularization ``` julia using Random Random.seed!(42) # Good representations: diverse (high variance, low correlation) Z_good = randn(4, 10) # 4 dims, 10 samples var_good = variance_regularization(Z_good; gamma=1.0) cov_good = covariance_regularization(Z_good) # Collapsed representations: all identical Z_bad = ones(4, 10) var_bad = variance_regularization(Z_bad; gamma=1.0) cov_bad = covariance_regularization(Z_bad) println("Diverse representations:") println(" Variance penalty: ", round(var_good; digits=4)) println(" Covariance penalty: ", round(cov_good; digits=4)) println("\nCollapsed representations:") println(" Variance penalty: ", round(var_bad; digits=4), " (high = BAD)") println(" Covariance penalty: ", round(cov_bad; digits=4)) ``` Diverse representations: Variance penalty: 0.4483 Covariance penalty: 0.7101 Collapsed representations: Variance penalty: 3.96 (high = BAD) Covariance penalty: 0.0 ## Part 7: Lean 4 Proof Certificates FunctorFlow.jl can generate Lean 4 proof certificates for JEPA diagrams, formally verifying the categorical structure: ``` julia D_proof = jepa_block(; name=:ProvedJEPA) add_energy_function!(D_proof, :compat; domain=[:ContextRepr, :TargetRepr], energy_type=:l2) add_bisimulation!(D_proof, :enc_equiv; coalgebra_a=:jepa_dynamics, coalgebra_b=:jepa_dynamics, relation=:predictor) cert = render_jepa_certificate(D_proof) println(cert) ``` -- Auto-generated by FunctorFlow.jl namespace FunctorFlowProofs.Generated.ProvedJEPA open FunctorFlowProofs in def exportedDiagram : DiagramDecl := { name := "ProvedJEPA", objects := ["Observation", "Target", "ContextRepr", "TargetRepr"], operations := [{ name := "context_encoder", kind := OperationKind.morphism, refs := ["Observation", "ContextRepr"] }, { name := "target_encoder", kind := OperationKind.morphism, refs := ["Target", "TargetRepr"] }, { name := "predictor", kind := OperationKind.morphism, refs := ["ContextRepr", "TargetRepr"] }, { name := "prediction_path", kind := OperationKind.composition, refs := ["context_encoder", "predictor", "Observation", "TargetRepr"] }], ports := [{ name := "context_input", ref := "Observation", kind := "object", portType := "observation", direction := "input" }, { name := "target_input", ref := "Target", kind := "object", portType := "observation", direction := "input" }, { name := "context_embedding", ref := "ContextRepr", kind := "object", portType := "representation", direction := "output" }, { name := "target_embedding", ref := "TargetRepr", kind := "object", portType := "representation", direction := "output" }, { name := "prediction", ref := "prediction_path", kind := "operation", portType := "representation", direction := "output" }, { name := "loss", ref := "prediction_loss", kind := "loss", portType := "loss", direction := "output" }] } def exportedArtifact : LoweringArtifact := { diagram := exportedDiagram, resolvedRefs := true, portsClosed := true } theorem exportedArtifact_checks : exportedArtifact.check = true := by native_decide theorem exportedArtifact_sound : exportedArtifact.Sound := LoweringArtifact.sound_of_check exportedArtifact_checks -- Coalgebra declarations def coalgebra_jepa_dynamics : CoalgebraDecl := { name := "jepa_dynamics", state := "ContextRepr", transition := "predictor", functorType := "identity" } -- JEPA prediction-as-obstruction theorems /-- The JEPA prediction loss is an obstruction to commutativity of the encoder/predictor square. When loss = 0, the square commutes and the encoder is an exact coalgebra morphism. -/ theorem jepa_prediction_loss_is_obstruction : exportedArtifact.lossIsObstruction "prediction_loss" := by exact LoweringArtifact.loss_obstruction_of_check exportedArtifact_checks /-- When the JEPA prediction loss is zero, the encoder-predictor path commutes with the target encoder path, making the encoder a coalgebra morphism (structure-preserving map). -/ theorem jepa_exact_implies_coalgebra_morphism (h : ∀ l ∈ exportedArtifact.losses, l.value = 0) : exportedArtifact.CoalgebraExact := by exact LoweringArtifact.coalgebra_exact_of_zero_loss h -- Bisimulation declarations def bisim_enc_equiv : BisimulationDecl := { name := "enc_equiv", coalgebraA := "jepa_dynamics", coalgebraB := "jepa_dynamics", relation := "predictor" } /-- Two coalgebras are bisimilar iff they map to the same element in the final coalgebra — behavioral equivalence. -/ theorem bisimilar_iff_final_coalgebra_equal (A B : CoalgebraDecl) (R : BisimulationDecl) (h : R.isBisimulation A B) : A.finalImage = B.finalImage := CoalgebraDecl.bisim_implies_final_eq h -- Energy function declarations def energy_compat : EnergyDecl := { name := "compat", domain := ["ContextRepr", "TargetRepr"], energyType := "l2" } /-- Energy functions are non-negative for L2 and cosine types. -/ theorem energy_nonneg (e : EnergyDecl) (h : e.energyType ∈ ["l2", "cosine"]) : 0 ≤ e.evaluate := by exact EnergyDecl.nonneg_of_standard h end FunctorFlowProofs.Generated.ProvedJEPA Key theorems generated: - **`jepa_prediction_loss_is_obstruction`**: The JEPA loss IS an obstruction to commutativity - **`jepa_exact_implies_coalgebra_morphism`**: Zero loss → encoder preserves structure - **`bisimilar_iff_final_coalgebra_equal`**: Bisimilar encoders are behaviorally equivalent - **`energy_nonneg`**: Energy functions are non-negative ## Part 8: EMA Update (Collapse Prevention) JEPA prevents representation collapse by using an exponential moving average (EMA) of the online encoder as the target encoder: ``` julia # Simulate EMA training dynamics target_params = [Float32[0.1, 0.2, 0.3]] online_params = [Float32[0.5, 0.6, 0.7]] println("Before EMA:") println(" Target: ", target_params[1]) println(" Online: ", online_params[1]) for step in 1:5 ema_update!(target_params, online_params; decay=0.99) end println("\nAfter 5 EMA steps (decay=0.99):") println(" Target: ", round.(target_params[1]; digits=4)) println(" Online: ", online_params[1], " (unchanged)") println("\nTarget slowly tracks online → prevents collapse without negatives") ``` Before EMA: Target: Float32[0.1, 0.2, 0.3] Online: Float32[0.5, 0.6, 0.7] After 5 EMA steps (decay=0.99): Target: Float32[0.1196, 0.2196, 0.3196] Online: Float32[0.5, 0.6, 0.7] (unchanged) Target slowly tracks online → prevents collapse without negatives ## Summary | Construction | FunctorFlow Type | Purpose | |----|----|----| | Coalgebra | `Coalgebra` | World model (state → F(state)) | | Coalgebra morphism | `CoalgebraMorphism` | Structure-preserving encoder | | Final coalgebra | `FinalCoalgebraWitness` | Optimal representation (Lambek) | | Bisimulation | `Bisimulation` | Behavioral equivalence of encoders | | JEPA block | `jepa_block()` | Encoder/predictor/target triple | | KAN-JEPA | `kan_jepa_block()` | Σ-attention as JEPA predictor | | H-JEPA | `hjepa_block()` | Multi-scale nested coalgebras | | Energy function | `EnergyFunction` | Compatibility measure in repr space | | Cost module | `CostModule` | IC + TC decomposition | | EMA update | `ema_update!()` | Collapse prevention | **The key insight**: JEPA’s prediction loss is an *obstruction to diagram commutativity*. Minimizing it drives the encoder toward being an exact coalgebra morphism — a structure-preserving map between the observation world model and the latent world model.