My failed attempt at AGI on the Tokio Runtime
Note: I am not, nor do I claim to be an expert in machine learning or neuroscience. This will become abundantly obvious as you continue reading.
A few weeks ago I decided to try and build AGI. OpenAI, Deepmind and xAI haven't delivered yet with the smartest researchers and billions in compute so I had to take matters into my own hands (what could go wrong?).
I bought a couple of books on Artificial Intelligence and Neuroscience and started:
cargo new agi
Meta Strategy
Assume you are racing a Formula 1 car. You are in last place. You are a worse driver in a worse car. If you follow the same strategy as the cars in front of you, pit at the same time and choose the same tires, you will certainly lose. The only chance you have is to pick a different strategy.
The same goes for me. If I go down the transformer / deep learning route I am outgunned. The only hope I have is to try something completely novel (or more precisely think I'm working on something novel only to discover this was done in the 1970s1).
Concrete Strategy
For reasons we'll cover in the following sections, I decided to go down the fully biologically inspired path. The rough idea was to build a fully asynchronous neural network and run it on a data center.
Neurons and Brains
When I started reading neuroscience I got the impression we don't really understand how the brain works. It's complicated and complex and the books I read model neuronal firing as partial differential equations. But before that a small primer.
At a high level a neuron consists of 3 main components.
The dendrites on the left act as inputs to the neuron from other neurons (we'll call those "pre-synaptic neurons"). The cell body has a cell wall which acts as a barrier between the internals of the neuron and the goop surrounding it. The axon on the right is connected to dendrites of other downstream neurons (we'll call those post-synaptic neurons).
When a neuron receives a signal from a pre-synaptic neuron it increases the potential in the neuron's cell body. If this potential increases past some threshold voltage (relative to the surrounding goop) it triggers a response where the neuron fires a signal down its axon to the post-synaptic neurons and resets its internal voltage. After firing a neuron has a rest period called the "refactory" period during which it does not respond to stimuli. After the refactory period the neuron is ready to fire again.
This is massively simplified. There are different types of neurons, a bunch of chemistry but I'm going to hand-wave those away and call them "implementation details". In fact I'm going to assume that the continuous nature of the signals fired is an implementation detail due to the substrate i.e. the underlying biological wetware and is not functionally important. There is no reason why signals can't be binary.
Conductance-Based Models
I didn't mention earlier that the cell body leaks potential into the surrounding goop over time. In 1963 Alan Hodgkin and Andrew Huxley received the Nobel Prize in Physiology and Medicine for describing this as a dynamical system described by a series of nonlinear differential equations. They modelled the relationship between the flow of ions across the neuron's cell membrane and the voltage of the cell. The experimental work for this was done on the squid giant axon because it was large enough for an electrode to be placed inside it.
Again I'm going to hand wave the chemistry away and call it an implementation detail using a simplified "Leaky integrate and fire" model.
This is also a differential equation over the capacitance, resistance and current across the neuron membrane and voltage of the cell. But really it boils down to:
- Pre-synaptic impulses increase membrane potential
- Time decreases membrane potential
Or in pseudocode:
let k = ... // some decay constant
let delta = ... // some potential difference constant
loop {
if signal.next() {
let now = time::now()
membrane_potential = membrane_potential * e^-k(now - previous_firing)
membrane_potential += delta
if membrane_potental > firing_threshold {
fire()
membrane_potential = 0
previous_firing = now
}
}
}
Encoding Information in Neuronal Signals
It looks like the jury is still out on how exactly neurons encode information. Namely is information encoded in neuron timings, i.e. when a neuron fires, or neuron firing rates, the rate at which a neuron fires. There's a bunch of statistics and math that's been developed to talk intelligently about neuronal firing rates, but I'm going to assume that I don't care because the firing rates are going to be emergent from the underlying neuron timings anyway.
Design
Meditating on the structure of a neuron described above and modern artificial neural networks like transformers, a few questions jump out at you.
Even if a network of these neurons is not being driven externally, there are certain configurations which allow for signals to propagate in cycles in your neuronal graph. There are configurations which sustain themselves without needing external stimuli to drive it while at the same time not having divergent outputs.
This is far-fetched but it feels like something that might implement consciousness rather than a pure feed-forward system.
Implementation
I decided to implement this network by employing something like an Actor Model on the Tokio runtime. Tokio is fast asynchronous runtime for Rust and exposes primitives which would make my life easier such as broadcast channels to implement synapses. Also it would be easy to hot-swap it for a non-local version if I want to run my AI across multiple machines.
Neurons
Neurons are implemented pretty much as described above.
pub struct Neuron {
#[allow(unused)]
index: usize,
membrane_potential: u32,
axon: broadcast::Sender<Impulse>,
dendrites: Vec<broadcast::Receiver<Impulse>>,
}
A broadcast::Sender is used to broadcast signals to post-synaptic neurons and signals from the pre-synaptic neurons which are broadcast::Receiver are used to drive the neuron.
An Impulse is just an empty tuple for now - we are assuming that the signal potential isn't important (or is constant) and information is encoded purely in the timing of firings (and consequently the firing rates).
To run the neuron we combine the dendrite receivers into a single stream and keep popping them implementing the leaky integrate and fire method:
impl Neuron {
async fn start(mut self) {
// Convert each receiver to a stream of messages
let streams = self
.dendrites
.drain(..)
.map(|mut rx| {
Box::pin(async_stream::stream! {
loop {
match rx.recv().await {
Ok(msg) => yield msg,
Err(broadcast::error::RecvError::Closed) => break,
Err(broadcast::error::RecvError::Lagged(skipped)) => {
// debug!("Receiver lagged by {} messages", skipped);
continue;
}
}
}
})
})
.collect::<Vec<_>>();
// Combine all streams into a single unified stream
let mut combined = stream::select_all(streams);
let mut last_fire = Instant::now();
// Process each message as it arrives from any receiver
while let Some(impulse) = combined.next().await {
let firings = FIRINGS.fetch_add(1, Ordering::Relaxed);
// Implement the "Integrate and fire" method.
let now = Instant::now();
if last_fire + ABSOLUTE_REFACTORY_PERIOD > now {
self.membrane_potential = self.membrane_potential + 1;
if self.membrane_potential > FIRING_THRESHOLD {
self.emit(Impulse);
self.membrane_potential = 0;
last_fire = now;
}
}
}
}
fn emit(&self, impulse: Impulse) {
if let Err(e) = self.axon.send(impulse) {
println!("{}", FIRINGS.fetch_add(0, Ordering::Relaxed));
panic!()
}
}
}
Brains
Brains are modelled as a bag of neurons with a set of inputs and outputs.
pub struct Brain {
neurons: Vec<Neuron>,
inputs: Vec<broadcast::Sender<Impulse>>,
outputs: Vec<broadcast::Receiver<Impulse>>,
}
The synapses for the neurons are already constructed beforehand as a brain is built from DNA.
impl From<&Dna> for Brain {
fn from(dna: &Dna) -> Self {
let mut neurons = Vec::new();
let mut broadcasts = Vec::new();
// Step 1: Initialize neurons and broadcast channels
for index in 0..Dna::num_neurons() {
let (tx, rx) = broadcast::channel(CHANNEL_CAPACITY);
neurons.push(Neuron {
membrane_potential: 0,
axon: tx.clone(),
dendrites: Vec::new(),
});
broadcasts.push((tx, rx));
}
let connectivity = dna.connectivity();
for (src, row) in connectivity.iter().enumerate() {
for (dest, &value) in row.iter().enumerate() {
if src == dest {
// do not allow neurons to wire back to themselves
continue;
}
if value == 1 {
let receiver = broadcasts[src].0.subscribe();
neurons[dest].dendrites.push(receiver);
}
}
}
let inputs = dna
.inputs()
.iter()
.map(|input_id| broadcasts[*input_id].0.clone())
.collect::<Vec<_>>();
let outputs = dna
.outputs()
.iter()
.map(|output_id| broadcasts[*output_id].0.subscribe())
.collect::<Vec<_>>();
Brain {
neurons,
inputs,
outputs,
}
}
}
DNA
The average brain of a human being has 85 billion neurons and over 100 trillion synaptic connections. If every neuron is connected to every other neuron you get synapses. Even in a sparsely connected brain you still get an unfeasibly large number of synapses for my 64 Gb RAM (neurons are thought to have 1,000-100,000 connections typically, depending to the type of neuron, its location etc.)
The sheer number of neurons and synapses mean that they are not deterministically encoded in your DNA. Instead your DNA defines rules for protein synthesis which generate these neurons and synapses.
This seems hard. I'm going to go down the road of the C. Elegans. nematode with exactly 302. I'm not sure I understand if its synapses are hard wired but mine will be.
pub struct Dna<const NUM_NEURONS: usize, const NUM_INPUT: usize, const NUM_OUTPUT: usize> {
potential_decay_ns: f64,
threshold: u16,
initiation_delay_ns: u64,
connectivity: Box<[[u8; NUM_NEURONS]; NUM_NEURONS]>,
// point to the input neurons of the connectivity matrix.
input_neurons: [usize; NUM_INPUT],
// point to the output neurons of the connectivity matrix.
output_neurons: [usize; NUM_OUTPUT],
}
We define a hard-coded connectivity matrix in our brain's DNA. The inputs and outputs point to specific neurons in the brain irrespective of positioning.
Games
Our brain is going to try to get better at playing a simple game I created for it. The game is basically snake. Your score increases every time you eat food. You can only go up, down, left and right. A higher score is better.
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Direction {
Up,
Down,
Left,
Right,
}
#[derive(Clone, PartialEq)]
pub struct Position {
x: i32,
y: i32,
}
pub struct Game {
pub width: usize,
pub height: usize,
pub snake: Position,
pub direction: Direction,
pub food: Position,
pub(crate) score: usize,
pub game_over: bool,
}
Organism
In order for our brain to play this game, it needs to be wrapped up in an organism. The organism is responsible for driving the inputs of the brain by reading the game state and playing the game using the brain's outputs.
The brain is constantly driven by the organism being fed the game's state even if it hasn't changed (much like you keep seeing an image in front of you even if it hasn't changed).
pub struct Organism {
pub(crate) dna: Dna,
inputs: Vec<broadcast::Sender<Impulse>>,
outputs: Vec<broadcast::Receiver<Impulse>>,
}
impl Organism {
pub fn new(dna: Dna) -> Organism {
let brain = Brain::from(&dna);
let (inputs, outputs) = brain.start();
Self {
dna,
inputs,
outputs,
}
}
// Given a 2D representation of the world state
// stimulates the appropriate input neurons.
pub(crate) fn drive_input(&self, state: Vec<Vec<u8>>) {
for (i, row) in state.iter().enumerate() {
for (j, val) in row.iter().enumerate() {
match val {
0 => continue,
_ => {
let index = i * row.len() + j;
self.inputs
.get(index)
.unwrap()
.send(Impulse)
.expect(&format!("Failed at index {}", index));
}
}
}
}
}
...
Training
Ok how the hell do we train this thing? Stochastic gradient descent with back-propagation won't work here (or if it does I have no idea how to implement it).
Instead I resorted to using genetic algorithms. Genetic algorithms are a class of optimisation algorithms inspired by nature using a combination of genetic darwinian selection based on individual fitness along with a small probability of genetic mutation to help explore the domain's search space and escape from local minima.
To do this for our Tokio brains requires a few steps:
- Initialise a population of DNA with random connectivity matrices
- Create brains from the DNA and put those brains in organisms and let them play our game.
- The individuals with the highest scores are bred with each other resulting in a new population.
- Breeding works by splitting the connectivity matrix into sections and randomly picking sections from each parent (along with any other relevant genes)
- Repeat
- Profit
pub fn train(&mut self) {
info!("Starting training.");
let mut population = self.initialize_population();
while self.epoch < self.max_epoch {
let runtime = tokio::runtime::Runtime::new().unwrap();
runtime.block_on(async {
info!("Starting epoch: {}", self.epoch);
let mut handles = vec![];
for (id, dna) in population.iter().enumerate() {
let dna = dna.clone();
let handle = tokio::spawn(async move { Simulation::simulate(id, dna).await });
handles.push(handle);
}
let population_with_scores = join_all(handles)
.await
.into_iter()
.filter_map(|handle| match handle {
Ok(dna_and_score) => Some(dna_and_score),
Err(e) => {
error!("{}", e);
None
}
})
.collect::<Vec<_>>();
let top_score = population_with_scores
.iter()
.map(|pop_with_score| pop_with_score.1)
.max();
info!("Epoch: {}, Top Score: {:?}", self.epoch, top_score);
population = self.reproduce_top_performers(population_with_scores);
println!("{}", population.get(0).unwrap());
self.epoch += 1;
});
}
}
Results
Nothing. Nada. I couldn't get this to work at all past a score of 3 which would disappear in the next epoch!
For reference, a human easily gets arbitrarily high scores. My brains have 512 neurons with up to ~13,000 synapses. I'm not sure if this is due to the lack of Neurons but I doubt it.
If I had to guess I would say the culprits are:
- A huge number of impulses being generated means that tokio struggled to process them all in a timely manner and these neurons are timing sensitive.
- Trying to do optimisation over a connectivity matrix by breaking it down into small chunks probably doesn't work.
Mother nature has defeated me once more. I'm going to put this project on ice for now. I'm going to continue reading neuroscience and pick it back up if / when inspiration strikes.
I later found out that what I was building has been known for at least 50 years and is called a spiking neural network.↩