File size: 2,830 Bytes
dbb04e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | # Omega-JEPA (Dream Stream) - Technical Specification
## 1. Overview
This document outlines the architecture for the Omega-JEPA Predictor Network, a clean-room implementation inspired by Joint Embedding Predictive Architectures (JEPA) but specialized for the "Dream Stream" environment. The core innovation is the integration of Protocol Omega metrics (Time-Reversal Asymmetry and Z-Scores) directly into the validation loop to ensure causal fidelity in state predictions.
## 2. Theoretical Foundation
### 2.1. The JEPA Paradigm
Unlike generative models that predict pixels (x), JEPA predicts representations (y) in an abstract space.
- **Context ($S_x$):** The current state representation.
- **Action ($a$):** The control or transition vector.
- **Latent ($z$):** A stochastic variable capturing uncertainty in the transition.
- **Prediction ($S_y$):** The predicted future state representation.
Equation: $S_y = Predictor(S_x, a, z)$
### 2.2. Protocol Omega Integration
To prevent the model from learning "easy" but non-causal shortcuts (representation collapse), we enforce Time-Reversal Asymmetry (TRA).
- **TRA Hypothesis:** A valid causal transition $A \to B$ should be distinguishable from $B \to A$ in the latent energy landscape.
- **Omega Score:** A composite metric combining prediction error (L2) with TRA violation penalties.
## 3. Architecture Design
### 3.1. `OmegaJEPA_Predictor` (The Brain)
- **Type:** PyTorch `nn.Module`.
- **Structure:** Multi-Layer Perceptron (MLP) with residual connections.
- **Inputs:**
- `context_embedding`: Tensor [Batch, Dim]
- `action_vector`: Tensor [Batch, ActionDim]
- `latent_z`: Tensor [Batch, LatentDim] (Optional/sampled)
- **Mechanism:**
1. Concatenate $S_x$, $a$, and $z$.
2. Pass through a high-dimensional projection layer.
3. Apply LayerNorm and GeLU activations.
4. Output projected state $\hat{S}_y$.
### 3.2. `OmegaMetrics` (The Auditor)
- **Purpose:** Stateless validator class to compute physics-inspired metrics.
- **Key Metrics:**
- `compute_tra(state_t, state_t1)`: Measures asymmetry magnitude.
- `compute_z_score(residuals)`: Standard deviation based anomaly detection.
- `energy_function(state)`: Helper for TRA computation (e.g., magnitude or learned energy).
## 4. Implementation Constraints
- **Framework:** PyTorch (v2.x).
- **Style:** Strict typing, modular design.
- **Clean Room:** No usage of Meta/Facebook source code; strictly first-principles implementation based on LeCun's 2022 paper and Protocol Omega specs.
## 5. Usage
```python
model = OmegaJEPA_Predictor(dim=256, action_dim=64)
metrics = OmegaMetrics()
# Forward pass
s_next_pred = model(s_curr, action, z)
# Validation
tra_score = metrics.compute_tra(s_curr, s_next_pred)
```
|