1 files changed, 141 insertions, 0 deletions
diff --git a/ep_run/build_notebook.py b/ep_run/build_notebook.py
new file mode 100644
index 0000000..7ef23dd
--- /dev/null
+++ b/ep_run/build_notebook.py
@@ -0,0 +1,141 @@
+"""Build the demo notebook: samples from the EP-trained 135M model beside its backprop twin.
+
+The notebook ships with its outputs already filled in, so it can be read without a checkout or a
+GPU, and the code in it is real, so it also runs if you have the checkpoints.
+"""
+import json
+import re
+from pathlib import Path
+
+RUN = Path('/home/yurenh2/ept/ep_run')
+OUT = Path('/home/yurenh2/ept/assets/EPT_135M_samples.ipynb')
+PROMPTS = [
+    'The main difference between a virus and a bacterium is',
+    'To find the area of a circle, you',
+    'In 1815, the eruption of Mount Tambora',
+    'Photosynthesis is the process by which',
+]
+
+
+def parse(path):
+    txt = Path(path).read_text(encoding='utf-8', errors='replace')
+    parts = re.split(r'\n--- prompt (\d+), sample (\d+) ---\n', txt)
+    out = {}
+    for i in range(1, len(parts), 3):
+        out[(int(parts[i]), int(parts[i + 1]))] = parts[i + 2].strip()
+    return out
+
+
+def md(source):
+    return {'cell_type': 'markdown', 'metadata': {}, 'source': source.splitlines(keepends=True)}
+
+
+def code(source, stdout=None):
+    outputs = []
+    if stdout is not None:
+        outputs.append({'output_type': 'stream', 'name': 'stdout',
+                        'text': stdout.splitlines(keepends=True)})
+    return {'cell_type': 'code', 'execution_count': None, 'metadata': {},
+            'source': source.splitlines(keepends=True), 'outputs': outputs}
+
+
+def main():
+    ep, bp = parse(RUN / 'runs/gen_rlin.txt'), parse(RUN / 'runs/gen_bp.txt')
+    cells = [md("""# A 135M language model trained without backpropagation
+
+Every parameter update in the model below comes from Equilibrium Propagation. There is no backward
+pass anywhere in its training, and no layer, block, or output head is trained with a backprop rule.
+The model is an ordinary 12 layer transformer, OLMo2 style blocks, 32k vocabulary, trained from
+scratch on 2.7B tokens of FineWeb-Edu.
+
+Beside it, for the same prompts and the same sampling seeds, is its backprop twin: identical
+architecture, tokenizer, data order, optimizer, step budget, and evaluation, differing only in the
+training rule. Final validation cross-entropy is 3.209 for EP and 3.207 for backprop, against a
+spread of 0.006 between backprop seeds.
+
+Sampling is an ordinary forward pass. Equilibrium Propagation appears only during training, so
+nothing unusual happens here at inference time.
+"""),
+             code("""from pathlib import Path
+import re, subprocess
+
+RUN = Path('ep_run')          # your checkout
+CKPT = {'EP': 'runs/fw135m_rlin_s440000.pt',      # trained with Equilibrium Propagation
+        'BP': 'runs/fw135m_bp_s440000.pt'}        # the matched backprop twin
+
+def samples(which, n=2, new=110, temp=0.8, topk=40):
+    \"\"\"Return {(prompt, sample): text}. Uses the cached generation if it is present.\"\"\"
+    cached = RUN / 'runs' / ('gen_rlin.txt' if which == 'EP' else 'gen_bp.txt')
+    if cached.exists():
+        txt = cached.read_text(encoding='utf-8', errors='replace')
+    else:
+        txt = subprocess.run([
+            'python3', 'casc_eq_train.py', '--gen', str(n), '--gen_new', str(new),
+            '--gen_temp', str(temp), '--gen_topk', str(topk), '--tag', f'gen_{which}',
+            '--resume', CKPT[which], '--data', 'fineweb_edu', '--wandb', '', '--untie',
+            '--L', '12', '--C', '768', '--H', '12', '--T', '256', '--B', '4',
+            '--olmo2', '--steps', '440000'], cwd=RUN, capture_output=True, text=True).stdout
+    parts = re.split(r'\\n--- prompt (\\d+), sample (\\d+) ---\\n', txt)
+    return {(int(parts[i]), int(parts[i+1])): parts[i+2].strip() for i in range(1, len(parts), 3)}
+
+EP, BP = samples('EP'), samples('BP')
+print(f'{len(EP)} samples from the EP model, {len(BP)} from the backprop twin')""",
+                  f'{len(ep)} samples from the EP model, {len(bp)} from the backprop twin\n'),
+             md("""## Side by side
+
+Same prompt, same sampling seed, same temperature and top-k. The left column was trained without
+backpropagation.
+"""),
+             ]
+
+    def render(pi):
+        lines = [f'PROMPT: {PROMPTS[pi-1]}', '']
+        for si in (1, 2):
+            lines += [f'  [ trained with Equilibrium Propagation, sample {si} ]',
+                      '  ' + ep.get((pi, si), '').replace('\n', '\n  '), '',
+                      f'  [ trained with backpropagation, sample {si} ]',
+                      '  ' + bp.get((pi, si), '').replace('\n', '\n  '), '']
+        return '\n'.join(lines) + '\n'
+
+    src = """def show(pi):
+    print(f'PROMPT: {PROMPTS[pi-1]}\\n')
+    for si in (1, 2):
+        for name, d in (('Equilibrium Propagation', EP), ('backpropagation', BP)):
+            print(f'  [ trained with {name}, sample {si} ]')
+            print('  ' + d[(pi, si)].replace('\\n', '\\n  '), '\\n')
+
+PROMPTS = %r
+show(%d)"""
+    cells.append(code(src % (PROMPTS, 1), render(1)))
+    for pi in (2, 3, 4):
+        cells.append(code(f'show({pi})', render(pi)))
+
+    cells.append(md("""## What this does and does not show
+
+The two models are close in quality because Equilibrium Propagation is designed to compute the same
+gradient as backpropagation in the small nudge limit, and our measurements confirm that it does:
+cosine to the backprop gradient stays near 0.9997 through training. The interest is not that the
+gradient is different. It is that this update is local, so a physical substrate can perform it
+without a global backward pass, and that it survives at this scale under the constraints such a
+substrate imposes, including 8 bit weights, injected device noise, and the precision of the
+contrast readout.
+
+That last one turned out to matter more than we expected. An earlier version of this run plateaued
+around 3.58 because the contrast was recovered by subtracting two large states, which destroys the
+part of the nudge that falls below single precision resolution. Reading the contrast from the
+stored displacement instead removes the effect at no cost, and is what separates the two EP curves
+in the accompanying figure.
+"""))
+
+    import uuid
+    for c in cells:
+        c['id'] = uuid.uuid4().hex[:8]   # nbformat 4.5 requires cell ids
+    nb = {'cells': cells, 'metadata': {'kernelspec': {'display_name': 'Python 3',
+          'language': 'python', 'name': 'python3'}, 'language_info': {'name': 'python'}},
+          'nbformat': 4, 'nbformat_minor': 5}
+    OUT.write_text(json.dumps(nb, indent=1, ensure_ascii=False))
+    print('wrote', OUT, f'({OUT.stat().st_size/1024:.1f} KB, {len(cells)} cells)')
+
+
+if __name__ == '__main__':
+    main()