path: root/assets/EPT_135M_samples.ipynb

{
 "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
}