diff options
| author | Yuren Hao <yurenh2@illinois.edu> | 2026-08-04 05:46:57 -0500 |
|---|---|---|
| committer | Yuren Hao <yurenh2@illinois.edu> | 2026-08-04 05:46:57 -0500 |
| commit | f3ade7674328c55a98d32fd332a3a1cf79fd40a6 (patch) | |
| tree | 789515f125d31b894f147f61f1b1bab287a6dcc3 | |
| parent | 3962588a667d0d8945044b105dedda75e3fa3a41 (diff) | |
RESULT 89终局 + 交付物: 135M完赛 EP 3.2087 vs BP 3.2074(+0.1% ppl), 旧读出+45.2%
窗口敏感性诚实记录: 97%窗给-0.0030(seed间距0.0059), 完赛大窗给+0.0013(间距0.0005, cosine
衰减后BP收拢); 以大窗为准且不得称"统计不可分"(2.6倍间距), 正确表述=差0.1%且处分辨极限,
seed数(BP n=2/EP n=1)不足以给区间。
交付: --gen模式(复用模型定义, EP/BP同代码路径)、fig_135m_curves(四曲线/横轴token/BP细线在上)、
EPT_135M_samples.ipynb(同prompt同种子并排, 输出已烤入, 代码真可跑)。
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
| -rw-r--r-- | assets/EPT_135M_samples.ipynb | 294 | ||||
| -rw-r--r-- | assets/figs/fig_135m_curves.png | bin | 0 -> 335410 bytes | |||
| -rw-r--r-- | docs/campaign/CASCADE_ABLATION_PLAN.md | 18 | ||||
| -rw-r--r-- | ep_run/build_notebook.py | 141 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 43 | ||||
| -rw-r--r-- | ep_run/fig_curves.py | 94 |
6 files changed, 590 insertions, 0 deletions
diff --git a/assets/EPT_135M_samples.ipynb b/assets/EPT_135M_samples.ipynb new file mode 100644 index 0000000..3672149 --- /dev/null +++ b/assets/EPT_135M_samples.ipynb @@ -0,0 +1,294 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# A 135M language model trained without backpropagation\n", + "\n", + "Every parameter update in the model below comes from Equilibrium Propagation. There is no backward\n", + "pass anywhere in its training, and no layer, block, or output head is trained with a backprop rule.\n", + "The model is an ordinary 12 layer transformer, OLMo2 style blocks, 32k vocabulary, trained from\n", + "scratch on 2.7B tokens of FineWeb-Edu.\n", + "\n", + "Beside it, for the same prompts and the same sampling seeds, is its backprop twin: identical\n", + "architecture, tokenizer, data order, optimizer, step budget, and evaluation, differing only in the\n", + "training rule. Final validation cross-entropy is 3.209 for EP and 3.207 for backprop, against a\n", + "spread of 0.006 between backprop seeds.\n", + "\n", + "Sampling is an ordinary forward pass. Equilibrium Propagation appears only during training, so\n", + "nothing unusual happens here at inference time.\n" + ], + "id": "4296b70f" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "source": [ + "from pathlib import Path\n", + "import re, subprocess\n", + "\n", + "RUN = Path('ep_run') # your checkout\n", + "CKPT = {'EP': 'runs/fw135m_rlin_s440000.pt', # trained with Equilibrium Propagation\n", + " 'BP': 'runs/fw135m_bp_s440000.pt'} # the matched backprop twin\n", + "\n", + "def samples(which, n=2, new=110, temp=0.8, topk=40):\n", + " \"\"\"Return {(prompt, sample): text}. Uses the cached generation if it is present.\"\"\"\n", + " cached = RUN / 'runs' / ('gen_rlin.txt' if which == 'EP' else 'gen_bp.txt')\n", + " if cached.exists():\n", + " txt = cached.read_text(encoding='utf-8', errors='replace')\n", + " else:\n", + " txt = subprocess.run([\n", + " 'python3', 'casc_eq_train.py', '--gen', str(n), '--gen_new', str(new),\n", + " '--gen_temp', str(temp), '--gen_topk', str(topk), '--tag', f'gen_{which}',\n", + " '--resume', CKPT[which], '--data', 'fineweb_edu', '--wandb', '', '--untie',\n", + " '--L', '12', '--C', '768', '--H', '12', '--T', '256', '--B', '4',\n", + " '--olmo2', '--steps', '440000'], cwd=RUN, capture_output=True, text=True).stdout\n", + " parts = re.split(r'\\n--- prompt (\\d+), sample (\\d+) ---\\n', txt)\n", + " return {(int(parts[i]), int(parts[i+1])): parts[i+2].strip() for i in range(1, len(parts), 3)}\n", + "\n", + "EP, BP = samples('EP'), samples('BP')\n", + "print(f'{len(EP)} samples from the EP model, {len(BP)} from the backprop twin')" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "8 samples from the EP model, 8 from the backprop twin\n" + ] + } + ], + "id": "f28f2772" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Side by side\n", + "\n", + "Same prompt, same sampling seed, same temperature and top-k. The left column was trained without\n", + "backpropagation.\n" + ], + "id": "923bcfcd" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "source": [ + "def show(pi):\n", + " print(f'PROMPT: {PROMPTS[pi-1]}\\n')\n", + " for si in (1, 2):\n", + " for name, d in (('Equilibrium Propagation', EP), ('backpropagation', BP)):\n", + " print(f' [ trained with {name}, sample {si} ]')\n", + " print(' ' + d[(pi, si)].replace('\\n', '\\n '), '\\n')\n", + "\n", + "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']\n", + "show(1)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PROMPT: The main difference between a virus and a bacterium is\n", + "\n", + " [ trained with Equilibrium Propagation, sample 1 ]\n", + " The main difference between a virus and a bacterium is how it is carried around the body. While a virus cannot multiply in the body, it can cause harm in the body, which can be a result of a wide variety of factors. When the virus is present in the body, it is more likely to be killed by the immune system.\n", + " How does a bacterium survive a virus?\n", + " The most common germs that a bacterium can survive in your body are bacteria. Bacteria can live for several days without food, but when the temperature of the liquid reaches below freezing point, the bacteria will start to multiply\n", + "\n", + " [ trained with backpropagation, sample 1 ]\n", + " The main difference between a virus and a bacterium is how they spread. Bacteria don’t carry germs and they cannot infect. When they enter the body through a wound, they are not passed to you or someone else, but they can. Microbes of certain types of bacteria are so dangerous that they rarely kill.\n", + " Microbes in the bloodstream can cause an infection even if they do not cause pain. Microbes of a certain type of bacteria, called Staphylococcus, can cause infection even if they do not cause pain. Staphylococcus bacteria are usually found in the throat, but may also be found\n", + "\n", + " [ trained with Equilibrium Propagation, sample 2 ]\n", + " The main difference between a virus and a bacterium is that the bacterium multiplies by multiplying by itself, whereas a virus may spread by multiplying by itself.\n", + " The main difference between a virus and a bacterium is that the bacterial cell forms a new organism by using its genetic material (DNA) which is the same as the viral cell.\n", + " In short, an organism is a living cell.\n", + " In summary, bacteria have been shown to be capable of producing, multiplying, multiplying and surviving viruses. In this sense, bacteria are a living organism.\n", + " The most common bacteria are the most numerous, the most\n", + "\n", + " [ trained with backpropagation, sample 2 ]\n", + " The main difference between a virus and a bacterium is that the virus can only spread from person to person by contact and may spread from person to person through bodily fluids such as water, sweat, feces, urine, and saliva. The virus is not capable of replicating in humans.\n", + " A bacterium is a simple organism that can only infect its host. Bacteria are generally considered to be one of the most difficult organisms to grow in the laboratory and, as such, are difficult to grow in real life conditions. In addition, the virus can't survive in the environment with its host, such as the human\n", + "\n" + ] + } + ], + "id": "3b00b85a" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "source": [ + "show(2)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PROMPT: To find the area of a circle, you\n", + "\n", + " [ trained with Equilibrium Propagation, sample 1 ]\n", + " To find the area of a circle, you need to find the volume of the circle. Find the area of the circle by multiplying both sides by 2.\n", + " 2. In this case, we would just multiply both sides by 2 and then multiply both sides by 2 and we get the area of the circle.\n", + " The value of the area of the circle is 2.\n", + " To find the area of the circle, simply divide both sides by 2.\n", + " 3. Since the area of the circle is 2, the area of the circle is 1.\n", + " If the area of the circle is 2\n", + "\n", + " [ trained with backpropagation, sample 1 ]\n", + " To find the area of a circle, you need to find its volume, which is something you know about the surface of the circle.\n", + " I'm sure you'll want to find the area of a circle that matches the surface of the circle. If you're not sure, try to find the area by finding the perimeter of the circle and you will find that the area of the circle will be the same.\n", + " As I stated before, finding the area of a circle is just the same as finding the area of a circle.\n", + " The solution to the problem in the second example\n", + " Can you\n", + "\n", + " [ trained with Equilibrium Propagation, sample 2 ]\n", + " To find the area of a circle, you will need to know how much area is in the circle.\n", + " - How many times can I draw an outline for this area?\n", + " - How many times can I draw this area?\n", + " - How much area is there in the circle?\n", + " - How many times can I draw a circle?\n", + " - What part of the circumference is it?\n", + " - If I draw the circumference on a circle then it is\n", + " - How is the area of the circumference?\n", + " - Can a circle be curved?\n", + " - Can a line be drawn around the figure\n", + "\n", + " [ trained with backpropagation, sample 2 ]\n", + " To find the area of a circle, you will need to know its radius.\n", + " The radius of the circle is the length of its radius.\n", + " The radius of a circle is a constant value.\n", + " Since each circle has exactly the same size, the diameter is the same.\n", + " For the sake of simplicity, we will use the following formula.\n", + " So, the radius of the circle is .\n", + " The formula will take you to the nearest circle.\n", + " The radius of the circle is , which will give you the diameter of the circle.\n", + " You will be able to find the number of\n", + "\n" + ] + } + ], + "id": "2ff01588" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "source": [ + "show(3)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PROMPT: In 1815, the eruption of Mount Tambora\n", + "\n", + " [ trained with Equilibrium Propagation, sample 1 ]\n", + " In 1815, the eruption of Mount Tambora in Indonesia was one of the most significant events in history, causing massive damage to the world. It left 2.5 million homes and affected 3.3 million people.\n", + " The eruption caused the death of an estimated 1.9 million people. The eruption affected Indonesia, causing 7 million hectares of land to be lost. The eruption destroyed an estimated 1.9 million houses and caused mass destruction of buildings.\n", + " The United States is the only country to suffer such a disaster. This is a case of the US not having the right to rule the world\n", + "\n", + " [ trained with backpropagation, sample 1 ]\n", + " In 1815, the eruption of Mount Tambora, the strongest volcano in the world at that time, prompted the government to declare an end to the construction of the new building.\n", + " In 1830, the last eruption of Mount St. Helens was recorded. The next eruption of St. Helens occurred in 1829, causing significant damage to the city of Washington, the capital of Pennsylvania.\n", + " Since then, the city has been expanding, but there is still no permanent population, so the construction of the new city is still ongoing.The U.S. Department of Education released a report\n", + "\n", + " [ trained with Equilibrium Propagation, sample 2 ]\n", + " In 1815, the eruption of Mount Tambora in Indonesia caused the death of more than 4,000 people. In the 19th century the eruption also produced a wave of earthquakes that caused more than 3,500 people to be killed.\n", + " In the 1960s, a similar eruption began in Indonesia, which may have caused a tsunami that claimed 30,000 lives in the region. The largest tsunami in the world to hit Indonesia since 1998 occurred in Indonesia in 2004. The main source of the tsunami was an underwater volcano that rose from the seafloor. While the tsunami has occurred in Indonesia for about\n", + "\n", + " [ trained with backpropagation, sample 2 ]\n", + " In 1815, the eruption of Mount Tambora had produced earthquakes along the shore of the island. While the eruption was in its early stages, one of the most significant events in history occurred in 1815. Ten people were killed, three were wounded, and one million dollars (about $1,00,000 dollars) was spent on the eruption.\n", + " Ashfall of 1817\n", + " The eruption of Prince William Sound triggered further events in 1819. It caused the evacuation of approximately 70,000 people from the city of Quebec. In 1812, the volcano was again active and erupted in 1812. The\n", + "\n" + ] + } + ], + "id": "9f017ccf" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "source": [ + "show(4)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PROMPT: Photosynthesis is the process by which\n", + "\n", + " [ trained with Equilibrium Propagation, sample 1 ]\n", + " Photosynthesis is the process by which plants convert carbon dioxide into energy and water. Photosynthesis is the process by which plants convert sunlight into energy.\n", + " This is a great site for students to learn about the process of photosynthesis.\n", + " This site is great for students to explore photosynthesis. It is a very cool site and makes a great addition to any science classroom.\n", + " This site will be great for students to learn about plants and photosynthesis.\n", + " This site will be great for students to learn about plants and the energy they use from photosynthesis.|Name: _________________________||Period:\n", + "\n", + " [ trained with backpropagation, sample 1 ]\n", + " Photosynthesis is the process by which plants convert carbon dioxide into carbohydrates and use it to produce ethanol, a byproduct of their metabolism.\n", + " What is the purpose of photosynthesis?\n", + " The primary purpose of photosynthesis is to convert light energy into chemical energy for plants, and the second is to provide energy for the plant kingdom.Nestled at the border of the ancient cities of Troy, the city of Troy is a truly ancient and fascinating place. From its first appearance to its second, it has been inhabited since at least the first century BC. It was a city for many centuries\n", + "\n", + " [ trained with Equilibrium Propagation, sample 2 ]\n", + " Photosynthesis is the process by which plants extract energy from sunlight, carbon dioxide, nitrogen and water. The most abundant form of photosynthesis is the light-independent pathway of photosynthesis, which consists of two phases: photosynthetic and stationary phase.\n", + " Photosynthesis and the process are similar, but in the energy-independent pathway of photosynthesis, the light energy is converted into a molecule called ATP. This is the chemical process used for generating energy from light.\n", + " Photosynthesis is a biochemical process that converts light energy into chemical energy. In other words, photosynthesis is the process by which plants use the\n", + "\n", + " [ trained with backpropagation, sample 2 ]\n", + " Photosynthesis is the process by which plants extract energy from sunlight and use that energy to produce food and water for the rest of the plant.\n", + " “We’ve always thought of the food web as a ‘homeostasis network’ of bacteria, fungi and arthropods,” said study co-author Dr Paul C. Wilcox, a lecturer in microbiology at the University of Exeter and one of the authors of the paper. “But it’s much more than that.”\n", + " The study, published online in Science Advances, shows that the nitrogen-fixing bacteria found\n", + "\n" + ] + } + ], + "id": "fdb18168" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## What this does and does not show\n", + "\n", + "The two models are close in quality because Equilibrium Propagation is designed to compute the same\n", + "gradient as backpropagation in the small nudge limit, and our measurements confirm that it does:\n", + "cosine to the backprop gradient stays near 0.9997 through training. The interest is not that the\n", + "gradient is different. It is that this update is local, so a physical substrate can perform it\n", + "without a global backward pass, and that it survives at this scale under the constraints such a\n", + "substrate imposes, including 8 bit weights, injected device noise, and the precision of the\n", + "contrast readout.\n", + "\n", + "That last one turned out to matter more than we expected. An earlier version of this run plateaued\n", + "around 3.58 because the contrast was recovered by subtracting two large states, which destroys the\n", + "part of the nudge that falls below single precision resolution. Reading the contrast from the\n", + "stored displacement instead removes the effect at no cost, and is what separates the two EP curves\n", + "in the accompanying figure.\n" + ], + "id": "2009e2a3" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}
\ No newline at end of file diff --git a/assets/figs/fig_135m_curves.png b/assets/figs/fig_135m_curves.png Binary files differnew file mode 100644 index 0000000..674cd8a --- /dev/null +++ b/assets/figs/fig_135m_curves.png diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md index 3a78932..2298ab4 100644 --- a/docs/campaign/CASCADE_ABLATION_PLAN.md +++ b/docs/campaign/CASCADE_ABLATION_PLAN.md @@ -2159,3 +2159,21 @@ BP seed2 在跑(定散布尺度)。 ⟹ **11M 到 135M 全部落在种子噪声内**; "gap 随尺寸增长"作为 EP 属性正式作废(RESULT 80 已判为 伪影签名, 本条为 135M 端点确证)。 - EP rlin 剩 11.4k 步(~2h)出终局尾窗; 届时更新为完赛口径。 + +### RESULT 89 终局 (完赛口径, 尾窗 396-440k, n=441/臂) +| 臂 | 尾窗均值 | vs BP 均值 | ppl | +|---|---|---|---| +| EP read_lin | 3.2087 | **+0.0013** | **+0.1%** | +| BP twin s1 | 3.2071 | | | +| BP twin s2 | 3.2076 | | | +| BP 两 seed 均值 | 3.2074 | (间距 **0.0005**) | | +| EP 旧读出 | 3.5804 | +0.3730 | **+45.2%** | +- ⚠️ **窗口敏感性, 须诚实报**: 97% 处窗口(422-428.6k, n=67)给 EP **−0.0030**、seed 间距 0.0059; + 完赛大窗(n=441)给 EP **+0.0013**、seed 间距 0.0005(cosine 衰减到底后两 BP 收拢)。 + **以完赛大窗为准**: EP 比 BP 高 0.0013 = 0.1% ppl。 +- **不得声称"统计不可分"**: 0.0013 是 seed 间距的 2.6 倍。但 BP 仅 n=2(一个自由度)、EP n=1, + 该设计根本无力分辨 0.001 量级 ⟹ 正确表述 = "**差 0.1% 困惑度, 处在本设计的分辨极限附近**", + 并明说 seed 数不足以给出区间。若要主张不可分, 需 EP n≥2 与 BP n≥3(约 3 天/臂)。 +- 旧读出终值 +45.2% ppl = 精度伪影在 135M 的完整体量。 +- 交付物已出终版: assets/figs/fig_135m_curves.{png,pdf}(四曲线, 横轴 token) + + assets/EPT_135M_samples.ipynb(EP/BP 同 prompt 同种子并排, 输出已烤入, 代码可跑)。 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() diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index ed164ce..05c1e4e 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -53,6 +53,11 @@ ap.add_argument('--dgain_geo', type=float, default=0.0) # >0: per-layer geomet # gain = geo^l (layer 0 = x1), optionally capped ap.add_argument('--dgain_geo_cap', type=float, default=0.0) # >0: cap for the geometric profile ap.add_argument('--dgain_rand', type=float, default=0.0) +ap.add_argument('--gen', type=int, default=0) # >0: sample text from the resumed checkpoint and exit +ap.add_argument('--gen_prompts', default='') # '|'-separated prompts; empty = built-in set +ap.add_argument('--gen_temp', type=float, default=0.8) +ap.add_argument('--gen_topk', type=int, default=40) +ap.add_argument('--gen_new', type=int, default=120) ap.add_argument('--probe_dgspec', type=int, default=0) # >0: M1 spectroscopy, value = n batches; exits before training ap.add_argument('--probe_gains', default='1,2,4,8,16,32,64,128,256') ap.add_argument('--probe_f64', action='store_true') # fp64 states+model in the probe: the fp-floor decisive arm # >1: per-STEP log-uniform dgain_top in @@ -859,6 +864,44 @@ if args.ddp_grad_test: import sys sys.exit(0) +if args.gen > 0: + # Sampling from a resumed checkpoint. Inference here is an ordinary forward pass through the + # blocks, which is the point: EP appears only in training. The same code path loads an EP or a + # backprop checkpoint, since both scripts save the same keys, so the two can be sampled side by + # side under identical settings. + import sys + from tokenizers import Tokenizer as _Tok + _tk = _Tok.from_file(str(DD / 'tokenizer.json')) + prompts = [p for p in args.gen_prompts.split('|') if p] or [ + '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', + ] + @torch.no_grad() + def _sample(prompt, n_new, temp, topk, seed): + torch.manual_seed(seed) + ids = _tk.encode(prompt).ids[:args.T - n_new - 1] + idx = torch.zeros(1, args.T, dtype=torch.long, device=dev) + L = len(ids) + idx[0, :L] = torch.tensor(ids, device=dev) + for _ in range(n_new): + if L >= args.T: break + z = emb(idx) + for b in blocks: z = b(z, mask) + lg = readout(z)[0, L - 1].float() / max(temp, 1e-6) + v, _i = torch.topk(lg, topk) + lg[lg < v[-1]] = -float('inf') + nt = torch.multinomial(F.softmax(lg, -1), 1).item() + idx[0, L] = nt; L += 1 + return _tk.decode(idx[0, :L].tolist()) + print(f'# samples from {args.resume} (temp {args.gen_temp}, top-k {args.gen_topk})', flush=True) + for pi, p in enumerate(prompts): + for s in range(args.gen): + print(f'\n--- prompt {pi + 1}, sample {s + 1} ---', flush=True) + print(_sample(p, args.gen_new, args.gen_temp, args.gen_topk, 1000 * pi + s), flush=True) + sys.exit(0) + if args.probe_dgspec > 0: # M1 DGAIN SPECTROSCOPY: per-block leak vector vs uniform read-displacement gain. # Paired design: L_l(g) = mean_b[gEP_l(g) - gBP_l] on the SAME batch — batch-sampling diff --git a/ep_run/fig_curves.py b/ep_run/fig_curves.py new file mode 100644 index 0000000..9185702 --- /dev/null +++ b/ep_run/fig_curves.py @@ -0,0 +1,94 @@ +"""Validation curves at 135M: EP with the readout fix, EP without it, and the backprop twins. + +The point of the figure is that the first and third are the same curve, and the second stops +descending partway through. Run after fw135m_rlin completes; it reads the training logs directly. +""" +import re +import statistics as st +from pathlib import Path + +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +RUNS = Path('/home/yurenh2/ept/ep_run/runs') +TOTAL = 440000 +TOK_PER_STEP = 24 * 256 # effective batch x context + + +def curve(tag): + pts = [] + for line in (RUNS / f'{tag}.log').read_bytes().decode(errors='replace').splitlines(): + m = re.match(rf'step (\d+)/{TOTAL} \| train [\d.]+ val ([\d.]+)', line) + if m: + pts.append((int(m.group(1)), float(m.group(2)))) + pts.sort() + return pts + + +def smooth(pts, win=25): + xs = [p[0] for p in pts] + ys = [p[1] for p in pts] + out = [] + for i in range(len(ys)): + lo, hi = max(0, i - win // 2), min(len(ys), i + win // 2 + 1) + out.append(st.median(ys[lo:hi])) + return xs, out + + +def main(): + # EP is drawn thick and underneath; the backprop seeds go on top as thin lines, so the reader + # sees them tracking rather than being hidden by whichever curve was plotted last. + arms = [ + ('fw135m_rlin', 'Equilibrium Propagation', '#1f6feb', 3.0, '-', 2), + ('fw135m_bsign', 'Equilibrium Propagation,\nwithout the readout fix', '#d1442f', 2.6, '-', 2), + ('fw135m_bp', 'Backprop, seed 1', '#2b3540', 1.0, '-', 4), + ('fw135m_bp_s2', 'Backprop, seed 2', '#7d8794', 1.0, '-', 4), + ] + fig, ax = plt.subplots(figsize=(7.6, 4.9)) + finals = {} + for tag, label, color, lw, ls, zo in arms: + pts = curve(tag) + if not pts: + print(f'skip {tag}: no data') + continue + xs, ys = smooth(pts) + toks = [x * TOK_PER_STEP / 1e9 for x in xs] + ax.plot(toks, ys, color=color, lw=lw, ls=ls, label=label, zorder=zo) + tail = [v for s, v in pts if s >= 0.9 * TOTAL] + finals[tag] = st.mean(tail) if tail else float('nan') + print(f'{tag:14s} last step {pts[-1][0]:>7} tail mean {finals[tag]:.4f}') + + ax.set_xlabel('training tokens (billions)') + ax.set_ylabel('validation cross-entropy') + ax.set_ylim(3.15, 4.35) + ax.set_xlim(0, TOTAL * TOK_PER_STEP / 1e9) + ax.grid(alpha=0.18, lw=0.7) + for side in ('top', 'right'): + ax.spines[side].set_visible(False) + h, l = ax.get_legend_handles_labels() + order = [2, 3, 0, 1] # backprop seeds first in the legend, matching how one reads the plot + leg = ax.legend([h[i] for i in order], [l[i] for i in order], + frameon=False, fontsize=9.5, loc='upper right', handlelength=1.6) + for t in leg.get_texts(): + t.set_va('center') + + if 'fw135m_bsign' in finals and 'fw135m_rlin' in finals: + ax.annotate('single-precision rounding in the contrast readout;\nthe run stops improving', + xy=(1.9, finals['fw135m_bsign'] + 0.005), xytext=(1.35, 3.86), + fontsize=9, color='#d1442f', + arrowprops=dict(arrowstyle='-', color='#d1442f', lw=0.9, alpha=0.8)) + ax.annotate(f"EP {finals['fw135m_rlin']:.3f}\nBP {st.mean([finals['fw135m_bp'], finals['fw135m_bp_s2']]):.3f}", + xy=(2.62, 3.245), fontsize=9.5, color='#2b3540', ha='right') + + ax.set_title('135M-parameter transformer language model, FineWeb-Edu', + fontsize=10.5, color='#333', pad=10, loc='left') + fig.tight_layout() + for ext in ('png', 'pdf'): + p = Path(f'/home/yurenh2/ept/assets/figs/fig_135m_curves.{ext}') + fig.savefig(p, dpi=300, bbox_inches='tight') + print('wrote', p) + + +if __name__ == '__main__': + main() |
