diff --git a/README.md b/README.md index 9182100..7c12f6a 100644 --- a/README.md +++ b/README.md @@ -1,242 +1,187 @@

-

⚡ EB-JEPA

+

⚙️ JEPA-ASML

-

Energy-Based Joint-Embedding Predictive Architectures

+

JEPA pour la superoptimisation d'assembleur

-
- Github - ArXiv -
- -
-

- Meta AI Research, FAIR + Apprendre une représentation du sens d'un programme assembleur — invariante + à la forme produite par le compilateur — et un prédicteur world-model qui + transforme la représentation d'un code lent (O0) en celle de son équivalent + optimisé (O3).

-

- Basile Terver, - Randall Balestriero, - Megi Dervishi, - David Fan, - Quentin Garrido, - Tushar Nagarajan, -
- Koustuv Sinha, - Wancong Zhang, - Mike Rabbat, - Yann LeCun, - Amir Bar -

- -

- An open source library and tutorial for learning representations for
- prediction and planning using joint embedding predictive architectures. -

- -

- EB-JEPA Architecture -

- -> Each example is (almost) self-contained and training takes up to a few hours on a single GPU card. - --- -## 📚 Examples +## 🎯 Le projet en une page -### [Image JEPA](examples/image_jepa/README.md) +Un compilateur peut produire, à partir d'une même fonction, des assembleurs très +différents selon le niveau d'optimisation (`-O0` … `-O3`). Ces variantes ont le +**même sens** (mêmes entrées → mêmes sorties) mais des **formes** et des **coûts** +(vitesse) très différents. -Self-supervised representations from unlabeled images on CIFAR-10, evaluated on classification. +L'idée de ce projet est d'appliquer les **Joint-Embedding Predictive +Architectures** (JEPA) à ce problème, avec une représentation *factorisée* : -![Image JEPA Architecture](examples/image_jepa/assets/arch_figure.png) +| Facteur | Ce qu'il capture | Doit être… | +|---|---|---| +| **meaning** (sens) | la sémantique du bloc | invariant à l'optimisation (O0 ≈ O3) | +| **temporality** (coût) | la vitesse / le coût d'exécution | sensible à l'optimisation | +| **architecture** | la cible matérielle | conditionnement (travaux futurs) | -### [Video JEPA](examples/video_jepa/README.md) +Une fois cet espace appris, la **superoptimisation** devient un problème de +*planning* dans l'espace latent : partir de `repr(O0)`, appliquer une **action** +(« optimise vers O3 ») et atteindre `repr(O3)`, puis décoder vers de l'assembleur. -Predict next image representation in a sequence. +> Ce dépôt est un **fork** de la librairie [EB-JEPA](#-crédits--upstream) de Meta +> AI (FAIR). Tout le code propre au projet vit dans `eb_jepa/asm/` et +> `examples/asm_superopt/` ; le reste est la librairie JEPA d'origine, réutilisée +> telle quelle. -![Moving MNIST](examples/video_jepa/assets/viz.png) +--- -### [AC Video JEPA](examples/ac_video_jepa/README.md) +## 🧩 Le pipeline -JEPA for world modeling + planning in Two Rooms environment. +``` + fichier .c + │ clang -S -O{0,1,2,3} -masm=intel + ▼ + assembleur x86-64 (Intel) eb_jepa/asm/corpus.py + │ découpage en blocs de base (straight-line, sans branchements) + ▼ + graphe de flot de données (DataFlowGraph) eb_jepa/asm/dataflow.py + │ nœuds = instructions ; arêtes = dépendances def→use (REG / MEM / FLAGS) + ▼ + encodeur GNN (message passing typé + mean-pool) eb_jepa/asm/encoder.py + │ têtes factorisées → vecteur « meaning » + sortie « temporality » + ▼ + prédicteur world-model : g(O0) + action(→O3) ≈ g(O3) +``` -| Planning Episode | Task Definition | -|------------------|-----------------| -| Successful planning episode | Episode task definition | -| *Successful planning episode* | *From init to goal state* | +Le **graphe de flot de données** expose volontairement l'*indépendance* entre +instructions (deux instructions sans chemin def→use n'ont pas d'arête) — c'est +exactement ce qui rend détectable le parallélisme de bloc (packing SLP / AVX) en +aval. Aucun code n'est exécuté ni assemblé : tout part du texte assembleur. --- -## 🚀 Installation - -### HTW cluster — quick start (hackathon only) +## 🏋️ Entraînement en deux phases -> Skip this section unless you are on the HTW hackathon cluster — the generic install below is all you need locally. +Le point clé du design : l'encodeur n'utilise **jamais** les labels de niveau +d'optimisation. Le couplage O0↔O3 est réservé au prédicteur. -Please follow the [setup instructions](setup.md) before starting the project. +**Phase 1 — Encodeur (auto-supervisé, sans labels d'opt)** — `train_jepa.py` ---- +JEPA latent classique transposé sur le graphe : on masque des nœuds et on prédit +leur **représentation** (pas leurs mnémoniques). Encodeur en ligne `θ` + encodeur +cible `ξ = EMA(θ)`, stop-gradient, et garde-fou anti-collapse VICReg +(variance + covariance). Tous les blocs (O0…O3) sont dans le même tas, vus comme +des programmes bruts. -### Local / generic (start here) +**Phase 2 — Prédicteur (sur l'encodeur gelé)** — `train_predictor.py` -We use [uv](https://docs.astral.sh/uv/guides/projects/) for package management. +C'est ici, et seulement ici, qu'on utilise les paires O0↔O3 de la même fonction : -```bash -# Install dependencies -uv sync -# Option 1: Activate virtual environment -source .venv/bin/activate -python -m examples.image_jepa.main -# Option 2: Run directly with uv -uv run python -m examples.image_jepa.main ``` -If you need conda-specific packages, you can use **Conda + uv** - -```bash -# Create conda environment with Python 3.12 -conda create -n eb_jepa python=3.12 -y -conda activate eb_jepa -# Install package in editable mode with dev dependencies (pytest, black, isort, autoflake) -uv pip install -e . --group dev +état = un programme → g = encodeur_gelé(programme) +action = un niveau d'opt cible → embedding d'action appris +modèle : g(forme_i) + action(niveau_j) → prédiction de g(forme_j) (résidu = l'« édit ») ``` -Add these to your `~/.bashrc` for persistent configuration. - -```bash -# Where datasets are stored / looked up -export EBJEPA_DSETS=/path/to/eb_jepa/datasets -# Optional: Directory for checkpoints and logs -export EBJEPA_CKPTS=/path/to/checkpoints -``` +**Métrique honnête : battre l'identité.** Comme le sens est ~invariant, `g(O0)` +est déjà proche de `g(O3)` ; le prédicteur n'a de valeur que s'il retrouve le vrai +`g(O3)` *mieux* que le no-op `g(O0)`. On rapporte la retrieval@1 des deux + la +similarité cosinus. -Verify the install with `uv run pytest tests/`. +--- -## 🏋️ Training +## 🚀 Démarrage rapide -### Quick Start +On utilise [uv](https://docs.astral.sh/uv/) pour la gestion des dépendances. +Le pipeline asm a besoin de **clang** sur le `PATH`. ```bash -# Local training -python -m examples.{image_jepa,video_jepa,ac_video_jepa}.main -``` -> Our default configs are tuned for H100 GPUs. With older GPUs (e.g., A100, V100), you may need to reduce batch size to fit in memory. - -### 📂 Folder Structure +# 1. Dépendances +uv sync +source .venv/bin/activate -All experiments use a unified folder structure: +# 2. Générer un corpus C, le compiler et construire la base de graphes +python -m examples.asm_superopt.gen_corpus --n 4000 # → fichiers .c +python -m eb_jepa.asm.corpus # → data/corpus.jsonl -``` -checkpoints/ -└── {example_name}/ - ├── dev_2026-01-16_00-10/ # Single/local runs (dev_ prefix) - │ └── {exp_name}_seed1/ - │ - ├── sweep_2026-01-16_00-10/ # Auto-named 3-seed sweep - │ ├── {exp_name}_seed1/ - │ ├── {exp_name}_seed1000/ - │ └── {exp_name}_seed10000/ - │ - └── sweep_my_experiment/ # Custom-named sweep - └── ... +# 3. Pipeline complet (encodeur → prédicteur → figures de diagnostic) +bash examples/asm_superopt/run_all.sh ``` -`{exp_name}` encodes key hyperparameters to avoid folder collisions, e.g.: -- **image_jepa**: `resnet_vicreg_proj_bs256_ep300_ph2048_po2048_std1.0_cov80.0` -- **video_jepa**: `resnet_bs64_lr0.001_std10.0_cov100.0` -- **ac_video_jepa**: `impala_cov8_std16_simt12_idm1` +> `run_all.sh` tourne sur **CPU** par défaut : pour ce petit GNN, le fallback CPU +> des ops de scatter sur MPS rend le GPU local ~250× plus lent. Le vrai chemin +> CUDA est le cluster SLURM (Vivatech / HTW). -
-🖥️ SLURM Launcher (optional) +### Scripts du dossier `examples/asm_superopt/` -| Command | Description | -|---------|-------------| -| `--example {name}` | Choose: `image_jepa`, `video_jepa`, `ac_video_jepa`, `maze`, `fintime`, `ltsf`, `eeg`, `audio`, `pointcloud`, `gray_scott`, `intuitive_physics`, `factors_of_variation` | -| `--fname {path}` | Run the sweep specified in the config at `{path}` | -| `--single` | Launch single job (dev mode) | -| `--sweep {name}` | Custom sweep name | -| `--array-parallelism {N}` | Limits the maximum number of concurrent jobs to `N` | -| `--full-sweep` | Full hyperparameter sweep from config | -| `--use-wandb-sweep` | Enable wandb sweep UI | +| Script | Rôle | +|---|---| +| `gen_corpus.py` | génère un corpus C de fonctions feuilles straight-line | +| `train_jepa.py` | Phase 1 — encodeur JEPA latent (masking, sans labels d'opt) | +| `train_mask.py` | baseline générative (prédit les mnémoniques masqués) | +| `train_predictor.py` | Phase 2 — prédicteur world-model sur encodeur gelé | +| `analyze_opt_proximity.py` | diagnostic : O0/O1/O2/O3 d'une même fonction sont-ils proches ? | +| `compare_idempotence.py` | le prédicteur est-il idempotent (O3 → O3 fixe) ? | +| `study_o3_input.py` | comportement du prédicteur selon le niveau d'entrée | +| `bridge_specificity.py` | le « pont » O0→O3 est-il spécifique à la bonne fonction ? | -```bash -# 3 seeds with wandb averaging (recommended) -python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml - -# Custom sweep name -python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --sweep my_experiment - -# Single job -python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --single +--- -# Full hyperparameter sweep -python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --full-sweep +## 📁 Structure -# With wandb sweep UI for hyperparameter analysis -python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --use-wandb-sweep ``` - -Replace `image_jepa` with `ac_video_jepa`, `video_jepa`, or `maze` for other examples. - -**Full Sweep Configuration:** The `--full-sweep` flag reads the `sweep.param_grid` section from the example's YAML config file (e.g., `examples/image_jepa/cfgs/default.yaml`). Without this flag, only a 3-seed sweep is launched. To customize sweep parameters, edit the `sweep` section in the config: - -```yaml -# Example: examples/image_jepa/cfgs/default.yaml -sweep: - param_grid: - loss.cov_coeff: [0.1, 1.0, 10.0, 100.0] - loss.std_coeff: [1.0, 10.0] - meta.seed: [1, 1000, 10000] +eb_jepa/asm/ cœur du projet (lib) +├── dataflow.py assembleur x86-64 → graphe de flot de données +├── encoder.py encodeur GNN factorisé (meaning / temporality) +├── corpus.py .c → clang → blocs → base de graphes JSONL +└── synth.py blocs synthétiques + réécritures préservant le sens + +examples/asm_superopt/ expériences, entraînement, diagnostics +├── seed_corpus/ test_corpus/ sources C +├── data/ corpus JSONL, checkpoints (.pt), figures +└── run_all.sh orchestrateur retrain + ré-évaluation + +eb_jepa/ librairie JEPA upstream (jepa, losses, planning, …) +tests/ test_asm_dataflow.py · test_asm_encoder.py · test_asm_corpus.py +docs/ slides, schémas, barème ``` -### Wandb Seed Averaging - -Runs with the same hyperparameters but different seeds share the same wandb run name, enabling automatic averaging: - -1. Go to wandb web UI → Runs table -2. Click **"Group by"** → select **"Name"** - → Groups runs with identical hyperparameters (different seeds) together - -To filter runs from a specific sweep: -3. Click **"Filter"** → **"Group"** → select your sweep name - -For detailed wandb sweep analysis (parallel coordinates, hyperparameter importance): -1. Use `--use-wandb-sweep` flag when launching -2. Go to wandb web UI → left pane → **"Sweeps"** → click your sweep name - -**SLURM Configuration:** SLURM parameters default to the HTW cluster and are read from `EBJEPA_SLURM_*` env vars (set by `env.sh`, which also auto-detects your account/QOS per user). Override per launch with the CLI flags `--partition`/`--account`/`--cpus-per-task`/`--time-min`/`--gpus-per-node`, or export the matching `EBJEPA_SLURM_*` var. The `SLURM_DEFAULTS` dictionary at the top of `examples/launch_sbatch.py` holds the fallbacks. - -
- -## 🧪 Running test cases +--- -Libraries added to eb_jepa [must have their own test cases](/tests/). To run the tests: +## 🧪 Tests ```bash -# With uv sync installation uv run pytest tests/ -# With conda + uv installation (no .venv created) -pytest tests/ ``` -## 👩‍💻 Development +Les tests propres au projet : `test_asm_dataflow.py`, `test_asm_encoder.py`, +`test_asm_corpus.py`. -Before contributing, please format your code with the following tools: +## 👩‍💻 Développement + +Avant de contribuer, formatez le code : ```bash -# Remove unused imports autoflake --remove-all-unused-imports -r --in-place . -# Sort imports python -m isort eb_jepa examples tests -# Format code python -m black eb_jepa examples tests ``` -## 📚 Citing EB-JEPA +--- + +## 📚 Crédits & upstream -If you find this repository useful, please consider giving a ⭐ and citing: +Ce dépôt est construit sur **EB-JEPA** (*Energy-Based Joint-Embedding Predictive +Architectures*), la librairie open-source de **Meta AI Research (FAIR)** : +Basile Terver, Randall Balestriero, Megi Dervishi, David Fan, Quentin Garrido, +Tushar Nagarajan, Koustuv Sinha, Wancong Zhang, Mike Rabbat, Yann LeCun, Amir Bar. ```bibtex @misc{terver2026lightweightlibraryenergybasedjointembedding, @@ -250,6 +195,6 @@ If you find this repository useful, please consider giving a ⭐ and citing: } ``` -## 📄 License +## 📄 Licence -EB-JEPA is Apache licensed. See [LICENSE](LICENSE.md). +Sous licence Apache 2.0 (héritée d'EB-JEPA). Voir [LICENSE](LICENSE.md). diff --git a/docs/slides_structure.pdf b/docs/slides_structure.pdf index 9109fdd..3798b05 100644 Binary files a/docs/slides_structure.pdf and b/docs/slides_structure.pdf differ diff --git a/docs/slides_structure.tex b/docs/slides_structure.tex index 6d561bd..a7635ca 100644 --- a/docs/slides_structure.tex +++ b/docs/slides_structure.tex @@ -1,17 +1,31 @@ \documentclass[aspectratio=169]{beamer} \usetheme{metropolis} - -\usepackage[utf8]{inputenc} -\usepackage[T1]{fontenc} +\setbeamerfont{title}{size=\large,series=\bfseries} +\setbeamerfont{subtitle}{size=\normalsize} + +\usepackage{iftex} +\ifPDFTeX + \usepackage[utf8]{inputenc} + \usepackage[T1]{fontenc} +\fi \usepackage[french]{babel} \usepackage{tikz} \usepackage{listings} -\usetikzlibrary{arrows.meta, positioning, fit, backgrounds, calc} +\usepackage{array} +\usepackage{booktabs} +\usetikzlibrary{arrows.meta, positioning, fit, backgrounds, calc, decorations.pathreplacing, trees, shapes.geometric} \definecolor{cReg}{HTML}{2E86C1} \definecolor{cMem}{HTML}{B7950B} \definecolor{cFlag}{HTML}{922B21} \definecolor{cBox}{HTML}{F2F3F4} +\definecolor{camblue}{HTML}{1F4E79} +\definecolor{camsky}{HTML}{2E86C1} +\definecolor{camgrey}{HTML}{566573} + +% badge coloré : \chip{couleur}{texte} +\newcommand{\chip}[2]{\tikz[baseline]\node[fill=#1,text=white,rounded corners=2pt,% + inner xsep=5pt,inner ysep=2pt,font=\scriptsize\bfseries]{#2};} \lstset{ basicstyle=\ttfamily\small, @@ -22,12 +36,87 @@ \title{EB-JEPA pour l'assembleur} \subtitle{Structure du système} -\author{Armand Coiffe, Sami Rabinovitch, Ulysse Camille, Grégoire Rouvière} +\author{A. Coiffe, S. Rabinovitch, U. Camille, G. Rouvière} \date{\today} \begin{document} -\maketitle +% --- titre : frame manuelle (évite l'overfull du \maketitle metropolis) --- +\begin{frame}[plain,noframenumbering] + \centering + \vfill + {\usebeamerfont{title}\usebeamercolor[fg]{title}\inserttitle\par} + \vspace{2mm} + {\usebeamerfont{subtitle}\usebeamercolor[fg]{subtitle}\insertsubtitle\par} + \vspace{6mm} + {\usebeamerfont{author}\insertauthor\par} + \vspace{3mm} + {\usebeamerfont{date}\insertdate\par} + \vfill +\end{frame} + +% --------------------------------------------------------------- +\begin{frame}{L'intuition : monotâche (humain) \emph{vs} omniscient (world-model)} +\begin{columns}[T] +% ====== GAUCHE : structure du code, faisceau de focus humain (monotâche) ====== +\begin{column}{0.54\textwidth} +\centering\chip{camsky}{\;L'humain \textperiodcentered\ monot\^ache\;}\par\vspace{4pt} +\begin{tikzpicture}[font=\tiny,>=stealth, + every node/.style={draw,rounded corners=1pt,inner sep=2pt,minimum height=0.34cm,minimum width=1.25cm}, + faded/.style={draw=camgrey!40,fill=camgrey!6,text=camgrey!55}, + hot/.style={draw=camsky,fill=camsky!20,text=camblue,thick}, + hotedge/.style={draw=camsky,thick}, + fadededge/.style={draw=camgrey!40}] +% --- colonne surlignée : la branche suivie (drill-down vertical) --- +\node[hot] (root) at (0,0) {repo}; +\node[hot] (ma) at (0,-1.05) {module A}; +\node[hot] (ff) at (0,-2.1) {fonction f()}; +\node[hot] (leaf) at (0,-3.15) {bloc / var}; +% --- frères estompés, décalés à droite (ignorés par l'humain) --- +\node[faded] (mb) at (2.05,-1.05) {module B}; +\node[faded] (mc) at (3.75,-1.05) {module C}; +\node[faded] (fg) at (2.05,-2.1) {fichier g}; +% --- arêtes --- +\draw[hotedge] (root) -- (ma); +\draw[hotedge] (ma) -- (ff); +\draw[hotedge] (ff) -- (leaf); +\draw[fadededge] (root) -- (mb); +\draw[fadededge] (root) -- (mc); +\draw[fadededge] (ma) -- (fg); +% --- œil + regard qui descend la branche (monotâche) --- +\node[draw=camblue,ellipse,minimum width=0.5cm,minimum height=0.26cm,inner sep=0pt,line width=0.5pt,fill=white] (eye) at (-1.5,0.3) {}; +\fill[camblue] (eye.center) circle (0.045); +\draw[camsky,thick,dashed,->] (eye.south) to[out=-90,in=160] (leaf.west); +\end{tikzpicture} +\par\vspace{1pt}\scriptsize\textcolor{camgrey}{une branche, un niveau \`a la fois — le reste est ignoré} +\end{column} +% ====== DROITE : échelle d'abstraction, le world-model voit TOUT ====== +\begin{column}{0.46\textwidth} +\centering\chip{camblue}{\;Le world-model \textperiodcentered\ omniscient\;}\par\vspace{4pt} +\begin{tikzpicture}[font=\scriptsize,>=stealth, + box/.style={draw=camblue!55,rounded corners=1pt,minimum width=2.4cm,minimum height=0.42cm,inner sep=2pt}] +\node[box,fill=camblue!8] (a) {source haut niveau}; +\node[box,fill=camblue!8,below=3pt of a] (b) {AST / SSA}; +\node[box,fill=camsky!32,below=3pt of b] (c) {LLVM-IR \textperiodcentered\ CFG}; +\node[box,fill=camsky!32,below=3pt of c] (d) {assembleur}; +\node[box,fill=camblue!8,below=3pt of d] (e) {binaire}; +% flèche « compilation » à gauche, bien dégagée des boîtes +\draw[->,camblue,thick] ([xshift=-5mm]a.north west) -- ([xshift=-5mm]e.south west); +\node[font=\tiny,text=camblue,rotate=90] at ([xshift=-9mm]$(a.west)!0.5!(e.west)$) {compilation}; +% accolade « voit TOUT » à droite +\draw[decorate,decoration={brace,amplitude=4pt},camblue,thick] + (a.north east) -- (e.south east) + node[midway,right=4pt,align=left,font=\tiny,text=camblue] {le mod\`ele\\ voit \textbf{TOUT}}; +\end{tikzpicture} +\par\vspace{1pt}\scriptsize\textcolor{camsky}{$\blacksquare$}\,\textcolor{camgrey}{IR/asm = on apprend ici} +\end{column} +\end{columns} +\vspace{2pt} +\centering\footnotesize Là où l'humain ne voit qu'\textbf{une} branche et \textbf{un} +niveau à la fois, le world-model a une vue \emph{omnisciente} — toute la structure +\emph{et} toutes les échelles. +\par\textcolor{camblue}{$\Rightarrow$ on apprend près du \textbf{matériel} (IR/asm) $\to$ un \textbf{optimiseur de code}.} +\end{frame} % --------------------------------------------------------------- \begin{frame}{Vue d'ensemble du pipeline} @@ -70,7 +159,7 @@ % Préambule requis : \usepackage{array}, \usepackage{booktabs} \begin{frame}{État de l'art : positionnement} \scriptsize - \renewcommand{\arraystretch}{1.3} + \renewcommand{\arraystretch}{1.05} \setlength{\tabcolsep}{4pt} \begin{center} \begin{tabular}{@{}>{\raggedright\arraybackslash}p{0.22\linewidth} @@ -93,7 +182,7 @@ \end{tabular} \end{center} - \vfill + \vspace{1mm} \begin{center} \footnotesize\textbf{Constat :} tout l'état de l'art \emph{exécute} ou \emph{génère} ; aucune approche n'apprend une représentation latente du @@ -119,14 +208,12 @@ \texttt{+ - * \& | \textasciicircum{} << >> \textasciitilde} \item \textbf{Pas} de \texttt{if} / boucle / appel / division $\Rightarrow$ un \textbf{seul bloc de base} de O0 à O3. - \item Filtre \emph{non-trivial} : $\geq\!1$ param \textbf{et} $\geq\!1$ - opérateur (sinon repli en constante : rien à apprendre). \end{itemize} \end{column} \begin{column}{0.46\textwidth} \textbf{Rôle \& limite} \begin{itemize} - \item Volume et diversité contrôlés, sans headers ni dépendances. + \item Volume et diversité contrôlés. \item Chaque fonction $\rightarrow$ 4 formes asm (O0…O3) garanties single-BB. \item \emph{Peu représentatif du vrai code} $\rightarrow$ piste : corpus réel (AnghaBench). @@ -136,7 +223,7 @@ \end{frame} % --------------------------------------------------------------- -\begin{frame}[fragile]{Étape 2 — Représentation : C $\rightarrow$ asm $\rightarrow$ graphe} +\begin{frame}[fragile]{Étape 2 — Représentation : C $\rightarrow$ IR $\rightarrow$ graphe} \footnotesize \begin{columns}[T] \begin{column}{0.54\textwidth} @@ -145,7 +232,7 @@ int poly(int x){ return 3*x*x+5*x+7; } \end{lstlisting} {\scriptsize\textbf{asm \texttt{-O0}} — naïf (spills pile)} -\begin{lstlisting}[language={[x86masm]Assembler}, basicstyle=\ttfamily\scriptsize, aboveskip=1pt, belowskip=1pt] +\begin{lstlisting}[language={[x86masm]Assembler}, basicstyle=\ttfamily\scriptsize, aboveskip=0pt, belowskip=0pt] mov dword ptr [rsp-4], edi imul eax, dword ptr [rsp-4], 3 imul eax, dword ptr [rsp-4] @@ -154,7 +241,7 @@ add eax, 7 \end{lstlisting} {\scriptsize\textbf{asm \texttt{-O3}} — factorisé (Horner)} -\begin{lstlisting}[language={[x86masm]Assembler}, basicstyle=\ttfamily\scriptsize, aboveskip=1pt, belowskip=1pt] +\begin{lstlisting}[language={[x86masm]Assembler}, basicstyle=\ttfamily\scriptsize, aboveskip=0pt, belowskip=0pt] lea eax, [rdi + 2*rdi] add eax, 5 imul eax, edi @@ -163,37 +250,91 @@ \end{column} \begin{column}{0.42\textwidth} \centering -\textbf{Graphe de flot de données}\\[2mm] +\textbf{Graphe de flot de données (\texttt{-O0})}\\[2mm] \begin{tikzpicture}[ n/.style={draw, rounded corners, fill=cBox, font=\ttfamily\scriptsize, - minimum width=20mm, minimum height=5.5mm}, - e/.style={-{Stealth[length=2mm]}, thick, cReg}, - xe/.style={-{Stealth[length=2mm]}, thick, cMem, densely dashed}, - inp/.style={draw, circle, fill=cMem!15, draw=cMem, font=\ttfamily\scriptsize, - minimum size=6.5mm}] - \node[n] (i0) {lea eax}; - \node[n, below=7.5mm of i0] (i1) {add eax,5}; - \node[n, below=7.5mm of i1] (i2) {imul eax}; - \node[n, below=7.5mm of i2] (i3) {add eax,7}; - \draw[e] (i0) -- node[right, font=\tiny]{REG} (i1); - \draw[e] (i1) -- node[right, font=\tiny]{REG} (i2); - \draw[e] (i2) -- node[right, font=\tiny]{REG} (i3); - \node[inp, left=12mm of i1] (x) {x}; - \draw[xe] (x) to[bend left=15] (i0.west); - \draw[xe] (x) to[bend right=15] (i2.west); -\end{tikzpicture}\\[1mm] -\scriptsize Nœuds = instructions ($+\,$7 traits struct.) ; arêtes typées + minimum width=21mm, minimum height=5mm}, + re/.style={-{Stealth[length=2mm]}, thick, cReg}, + me/.style={-{Stealth[length=2mm]}, thick, cMem, densely dashed}] + \node[n] (n0) {n0 mov [rsp-4]}; + \node[n, below=5mm of n0] (n1) {n1 imul eax,3}; + \node[n, below=5mm of n1] (n2) {n2 imul eax}; + \node[n, right=12mm of n1] (n3) {n3 imul ecx,5}; + \node[n, below=5mm of n2] (n4) {n4 add eax,ecx}; + \node[n, below=5mm of n4] (n5) {n5 add eax,7}; + % MEM : le spill [rsp-4] (écrit par n0) relu par n1, n2, n3 + \draw[me] (n0) -- node[left, font=\tiny]{MEM} (n1); + \draw[me] (n0.south) to[bend left=12] (n2.north east); + \draw[me] (n0.east) to[bend left=20] node[above, font=\tiny]{MEM} (n3.north); + % REG : chaîne sur eax / ecx + \draw[re] (n1) -- node[right, font=\tiny]{REG} (n2); + \draw[re] (n2) -- node[left, font=\tiny]{REG} (n4); + \draw[re] (n3.south) to[bend left=15] node[right, font=\tiny]{REG} (n4.east); + \draw[re] (n4) -- node[right, font=\tiny]{REG} (n5); +\end{tikzpicture} +\end{column} +\end{columns} +{\scriptsize +Nœuds $=$ instructions ($+\,$7 traits struct.) ; arêtes typées \textcolor{cReg}{REG}\,/\,\textcolor{cMem}{MEM}\,/\,\textcolor{cFlag}{FLAG} (def$\rightarrow$use). -\textcolor{cMem}{\texttt{x}} (\texttt{rdi}) lue \textbf{2 fois} ; absence d'arête $=$ indépendance -$\Rightarrow$ parallélisme (SLP/AVX). +Le spill \texttt{[rsp-4]} (\textcolor{cMem}{MEM}, écrit par n0) est \textbf{relu 3 fois} ; +n3 sans arête \textcolor{cReg}{REG} entrante $=$ indépendant.} +\end{frame} + +% --------------------------------------------------------------- +\begin{frame}{Étape 3 — Le masquage du graphe} +\footnotesize +\begin{columns}[T] +\begin{column}{0.50\textwidth} +\textbf{Principe} +\begin{itemize}\setlength{\itemsep}{0pt} + \item On cache une \textbf{partie du graphe} : on choisit + \textbf{$\sim$15\,\% des nœuds} (choix \emph{empirique}). + \item Pour chaque nœud masqué, son contenu est + remplacé par un \textbf{jeton de masque appris}. + \item \textbf{Les arêtes ne sont pas touchées} : la structure + \textcolor{cReg}{REG}\,/\,\textcolor{cMem}{MEM}\,/\,\textcolor{cFlag}{FLAG} + reste visible $\Rightarrow$ le nœud « sait » où est le trou. +\end{itemize} +\textbf{Pourquoi des nœuds (et pas des arêtes)} +\begin{itemize}\setlength{\itemsep}{1pt} + \item Les arêtes se \emph{redéduisent} mécaniquement des nœuds. + \item Forcé de retrouver un nœud depuis ses seules dépendances, + l'encodeur doit \textbf{apprendre la structure} du programme. +\end{itemize} +\end{column} +\begin{column}{0.46\textwidth} +\centering +\textbf{$\sim$15\,\% des nœuds $\rightarrow$ jeton}\\[1mm] +\resizebox{0.92\linewidth}{!}{% +\begin{tikzpicture}[ + n/.style={draw, rounded corners, fill=cBox, font=\ttfamily\scriptsize, + minimum width=20mm, minimum height=5mm}, + mk/.style={draw, rounded corners, fill=cFlag!12, draw=cFlag, dashed, + font=\ttfamily\scriptsize, minimum width=20mm, minimum height=5mm}, + re/.style={-{Stealth[length=2mm]}, thick, cReg}, + me/.style={-{Stealth[length=2mm]}, thick, cMem, densely dashed}] + \node[n] (n0) {n0 mov [rsp-4]}; + \node[n, below=5mm of n0] (n1) {n1 imul eax,3}; + \node[mk, below=5mm of n1] (n2) {\textbf{?? (masqué)}}; + \node[n, right=12mm of n1] (n3) {n3 imul ecx,5}; + \node[n, below=5mm of n2] (n4) {n4 add eax,ecx}; + \node[n, below=5mm of n4] (n5) {n5 add eax,7}; + \draw[me] (n0) -- node[left, font=\tiny]{MEM} (n1); + \draw[me] (n0.west) to[bend right=55] node[left, font=\tiny]{MEM} (n2.west); + \draw[me] (n0.east) to[bend left=20] node[above, font=\tiny]{MEM} (n3.north); + \draw[re] (n1) -- node[right, font=\tiny]{REG} (n2); + \draw[re] (n2) -- node[left, font=\tiny]{REG} (n4); + \draw[re] (n3.south) to[bend left=15] (n4.east); + \draw[re] (n4) -- node[right, font=\tiny]{REG} (n5); +\end{tikzpicture}}\\[2mm] +{\scriptsize Le contenu de n2 disparaît, \textbf{ses arêtes restent}.} \end{column} \end{columns} -\vspace{0.5mm} -{\scriptsize\emph{Single-BB} = bloc en ligne droite (aucun saut) $\Rightarrow$ le flot de données suffit — scope v1.} \end{frame} % --------------------------------------------------------------- -\begin{frame}{Étape 3 — Encodeur GNN factorisé} +\begin{frame}{Étape 4 — Encodeur GNN factorisé} \centering \begin{tikzpicture}[ box/.style={draw, rounded corners, fill=cBox, align=center, @@ -222,9 +363,9 @@ \end{frame} % --------------------------------------------------------------- -\begin{frame}{Étape 4 — Phase 1 : EB-JEPA par masquage (encodeur)} +\begin{frame}{Étape 5 — Phase 1 : EB-JEPA par masquage (encodeur)} \centering -\resizebox{0.98\textwidth}{!}{% +\resizebox{0.56\textwidth}{!}{% \begin{tikzpicture}[ box/.style={draw, rounded corners, fill=cBox, align=center, font=\scriptsize, minimum height=8mm, minimum width=18mm}, @@ -248,30 +389,30 @@ \node[font=\tiny, below=0.5mm of D, align=center] {$D=1-\cos$\\(+ VICReg)}; \end{tikzpicture}} -\vspace{1mm} -{\footnotesize +{\scriptsize \[ \mathcal{L}_\theta= \underbrace{\;2-2\cos\!\big(\hat s,\;\mathrm{sg}\,s\big)\;}_{\text{régression latente }=\,\lVert\hat s-\mathrm{sg}\,s\rVert^2} \;+\;\lambda\, - \underbrace{\tfrac1d\textstyle\sum_k \max\!\big(0,\,1-\sigma_k(\hat s)\big)}_{\text{VICReg (anti-collapse)}} + \underbrace{\tfrac1d\textstyle\sum_k \max\!\big(0,\,1-\sigma_k(\hat s)\big)}_{\text{variance (anti-collapse)}} + \;+\;\nu\, + \underbrace{\tfrac1d\textstyle\sum_{i\neq j} \big[C(\hat s)\big]_{ij}^{2}}_{\text{covariance (décorrélation)}} \]} -\vspace{0.5mm} -\footnotesize +\scriptsize \textbf{Le rapport avec EB-JEPA} -\begin{itemize} +\begin{itemize}\setlength{\itemsep}{1pt} \item $x$ = contexte (programme \textbf{partiellement masqué}), $y$ = le \emph{même} programme complet. \item Prédiction dans l'\textbf{espace latent} (pas de reconstruction d'octets/tokens) : le cœur de JEPA. - \item Anti-collapse : cible \textbf{EMA + stop-grad} ($\xi\!\leftarrow\!\tau\xi+(1-\tau)\theta$) + variance VICReg. + \item Anti-collapse : cible \textbf{EMA + stop-grad} ($\xi\!\leftarrow\!\tau\xi+(1-\tau)\theta$) + VICReg (\emph{variance} + \emph{covariance}). \item \textbf{Aucune étiquette}, aucun appariement O0$\leftrightarrow$O3 — O0…O3 ne sont \emph{que} de l'asm. \end{itemize} \end{frame} % --------------------------------------------------------------- -\begin{frame}{Étape 5 — Phase 2 : prédicteur (modèle du monde)} +\begin{frame}{Étape 6 — Phase 2 : prédicteur (modèle du monde)} \centering -\resizebox{0.98\textwidth}{!}{% +\resizebox{0.66\textwidth}{!}{% \begin{tikzpicture}[ box/.style={draw, rounded corners, fill=cBox, align=center, font=\scriptsize, minimum height=8mm, minimum width=16mm}, @@ -296,8 +437,7 @@ \node[font=\tiny, below=0.5mm of D, align=center] {$D=1-\cos(\hat g, g_{\mathrm{O3}})$}; \end{tikzpicture}} -\vspace{1mm} -{\footnotesize +{\scriptsize \[ \mathcal{L}_\phi= 2-2\cos\!\big(\hat g,\;g_{\mathrm{O3}}\big) @@ -305,10 +445,9 @@ \qquad(\text{cible \textbf{gelée}}\;\Rightarrow\;\text{ni EMA ni VICReg}) \]} -\vspace{0.5mm} -\footnotesize +\scriptsize \textbf{Le rapport avec EB-JEPA} -\begin{itemize} +\begin{itemize}\setlength{\itemsep}{1pt} \item Même squelette JEPA, mais le prédicteur est \textbf{conditionné par une action} $a$ (variable latente $z$) $\Rightarrow$ \textbf{modèle du monde}. \item L'encodeur cible EMA devient l'encodeur \textbf{gelé} : cibles fixes $\Rightarrow$ pas d'anti-collapse. @@ -350,12 +489,12 @@ \begin{columns}[T] \begin{column}{0.5\textwidth} \centering -\includegraphics[height=0.5\textheight]{slides_structure_figs/opt_proximity_hist_bare.png}\\ +\includegraphics[height=0.44\textheight]{slides_structure_figs/opt_proximity_hist_bare.png}\\ \tiny Même code (O0 vs optimisé) \emph{vs} code sans rapport. \end{column} \begin{column}{0.5\textwidth} \centering -\includegraphics[height=0.5\textheight]{slides_structure_figs/opt_proximity_bridge_bare.png}\\ +\includegraphics[height=0.44\textheight]{slides_structure_figs/opt_proximity_bridge_bare.png}\\ \tiny cos(O0, O3) avant / après le prédicteur (fonctions \emph{jamais vues}). \end{column} \end{columns} @@ -399,38 +538,6 @@ \end{columns} \end{frame} -% --------------------------------------------------------------- -\begin{frame}{Cohérence : « si c'est déjà optimal, ne touche à rien »} -\footnotesize -\begin{columns}[T] -\begin{column}{0.46\textwidth} -\centering -\includegraphics[height=0.52\textheight]{slides_structure_figs/o3_idem_compare_bare.png}\\ -\tiny Fonctions \emph{jamais vues}. Point fixe O3$\to$O3 \emph{vs} optimisation O0$\to$O3. -\end{column} -\begin{column}{0.54\textwidth} -\textbf{Le problème.} Le prédicteur n'a vu que des entrées \emph{non -optimisées}. Donné un O3 déjà optimal (action~$\rightarrow$O3), il le -\textbf{déplace fortement} ($\cos(\text{sortie},\text{O3})=0{,}43$ au lieu de -$1$) : \textbf{pas de point fixe} $\Rightarrow$ on ne peut pas \emph{itérer} le -prédicteur (Phase 3 diverge). -\vspace{1.5mm} - -\textbf{Le correctif : paires identité.} On ajoute des exemples -$x \rightarrow x$ (édition nulle) : « déjà au niveau cible $\Rightarrow$ ne rien -faire ». Le modèle le déduit de la \emph{représentation} (pas d'étiquette de -niveau) --- ce qu'exige la recherche en Phase 3. -\vspace{1.5mm} - -\textbf{Résultat (held-out).} -\begin{itemize} - \item Point fixe O3$\to$O3 : $0{,}43 \rightarrow \mathbf{1{,}00}$. - \item Optimisation O0$\to$O3 \textbf{intacte} : $0{,}75 \rightarrow 0{,}75$. -\end{itemize} -\end{column} -\end{columns} -\end{frame} - % --------------------------------------------------------------- \begin{frame}{Suite — boucle de superoptimisation} \begin{itemize} @@ -440,4 +547,38 @@ \end{itemize} \end{frame} +% --------------------------------------------------------------- +\begin{frame}{Prochaines étapes} +\begin{itemize}\setlength{\itemsep}{6pt} + \item \textbf{Hiérarchisation} + \item \textbf{Module acteur} + \item \textbf{Hyperparameter tuning} +\end{itemize} +\vspace{4mm} +\begin{center} + \large\textcolor{camblue}{Si l'on parvient à mettre le JEPA à l'échelle $\Rightarrow$ \textbf{compilateur universel}.} +\end{center} +\end{frame} + +% --------------------------------------------------------------- +\begin{frame}{Représentation : graphe contrôle / données / mémoire} +\centering +\includegraphics[height=0.74\textheight]{slides_structure_figs/programgraph_sum_array.jpeg}\\[1mm] +{\scriptsize Exemple réel \texttt{sum\_array} : arêtes typées \textcolor{cFlag}{contrôle} / \textcolor{cReg}{données} / \textcolor{cMem!80!black}{mémoire}.} +\end{frame} + +% --------------------------------------------------------------- +\begin{frame}{Sans VICReg, l'espace latent s'effondre} +\centering +\includegraphics[height=0.76\textheight]{slides_structure_figs/vicreg_anticollapse.jpeg}\\[1mm] +{\scriptsize Ablation : écart-type des embeddings $1{,}07$ (sain) \emph{vs} $0{,}0003$ (effondré) ; le nuage PCA se réduit à un point.} +\end{frame} + +% --------------------------------------------------------------- +\begin{frame}{Le gain se concentre sur le seul vrai saut : O0\,$\rightarrow$\,O1} +\centering +\includegraphics[height=0.74\textheight]{slides_structure_figs/predictor_par_transition.jpeg}\\[1mm] +{\scriptsize Le predictor brille sur O0\,$\rightarrow$\,O1 ; O1\,$\rightarrow$\,O2\,$\rightarrow$\,O3 quasi triviaux (clang sature dès \texttt{-O1}).} +\end{frame} + \end{document} diff --git a/docs/slides_structure_figs/predictor_par_transition.jpeg b/docs/slides_structure_figs/predictor_par_transition.jpeg new file mode 100644 index 0000000..85cd533 Binary files /dev/null and b/docs/slides_structure_figs/predictor_par_transition.jpeg differ diff --git a/docs/slides_structure_figs/programgraph_sum_array.jpeg b/docs/slides_structure_figs/programgraph_sum_array.jpeg new file mode 100644 index 0000000..f5f36f7 Binary files /dev/null and b/docs/slides_structure_figs/programgraph_sum_array.jpeg differ diff --git a/docs/slides_structure_figs/vicreg_anticollapse.jpeg b/docs/slides_structure_figs/vicreg_anticollapse.jpeg new file mode 100644 index 0000000..d2fa2fa Binary files /dev/null and b/docs/slides_structure_figs/vicreg_anticollapse.jpeg differ diff --git a/examples/asm_superopt/train_jepa.py b/examples/asm_superopt/train_jepa.py index 2d2c3b7..aa2aa1d 100644 --- a/examples/asm_superopt/train_jepa.py +++ b/examples/asm_superopt/train_jepa.py @@ -13,7 +13,8 @@ Both the encoder and the (throwaway) predictor are trained jointly, end-to-end, by the SAME loss. The target encoder gets no gradient; it follows the online encoder by EMA. Collapse (encoder -> constant) is prevented by EMA + stop-grad -(the core, BYOL-style) plus a VICReg variance hinge (the safety belt). +(the core, BYOL-style) plus the VICReg anti-collapse pair (variance hinge + +covariance decorrelation) as the safety belt. NO opt-level labels are used: every block is just a raw program (all of O0..O3 go in the same pile). The O0<->O3 pairing is reserved for the *predictor* @@ -86,16 +87,27 @@ def ema_update(self, tau: float): bt.copy_(bo) -def jepa_loss(pred, tgt, var_coeff: float): - """Latent regression (cosine) + variance hinge. Returns (loss, inv, tgt_std).""" +def jepa_loss(pred, tgt, var_coeff: float, cov_coeff: float): + """Latent regression (cosine) + full VICReg anti-collapse (variance + covariance). + + Returns (loss, inv, cov, tgt_std). + """ p = F.normalize(pred, dim=-1) t = F.normalize(tgt.detach(), dim=-1) inv = (2.0 - 2.0 * (p * t).sum(-1)).mean() # == ||p - t||^2 on unit sphere - # VICReg-style hinge: keep per-dim std of predictions >= 1 (deter collapse). + # VICReg variance hinge: keep per-dim std of predictions >= 1 (deter collapse). std = pred.std(dim=0) var = torch.relu(1.0 - std).mean() - tgt_std = tgt.std(dim=0).mean().item() # collapse monitor (raw target reps) - return inv + var_coeff * var, inv.item(), tgt_std + # VICReg covariance: push off-diagonal covariances to 0 (decorrelate dims so + # they don't carry redundant info -> guards against informational collapse). + d = pred.shape[-1] + z = pred - pred.mean(dim=0) + cov = (z.T @ z) / (pred.shape[0] - 1) # (d, d) covariance matrix + off = cov - torch.diag(torch.diagonal(cov)) # zero the diagonal + cov_loss = off.pow(2).sum() / d # mean off-diagonal energy + tgt_std = tgt.std(dim=0).mean().item() # collapse monitor (raw target reps) + loss = inv + var_coeff * var + cov_coeff * cov_loss + return loss, inv.item(), cov_loss.item(), tgt_std def featurize_all(encoder, corpus, min_nodes=3): @@ -182,7 +194,9 @@ def main(): ap.add_argument("--wd", type=float, default=1e-4) ap.add_argument("--mask-frac", type=float, default=0.15) ap.add_argument("--ema", type=float, default=0.996, help="target EMA momentum tau") - ap.add_argument("--var-coeff", type=float, default=1.0, help="anti-collapse variance weight") + ap.add_argument("--var-coeff", type=float, default=1.0, help="VICReg variance (anti-collapse) weight") + ap.add_argument("--cov-coeff", type=float, default=0.04, + help="VICReg covariance (decorrelation) weight; VICReg uses cov:var ~ 1:25") ap.add_argument("--hidden", type=int, default=128) ap.add_argument("--meaning-dim", type=int, default=128) ap.add_argument("--layers", type=int, default=3) @@ -224,21 +238,21 @@ def main(): for ep in range(1, args.epochs + 1): model.train() rng.shuffle(train) - inv_sum = std_sum = nb = 0 + inv_sum = cov_sum = std_sum = nb = 0 for items in _chunks(train, args.batch): ids, feats, ei, et, midx = collate(items, gen, args.mask_frac, device) pred, tgt = model(ids, feats, ei, et, midx) - loss, inv, tgt_std = jepa_loss(pred, tgt, args.var_coeff) + loss, inv, cov, tgt_std = jepa_loss(pred, tgt, args.var_coeff, args.cov_coeff) opt.zero_grad() loss.backward() opt.step() model.ema_update(args.ema) - inv_sum += inv; std_sum += tgt_std; nb += 1 + inv_sum += inv; cov_sum += cov; std_sum += tgt_std; nb += 1 if ep % 2 == 0 or ep == 1: r1, cos = evaluate(model, val, eval_gen(), args.mask_frac, device) - print(f"epoch {ep:3d} inv {inv_sum/nb:.4f} tgt_std {std_sum/nb:.3f} " - f"val_r@1 {r1:.3f} val_cos {cos:.3f}") + print(f"epoch {ep:3d} inv {inv_sum/nb:.4f} cov {cov_sum/nb:.4f} " + f"tgt_std {std_sum/nb:.3f} val_r@1 {r1:.3f} val_cos {cos:.3f}") if r1 > best: best = r1 _save(args.out, encoder, vocab, args)