|
| 1 | +// ============================================================================= |
| 2 | +// jepa — a JEPA-flavoured, self-supervised proxy: how much does the text tell |
| 3 | +// you about the FUTURE of the feed? |
| 4 | +// |
| 5 | +// For a given projector config, a linear (ridge) predictor is trained to map |
| 6 | +// features of the projection of frames 0..T-1 to a target computed from frame T, |
| 7 | +// which the projector never saw. No labels: the feed supervises itself. Held-out |
| 8 | +// R^2 is the score. |
| 9 | +// |
| 10 | +// features the text read back as ink, pooled to 16x8; plus four 8x4 channels |
| 11 | +// counting each arrow direction, one for '~' trails, and (for |
| 12 | +// multi-frame projections) one of |ink change| between the oldest |
| 13 | +// and newest block, so stacking frames gets credit too |
| 14 | +// target 'energy' (default): |next-frame luma - current luma| at 16x8 — |
| 15 | +// WHERE the next change will happen. Predicting it needs to know |
| 16 | +// where things are and which way they move, which is what |
| 17 | +// "understanding over time" should mean. Sign-free on purpose: a |
| 18 | +// dark car on grass and a bright ball on a floor move the same way |
| 19 | +// with opposite-signed deltas, which a linear probe cannot untangle. |
| 20 | +// 'delta': the signed change (measured: a linear probe gets R^2 <= 0.02) |
| 21 | +// 'next': the next frame's luma itself (mostly static; R^2 ~0.98 for all) |
| 22 | +// |
| 23 | +// Honest scope: this is a linear probe on a 16x8 pixel target, not a learned |
| 24 | +// joint embedding. The real JEPA (predict a frozen image encoder's embedding of |
| 25 | +// frame T, e.g. CLIP via DeepInfra) is the proposed next step; the harness is |
| 26 | +// the same with a different target function. |
| 27 | +// ============================================================================= |
| 28 | +import { lumaGrid } from '../core.mjs'; |
| 29 | +import { registerScorer, project } from '../registry.mjs'; |
| 30 | +import { renderSequence, randomSpec } from '../scenes.mjs'; |
| 31 | +import { textGrid } from './inverse.mjs'; |
| 32 | + |
| 33 | +const TW = 16, TH = 8, AW = 8, AH = 4; |
| 34 | +const ARROW = { '>': 0, '→': 0, '<': 1, '←': 1, '^': 2, '↑': 2, 'v': 3, '↓': 3, '~': 4 }; |
| 35 | + |
| 36 | +export function features(text) { |
| 37 | + const blocks = text.split('\n\n'), lines = blocks[blocks.length - 1].split('\n').filter((l) => !l.startsWith('legend:')); |
| 38 | + const t = textGrid(text), f = new Float64Array(TW * TH + 6 * AW * AH + 1), cnt = new Float64Array(TW * TH); |
| 39 | + // multi-frame projections: |ink(newest) - ink(oldest block)| pooled to 8x4 |
| 40 | + if (blocks.length > 1) { |
| 41 | + const o = textGrid(blocks[0]); |
| 42 | + for (let y = 0; y < Math.min(t.rows, o.rows); y++) for (let x = 0; x < Math.min(t.cols, o.cols); x++) { |
| 43 | + const ax = Math.min(AW - 1, Math.floor(x * AW / t.cols)), ay = Math.min(AH - 1, Math.floor(y * AH / t.rows)); |
| 44 | + f[TW * TH + 5 * AW * AH + ay * AW + ax] += Math.abs(t.grid[y * t.cols + x] - o.grid[y * o.cols + x]); |
| 45 | + } |
| 46 | + } |
| 47 | + for (let y = 0; y < t.rows; y++) for (let x = 0; x < t.cols; x++) { |
| 48 | + const X = Math.min(TW - 1, Math.floor(x * TW / t.cols)), Y = Math.min(TH - 1, Math.floor(y * TH / t.rows)); |
| 49 | + f[Y * TW + X] += t.grid[y * t.cols + x]; cnt[Y * TW + X]++; |
| 50 | + const ch = [...(lines[y] || '')][x], a = ARROW[ch]; |
| 51 | + if (a !== undefined) { |
| 52 | + const ax = Math.min(AW - 1, Math.floor(x * AW / t.cols)), ay = Math.min(AH - 1, Math.floor(y * AH / t.rows)); |
| 53 | + f[TW * TH + a * AW * AH + ay * AW + ax] += 1; |
| 54 | + } |
| 55 | + } |
| 56 | + for (let i = 0; i < TW * TH; i++) f[i] = cnt[i] ? f[i] / cnt[i] : 0; |
| 57 | + f[f.length - 1] = 1; |
| 58 | + return f; |
| 59 | +} |
| 60 | + |
| 61 | +export function target(seq, kind = 'delta') { |
| 62 | + const n = seq.frames.length, next = lumaGrid(seq.frames[n - 1], TW, TH); |
| 63 | + if (kind === 'next') return Float64Array.from(next, (v) => v / 255); |
| 64 | + const cur = lumaGrid(seq.frames[n - 2], TW, TH); |
| 65 | + if (kind === 'energy') return Float64Array.from(next, (v, i) => Math.abs(v - cur[i]) / 255); |
| 66 | + return Float64Array.from(next, (v, i) => (v - cur[i]) / 255); |
| 67 | +} |
| 68 | + |
| 69 | +// Solve (X^T X + lam I) W = X^T Y by Cholesky. X: n x d, Y: n x k. |
| 70 | +export function ridge(X, Y, lam = 1e-1) { |
| 71 | + const n = X.length, d = X[0].length, k = Y[0].length; |
| 72 | + const A = Array.from({ length: d }, () => new Float64Array(d)), B = Array.from({ length: d }, () => new Float64Array(k)); |
| 73 | + for (let r = 0; r < n; r++) { |
| 74 | + const x = X[r], y = Y[r]; |
| 75 | + for (let i = 0; i < d; i++) { if (!x[i]) continue; const xi = x[i]; |
| 76 | + for (let j = i; j < d; j++) A[i][j] += xi * x[j]; |
| 77 | + for (let j = 0; j < k; j++) B[i][j] += xi * y[j]; } |
| 78 | + } |
| 79 | + for (let i = 0; i < d; i++) { for (let j = 0; j < i; j++) A[i][j] = A[j][i]; A[i][i] += lam; } |
| 80 | + const L = Array.from({ length: d }, () => new Float64Array(d)); |
| 81 | + for (let i = 0; i < d; i++) for (let j = 0; j <= i; j++) { |
| 82 | + let s = A[i][j]; for (let q = 0; q < j; q++) s -= L[i][q] * L[j][q]; |
| 83 | + L[i][j] = i === j ? Math.sqrt(Math.max(s, 1e-12)) : s / L[j][j]; |
| 84 | + } |
| 85 | + const W = Array.from({ length: d }, () => new Float64Array(k)); |
| 86 | + for (let c = 0; c < k; c++) { |
| 87 | + const z = new Float64Array(d); |
| 88 | + for (let i = 0; i < d; i++) { let s = B[i][c]; for (let q = 0; q < i; q++) s -= L[i][q] * z[q]; z[i] = s / L[i][i]; } |
| 89 | + for (let i = d - 1; i >= 0; i--) { let s = z[i]; for (let q = i + 1; q < d; q++) s -= L[q][i] * W[q][c]; W[i][c] = s / L[i][i]; } |
| 90 | + } |
| 91 | + return W; |
| 92 | +} |
| 93 | + |
| 94 | +export function predict(W, x) { |
| 95 | + const k = W[0].length, out = new Float64Array(k); |
| 96 | + for (let i = 0; i < x.length; i++) if (x[i]) for (let c = 0; c < k; c++) out[c] += x[i] * W[i][c]; |
| 97 | + return out; |
| 98 | +} |
| 99 | + |
| 100 | +export function r2(P, Y) { |
| 101 | + const k = Y[0].length, mean = new Float64Array(k); |
| 102 | + for (const y of Y) for (let c = 0; c < k; c++) mean[c] += y[c] / Y.length; |
| 103 | + let sse = 0, sst = 0; |
| 104 | + for (let r = 0; r < Y.length; r++) for (let c = 0; c < k; c++) { sse += (P[r][c] - Y[r][c]) ** 2; sst += (Y[r][c] - mean[c]) ** 2; } |
| 105 | + return sst ? 1 - sse / sst : 0; |
| 106 | +} |
| 107 | + |
| 108 | +// Train + evaluate the probe for one projector config. Deterministic (seeded feeds). |
| 109 | +const memo = new Map(); |
| 110 | +export function probe(projector, params, { train = 600, test = 150, kind = 'energy', frames = 12, lam = 10 } = {}) { |
| 111 | + const key = JSON.stringify([projector, params, train, test, kind, frames, lam]); |
| 112 | + if (memo.has(key)) return memo.get(key); |
| 113 | + const sample = (seed) => { |
| 114 | + const seq = renderSequence(randomSpec(seed), frames + 1); |
| 115 | + const past = { frames: seq.frames.slice(0, frames), truth: seq.truth }; |
| 116 | + return { x: features(project(projector, past, params).text_projection), y: target(seq, kind) }; |
| 117 | + }; |
| 118 | + const tr = Array.from({ length: train }, (_, i) => sample(1000 + i)); |
| 119 | + const te = Array.from({ length: test }, (_, i) => sample(900000 + i)); |
| 120 | + const W = ridge(tr.map((s) => s.x), tr.map((s) => s.y), lam); |
| 121 | + const res = { r2_test: r2(te.map((s) => predict(W, s.x)), te.map((s) => s.y)), |
| 122 | + r2_train: r2(tr.map((s) => predict(W, s.x)), tr.map((s) => s.y)), train, test, kind }; |
| 123 | + memo.set(key, res); |
| 124 | + return res; |
| 125 | +} |
| 126 | + |
| 127 | +registerScorer({ |
| 128 | + name: 'jepa', |
| 129 | + cost: 'cheap', |
| 130 | + about: 'self-supervised: held-out R^2 of a ridge probe predicting the next frame\'s change from the text', |
| 131 | + async score({ projection, ctx = {} }) { |
| 132 | + const r = probe(projection.projector, projection.params, ctx.jepa || {}); |
| 133 | + return { fidelity: Math.max(0, Math.min(1, r.r2_test)), detail: r }; |
| 134 | + }, |
| 135 | +}); |
0 commit comments