Gradient Descent in Lenia
¶
Lenia is a continuous cellular automaton, and its update rule is built from smooth pieces —
convolutions, Gaussian bells, clipping. This notebook takes jax.grad through entire Lenia
simulations and uses it three ways: to find the fastest travelling creature a real Lenia
rule can support, to grow a target image, and to find the smallest perturbation that kills
a soliton.
Along the way it maps out what a gradient can and cannot reach in a Lenia rule — which is worth knowing before optimizing anything.
Installation¶
You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:
%pip install -U "jax[cuda]"
/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip
Note: you may need to restart the kernel to use updated packages.
Then, install CAX from PyPi:
%pip install -U "cax[examples]"
/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip
Note: you may need to restart the kernel to use updated packages.
Import¶
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import mediapy
import numpy as np
import optax
from flax import nnx
from cax.cs.lenia import (
FreeKernelParams,
Lenia,
LeniaGrowthParams,
LeniaKernelParams,
LeniaRuleParams,
exponential_kernel_fn,
free_kernel_fn,
gaussian_kernel_fn,
load_pattern,
)
from cax.cs.lenia.kernel import exponential_kernel_core, gaussian_kernel_core
from cax.utils import get_emoji_array
Configuration¶
seed = 0
spatial_dims = (96, 96)
channel_size = 1
R = 13
T = 10
state_scale = 1
key = jax.random.key(seed)
The specimen¶
Orbium is the smallest well-known Lenia creature: one channel, one kernel, and a growth mapping with two numbers. It travels in a straight line, which makes it the reference point for everything below.
def make_lenia(rule_params, **kwargs):
"""Build a Lenia system, with keyword overrides for the shared configuration."""
options = {
"spatial_dims": spatial_dims,
"channel_size": channel_size,
"R": R,
"T": T,
"state_scale": state_scale,
}
options.update(kwargs)
return Lenia(rule_params=rule_params, **options)
def sample_state(pattern, num_channels=channel_size, dims=spatial_dims):
"""Place a pattern at the center of an empty grid."""
mid = tuple(dim // 2 for dim in dims)
slices = tuple(
slice(m - c // 2, m + c - c // 2) for m, c in zip(mid, pattern.shape[:-1])
)
return jnp.zeros((*dims, num_channels)).at[slices].set(pattern)
orbium_pattern, orbium_rule_params = load_pattern("Orbium")
state_init = sample_state(orbium_pattern)
Simulations run through the system's own driver, cs(state, num_steps=...), which scans
the perceive/update pair and returns the final state — or the whole trajectory with
return_states=True. It composes with grad, jvp and jit, which is what the rest of
this notebook relies on.
def show_run(cs, state, *, num_steps=192, every=4, **kwargs):
"""Simulate and show the run as a video."""
_, states = cs(state, num_steps=num_steps, return_states=True)
states = jnp.concatenate([state[None], states])[::every]
frames = nnx.vmap(lambda cs, state: cs.render(state), in_axes=(None, 0))(cs, states)
mediapy.show_video(frames.repeat(2, axis=-3).repeat(2, axis=-2), fps=12, **kwargs)
return states
_ = show_run(make_lenia(orbium_rule_params), state_init)
What can a gradient reach?¶
A Lenia rule is a LeniaRuleParams pytree: which channel feeds which
(channel_source, channel_target), a mixing weight per kernel, kernel parameters
(r, beta), and growth parameters (mean, std).
jax.grad refuses integer inputs — and the channel wiring is integer. That is not a
technicality: the wiring of the rule graph is discrete, so a gradient can reshape a rule but
never rewire it. We split the pytree into its float leaves, which we differentiate, and its
integer leaves, which we carry along untouched.
def is_inexact(leaf):
"""Check whether a leaf holds floating point values."""
return jnp.issubdtype(jnp.asarray(leaf).dtype, jnp.inexact)
def partition(tree):
"""Split a pytree into its float leaves and its remaining leaves."""
return (
jax.tree.map(lambda x: x if is_inexact(x) else None, tree),
jax.tree.map(lambda x: None if is_inexact(x) else x, tree),
)
def combine(params, static):
"""Reassemble a pytree split by `partition`."""
return jax.tree.map(
lambda x, y: y if x is None else x, params, static, is_leaf=lambda x: x is None
)
params, static = partition(orbium_rule_params)
With that split, a whole simulation differentiates in one line. The loss here is just the mean of the state after 32 steps — enough to see which parameters the gradient reaches.
def mean_state(params, static, state, *, num_steps, **kwargs):
"""Compute the mean of the state after a rollout, as a scalar loss."""
cs = make_lenia(combine(params, static), **kwargs)
return jnp.mean(cs(state, num_steps=num_steps))
grads = jax.grad(mean_state)(params, static, state_init, num_steps=32)
for path, leaf in jax.tree_util.tree_flatten_with_path(grads)[0]:
print(f"{jax.tree_util.keystr(path):<22} {np.asarray(leaf).ravel()}")
.weight [0.] .kernel_params.r [0.01354898] .kernel_params.beta [-4.7857895e-10] .growth_params.mean [-0.00084772] .growth_params.std [0.28612816]
Reading this gradient teaches most of what there is to know:
growth_params.meanandgrowth_params.stdcarry a clean gradient. The growth mapping is a Gaussian bell — genuinely smooth.weightandbetaare exactly zero. Both are normalized before use — the weights by their sum, the kernel by its own integral — so scaling either changes nothing, and the gradient along that scale direction is exactly zero. Orbium has a single kernel with a single ring, so the scale direction is the whole parameter. In rules with several kernels or rings, the repartition between them carries gradient normally; only the overall scale of each normalized group is flat.kernel_params.rreturns a number — but it is lying. The next section shows why.
The kernel core decides whether r is differentiable¶
The kernel is built as mask * beta[segment] * kernel_core(position), where the mask is a
hard threshold radius < r. A hard threshold is only a problem if it cuts something — and
what it cuts is the value of the core at the edge of the support.
CAX ships several cores. The Gaussian core, which the catalogued creatures use, does not vanish at the edge. The exponential core — the canonical one from Chan's original Lenia paper — does.
positions = jnp.linspace(0.0, 1.0, 512)
print(f"{'core':<14}{'value at the support edge':>28}")
for name, core in [
("gaussian", gaussian_kernel_core),
("exponential", exponential_kernel_core),
]:
print(f"{name:<14}{float(core(jnp.array(1.0))):>28.6f}")
core value at the support edge
gaussian 0.003866
exponential 0.000000
The consequence: as r grows continuously, grid cells cross the threshold one at a time.
Under the Gaussian core each one pops into the kernel with weight 0.0039 — a jump. Under
the exponential core each one enters with weight zero — nothing happens.
Watching how much the normalized kernel changes for a small change in r makes this
visible directly. Every spike below is one lattice cell entering the support.
grid = jnp.mgrid[-48:48, -48:48] / R
distance = jnp.linalg.norm(grid, axis=0)
def kernel_of(r, kernel_fn):
"""Build the normalized kernel at radius `r`."""
kernel = kernel_fn(
distance, LeniaKernelParams(r=jnp.array(r), beta=jnp.array([1.0]))
)
return kernel / jnp.sum(kernel)
radii = jnp.linspace(0.95, 1.05, 401)
delta = 1e-4
fig, axes = plt.subplots(1, 2, figsize=(11, 3.5))
for name, kernel_fn in [
("gaussian", gaussian_kernel_fn),
("exponential", exponential_kernel_fn),
]:
profile = kernel_fn(
positions, LeniaKernelParams(r=jnp.array(1.0), beta=jnp.array([1.0]))
)
axes[0].plot(positions, profile, label=name)
response = [
float(
jnp.max(jnp.abs(kernel_of(r + delta, kernel_fn) - kernel_of(r, kernel_fn)))
/ delta
)
for r in radii
]
axes[1].plot(radii, response, label=name)
axes[0].set_xlabel("radius")
axes[0].set_ylabel("kernel value")
axes[0].set_title("Kernel profile — the Gaussian core does not vanish at the edge")
axes[0].legend()
axes[1].set_xlabel("r")
axes[1].set_ylabel(r"$\|\partial K / \partial r\|_\infty$")
axes[1].set_title("Kernel response to r — each spike is one cell entering the support")
axes[1].legend()
plt.tight_layout()
plt.show()
So plain, official Lenia is differentiable in every continuous parameter — including the kernel radius — provided the kernel core vanishes at its support boundary. The Gaussian core breaks that one property; the exponential core has it. No modification of Lenia is needed, only the right choice among its own cores.
The swap is drop-in — Orbium does not notice:
_ = show_run(
make_lenia(orbium_rule_params, kernel_fn=exponential_kernel_fn), state_init
)
Everything below uses the exponential core, and measures its baselines under it.
One more kernel is worth knowing about: free_kernel_fn, introduced in the Flow Lenia line
of work, parameterizes the kernel as a sum of Gaussian bumps with explicit position and
width parameters. It is not needed for differentiability — but when the kernel's shape is
the thing being optimized, explicit shape parameters are the right handle, and §3 uses it
for exactly that.
The fastest soliton¶
First use of the gradient: a race. Given this world and a fixed budget of steps, how far can a creature travel in a straight line without vanishing or exploding?
The score is the net displacement of the center of mass — the norm of the sum of per-step
displacements. A creature spinning in place accumulates a long path but a net displacement
of nearly zero, so the objective cannot be won by spinning; straightness is the ratio of
the two. Summing wrapped per-step displacements also keeps the distance well-defined on the
torus even after the creature laps it. The clock starts at step 0, so the settling phase is
part of the budget.
def center_of_mass(state):
"""Compute the circular center of mass, with a finite gradient at zero mass."""
mass_grid = jnp.sum(state, axis=-1)
centers = []
for axis in range(2):
mass_axis = jnp.sum(mass_grid, axis=1 - axis)
angles = 2 * jnp.pi * jnp.arange(state.shape[axis]) / state.shape[axis]
cosine = jnp.sum(mass_axis * jnp.cos(angles))
sine = jnp.sum(mass_axis * jnp.sin(angles))
# Double-`where` so the gradient stays finite when the creature dies.
is_alive = cosine**2 + sine**2 > 1e-8
angle = jnp.where(
is_alive,
jnp.arctan2(
jnp.where(is_alive, sine, 0.0), jnp.where(is_alive, cosine, 1.0)
),
0.0,
)
centers.append((angle % (2 * jnp.pi)) / (2 * jnp.pi) * state.shape[axis])
return jnp.stack(centers)
def wrap(delta, size):
"""Wrap a grid offset into (-size/2, size/2]."""
return delta - jnp.round(delta / size) * size
def spread_of(state):
"""Compute the mass-weighted RMS radius about the center of mass, in units of R."""
mass_grid = jnp.sum(state, axis=-1)
total = jnp.sum(mass_grid)
center = center_of_mass(state)
dy = wrap(jnp.arange(state.shape[0]) - center[0], state.shape[0])[:, None]
dx = wrap(jnp.arange(state.shape[1]) - center[1], state.shape[1])[None, :]
is_alive = total > 1e-6
variance = jnp.sum(mass_grid * (dy**2 + dx**2)) / jnp.where(is_alive, total, 1.0)
return jnp.where(is_alive, jnp.sqrt(variance), 0.0) / R
def trajectory_stats(states):
"""Compute displacement, straightness, mass and spread over a stack of states."""
centers = jax.vmap(center_of_mass)(states)
steps = wrap(jnp.diff(centers, axis=0), jnp.array(states.shape[1:3], jnp.float32))
net = jnp.linalg.norm(jnp.sum(steps, axis=0)) / R
path = jnp.sum(jnp.linalg.norm(steps, axis=-1)) / R
return {
"net": net,
"straightness": net / jnp.maximum(path, 1e-6),
"mass": jnp.sum(states, axis=(1, 2, 3)) / R**2,
"spread": jax.vmap(spread_of)(states),
}
def speed_of(net, num_steps):
"""Convert a net displacement in units of R into R per unit of Lenia time."""
return net * T / num_steps
The search space¶
Orbium's own rule, made differentiable in place. Its weight and beta are pure gauge, so
the rule contributes three numbers: the kernel radius (through a sigmoid, so the kernel can
never outgrow the neighbourhood it is defined on) and the two growth parameters (through
exponentials, so one Adam step is the same relative move for parameters two orders of
magnitude apart). The initial state is optimized too: a patch of logits behind a sigmoid,
masked by a disk so the seed starts localized.
patch_size = 40
offsets = jnp.arange(patch_size) - (patch_size - 1) / 2
disk = (
(offsets[:, None] ** 2 + offsets[None, :] ** 2) <= (patch_size / 2) ** 2
).astype(jnp.float32)[..., None]
def logit(p, eps=1e-4):
"""Invert the logistic function, clipped away from its poles."""
p = jnp.clip(p, eps, 1 - eps)
return jnp.log(p / (1 - p))
def build_rule(params):
"""Assemble Orbium's rule with the differentiable parameters in place."""
return LeniaRuleParams(
channel_source=orbium_rule_params.channel_source,
channel_target=orbium_rule_params.channel_target,
weight=orbium_rule_params.weight,
kernel_params=LeniaKernelParams(
r=jax.nn.sigmoid(params["r"]),
beta=orbium_rule_params.kernel_params.beta,
),
growth_params=LeniaGrowthParams(
mean=jnp.exp(params["log_growth_mean"]),
std=jnp.exp(params["log_growth_std"]),
),
)
def build_seed(params):
"""Assemble the optimized initial state."""
return sample_state(jax.nn.sigmoid(params["seed"]) * disk)
def simulate(params, *, num_steps):
"""Roll the candidate forward from its optimized seed."""
cs = make_lenia(build_rule(params), kernel_fn=exponential_kernel_fn)
return cs(build_seed(params), num_steps=num_steps, return_states=True)
params_init = {
"r": logit(orbium_rule_params.kernel_params.r),
"log_growth_mean": jnp.log(orbium_rule_params.growth_params.mean),
"log_growth_std": jnp.log(orbium_rule_params.growth_params.std),
"seed": logit(sample_state(orbium_pattern, dims=(patch_size, patch_size))),
}
The objective, and why candidates are re-simulated¶
Maximize net displacement over a 128-step budget, with two hinge penalties: the spread must stay under one kernel radius (no smearing into a travelling wave) and the mass must stay in a band (no evaporating, no exploding). One full backpropagation runs through all 128 steps.
The budget alone is not enough. An earlier version of this search selected candidates on their training window and found a creature that scored 0.97 — double Orbium — by sprinting and then evaporating at step 200, just after the window closed. Nothing rewarded dying; the budget simply ended first, and dying pays, because the balance between growth at the front and decay at the back is exactly what caps a soliton's speed. So every candidate is re-simulated for 448 steps, without gradients, and only counts if it is still a localized, straight, travelling creature at the end.
budget = 128
long_run = 448
def objective(params):
"""Compute the negative net displacement, with localization and survival hinges."""
_, states = simulate(params, num_steps=budget)
stats = trajectory_stats(states)
spread_penalty = jnp.mean(jax.nn.relu(stats["spread"] - 1.0) ** 2)
mass_penalty = jnp.mean(
jax.nn.relu(0.2 - stats["mass"]) ** 2 + jax.nn.relu(stats["mass"] - 2.0) ** 2
)
return -stats["net"] + 30.0 * spread_penalty + 20.0 * mass_penalty
def evaluate(params):
"""Score a candidate on the long run and check that it survives it."""
_, states = simulate(params, num_steps=long_run)
stats = trajectory_stats(states)
report = {
"speed": float(speed_of(stats["net"], long_run)),
"straightness": float(stats["straightness"]),
"mass_min": float(jnp.min(stats["mass"])),
"spread_max": float(jnp.max(stats["spread"])),
}
report["valid"] = bool(
np.isfinite(report["speed"])
and report["mass_min"] > 0.1
and report["spread_max"] < 1.0
and report["straightness"] > 0.9
)
return report
Training¶
Adam with gradient clipping and a cosine-decayed learning rate. Two habits matter here,
both learned from failed runs: keep the best surviving candidate rather than the last
one, because these optima are sharp and a run often ends worse than its middle; and when a
round diverges to nan, restart from the last good checkpoint at a lower rate instead of
giving up.
def train(params, *, rounds=6, iterations=150, learning_rate=4e-3):
"""Run rounds of Adam with back-off, keeping the best surviving candidate."""
best = (evaluate(params), params)
print(f"Start | Speed {best[0]['speed']:.4f}")
for round_index in range(rounds):
schedule = optax.cosine_decay_schedule(learning_rate, iterations, alpha=0.2)
optimizer = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(schedule))
opt_state = optimizer.init(params)
@jax.jit
def step(params, opt_state, optimizer=optimizer):
loss, grads = jax.value_and_grad(objective)(params)
updates, opt_state = optimizer.update(grads, opt_state, params)
return optax.apply_updates(params, updates), opt_state, loss
found = None
for iteration in range(iterations):
params, opt_state, loss = step(params, opt_state)
if not jnp.isfinite(loss):
break
if iteration % 25 == 24:
report = evaluate(params)
if report["valid"] and (
found is None or report["speed"] > found[0]["speed"]
):
found = (report, jax.tree.map(jnp.copy, params))
if found is None:
params = jax.tree.map(lambda x: x, best[1])
learning_rate *= 0.4
print(
f"Round {round_index}/{rounds} | diverged, "
f"backing off to {learning_rate:.1e}"
)
continue
if found[0]["speed"] > best[0]["speed"]:
best = found
report, params = found
print(
f"Round {round_index}/{rounds} | Speed {report['speed']:.4f} "
f"| Straightness {report['straightness']:.3f} "
f"| Spread {report['spread_max']:.2f}"
)
return best
baseline = evaluate(params_init)
report, params_best = train(params_init)
print(
f"\n✨ Orbium {baseline['speed']:.4f} -> trained {report['speed']:.4f} "
f"({report['speed'] / baseline['speed'] - 1:+.1%})"
)
Start | Speed 0.4720
Round 0/6 | Speed 0.5259 | Straightness 1.000 | Spread 0.49
Round 1/6 | Speed 0.5656 | Straightness 1.000 | Spread 0.51
Round 2/6 | Speed 0.5654 | Straightness 1.000 | Spread 0.50
Round 3/6 | Speed 0.5645 | Straightness 1.000 | Spread 0.51
Round 4/6 | diverged, backing off to 1.6e-03
Round 5/6 | Speed 0.5690 | Straightness 1.000 | Spread 0.51 ✨ Orbium 0.4720 -> trained 0.5690 (+20.6%)
rule_best = build_rule(params_best)
print(f"{'parameter':<16}{'Orbium':>10}{'trained':>10}")
print(
f"{'kernel r':<16}{float(orbium_rule_params.kernel_params.r[0]):>10.4f}"
f"{float(rule_best.kernel_params.r[0]):>10.4f}"
)
print(
f"{'growth mean':<16}{float(orbium_rule_params.growth_params.mean[0]):>10.4f}"
f"{float(rule_best.growth_params.mean[0]):>10.4f}"
)
print(
f"{'growth std':<16}{float(orbium_rule_params.growth_params.std[0]):>10.4f}"
f"{float(rule_best.growth_params.std[0]):>10.4f}"
)
parameter Orbium trained kernel r 1.0000 0.9999 growth mean 0.1500 0.1309 growth std 0.0150 0.0174
_ = show_run(
make_lenia(rule_best, kernel_fn=exponential_kernel_fn), build_seed(params_best)
)
About 20% faster than Orbium, verified over 448 steps — and the interesting part is how little changed. The radius is untouched; the entire speed-up lives in the growth mapping, whose mean drops and whose band widens. Cells begin growing at a thinner local density, so the creature starts building its next body further out in front of itself. Orbium is a well-tuned creature — it is just not tuned for speed.
Growing a picture¶
Second use: expressivity. A Lenia rule is a few dozen numbers; a picture is thousands of pixels. Can gradient descent find a rule under which an arbitrary image is a living pattern?
This section works in colour — three coupled channels on a smaller grid — and optimizes the
kernel shapes, which is what free_kernel_fn is for: each kernel is a Gaussian bump with
an explicit position a and width w, under a soft support mask. The channel wiring (two
self-kernels per channel, one kernel per ordered pair of channels) is discrete, so it is
chosen by hand; everything else is learned.
image_dims = (64, 64)
image_channels = 3
target = get_emoji_array("🦎", size=40, pad_width=12)[..., :3]
mediapy.show_image(np.asarray(target).repeat(3, axis=-3).repeat(3, axis=-2))
pairs = [(i, j) for i in range(image_channels) for j in range(image_channels) if i != j]
image_source = jnp.array(
[c for c in range(image_channels) for _ in range(2)] + [i for i, _ in pairs],
jnp.int32,
)
image_target = jnp.array(
[c for c in range(image_channels) for _ in range(2)] + [j for _, j in pairs],
jnp.int32,
)
num_kernels = image_source.shape[0]
def build_image_rule(params):
"""Assemble a free-kernel rule over the three channels."""
return LeniaRuleParams(
channel_source=image_source,
channel_target=image_target,
weight=jax.nn.softmax(params["weight"]),
kernel_params=FreeKernelParams(
r=jnp.ones((num_kernels,)),
b=jnp.ones((num_kernels, 1)),
a=jax.nn.sigmoid(params["a"]),
w=jnp.exp(params["log_w"]),
),
growth_params=LeniaGrowthParams(
mean=jnp.exp(params["log_growth_mean"]),
std=jnp.exp(params["log_growth_std"]),
),
)
def sample_image_params(key):
"""Sample initial parameters for the image rule."""
keys = jax.random.split(key, 4)
return {
"weight": jnp.zeros((num_kernels,)),
"a": 0.5 * jax.random.normal(keys[0], (num_kernels, 1)),
"log_w": jnp.log(
jax.random.uniform(keys[1], (num_kernels, 1), minval=0.05, maxval=0.25)
),
"log_growth_mean": jnp.log(
jax.random.uniform(keys[2], (num_kernels,), minval=0.05, maxval=0.3)
),
"log_growth_std": jnp.log(
jax.random.uniform(keys[3], (num_kernels,), minval=0.01, maxval=0.1)
),
}
def evolve(params, state, *, num_steps=32):
"""Roll an image rule forward from a given state."""
cs = make_lenia(
build_image_rule(params),
spatial_dims=image_dims,
channel_size=image_channels,
kernel_fn=free_kernel_fn,
)
return cs(state, num_steps=num_steps)
def r_squared(state):
"""Compute the fraction of the target's variance explained by a state."""
return float(1 - jnp.mean(jnp.square(state - target)) / jnp.var(target))
def fit_image_rule(loss_fn, key, *, iterations=250):
"""Fit an image rule by Adam on a given loss."""
params = sample_image_params(key)
optimizer = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(1e-2))
opt_state = optimizer.init(params)
@jax.jit
def step(params, opt_state):
loss, grads = jax.value_and_grad(loss_fn)(params)
updates, opt_state = optimizer.update(grads, opt_state, params)
return optax.apply_updates(params, updates), opt_state, loss
for _ in range(iterations):
params, opt_state, _ = step(params, opt_state)
return params
First attempt: hold the picture¶
The target is the initial state, and the loss asks for it to still be there 32 steps later.
key, subkey = jax.random.split(key)
hold_params = fit_image_rule(
lambda params: jnp.mean(jnp.square(evolve(params, target) - target)), subkey
)
for num_steps in [32, 128, 256]:
held_state = evolve(hold_params, target, num_steps=num_steps)
print(f"R2 after {num_steps:>3} steps: {r_squared(held_state):+.3f}")
print(
"\nfitted growth std:", np.asarray(jnp.exp(hold_params["log_growth_std"])).round(3)
)
print(
"a living creature's, for comparison:",
float(orbium_rule_params.growth_params.std[0]),
)
W0904 22:14:00.121391 521456 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.
R2 after 32 steps: +0.989 R2 after 128 steps: +0.813
R2 after 256 steps: +0.238 fitted growth std: [0.134 0.099 0.096 0.113 0.117 0.058 0.08 0.034 0.018 0.134 0.054 0.034] a living creature's, for comparison: 0.014999999664723873
A high score at the trained horizon that decays steadily beyond it — because the optimizer cheated. There is a trivial way to keep a picture where it is: stop the dynamics. Most of the fitted growth bands come out several times wider than a living creature's, which makes the growth mapping nearly flat, and the picture persists because nothing happens to it.
The fix is to give the rule work to do: erase half the image and ask the rule to put it back, while still holding the intact version. A do-nothing rule now scores only whatever the half-image scores.
corrupted = target * (jnp.arange(image_dims[0]) < image_dims[0] // 2)[:, None, None]
print(f"a rule that does nothing scores R2 = {r_squared(corrupted):.3f}")
def regrow_loss(params):
"""Ask the rule to regrow the target from a corrupted seed, and hold the intact."""
return (
jnp.mean(jnp.square(evolve(params, corrupted) - target))
+ jnp.mean(jnp.square(evolve(params, target) - target))
) / 2
key, subkey = jax.random.split(key)
regrow_params = fit_image_rule(regrow_loss, subkey)
held = evolve(regrow_params, target)
regrown = evolve(regrow_params, corrupted)
print(f"R2 hold = {r_squared(held):+.3f}")
print(f"R2 regrow = {r_squared(regrown):+.3f}")
a rule that does nothing scores R2 = 0.486
R2 hold = +0.974 R2 regrow = +0.650
images = {
"target": np.asarray(target),
"held, 32 steps": np.asarray(jnp.clip(held, 0.0, 1.0)),
"corrupted seed": np.asarray(corrupted),
"regrown, 32 steps": np.asarray(jnp.clip(regrown, 0.0, 1.0)),
}
mediapy.show_images(
{
title: image.repeat(3, axis=-3).repeat(3, axis=-2)
for title, image in images.items()
},
columns=4,
)
target | held, 32 steps | corrupted seed | regrown, 32 steps |
Holding an arbitrary image is easy — three coupled channels park the lizard as a
near-stationary state almost exactly. Regrowing is much harder: the rule pulls mass back
into the erased half, clearly beating the do-nothing baseline, but the recovered half is a
soft approximation rather than a tail. A Lenia rule mostly controls which configurations
persist; controlling which configurations get reached is what neural cellular automata
spend thousands of parameters on (see 40_growing_nca.ipynb).
How fragile is a creature?¶
Last use of the gradient: point it at the creature. Orbium survives 128 steps with 97% of its mass. What is the smallest perturbation of its initial state that kills it?
This is projected gradient descent — the construction behind adversarial examples for image classifiers — descending on the surviving mass, with the perturbation projected back into an L2 ball of radius epsilon after every step. Random perturbations of the same size are the control.
mass_init = float(jnp.sum(state_init))
seed_norm = float(jnp.linalg.norm(state_init))
@jax.jit
def surviving_mass(delta):
"""The fraction of the initial mass left 128 steps after a perturbed start."""
cs = make_lenia(orbium_rule_params, kernel_fn=exponential_kernel_fn)
final_state = cs(jnp.clip(state_init + delta, 0.0, 1.0), num_steps=128)
return jnp.sum(final_state) / mass_init
def project(delta, eps):
"""Project a perturbation into the L2 ball of radius `eps`."""
norm = jnp.linalg.norm(delta)
return jnp.where(norm > eps, delta * eps / norm, delta)
@jax.jit
def attack_step(delta, eps):
"""Take one projected gradient step against the creature's survival."""
grads = jax.grad(surviving_mass)(delta)
return project(delta - (eps / 10) * grads / (jnp.linalg.norm(grads) + 1e-12), eps)
def attack(eps, key, *, iterations=60, restarts=2):
"""Find the deadliest perturbation of norm `eps` by restarted PGD."""
best = None
for subkey in jax.random.split(key, restarts):
delta = project(
1e-3 * jax.random.normal(subkey, state_init.shape), jnp.array(eps)
)
for _ in range(iterations):
delta = attack_step(delta, jnp.array(eps))
if best is None or surviving_mass(delta) < surviving_mass(best):
best = delta
return best
print(f"{'eps':>6}{'% of seed norm':>16}{'adversarial':>13}{'random (5)':>12}")
for eps in [1.0, 2.0, 4.0, 8.0]:
key, attack_key, noise_key = jax.random.split(key, 3)
adversarial = float(surviving_mass(attack(eps, attack_key)))
random_masses = [
float(surviving_mass(eps * noise / jnp.linalg.norm(noise)))
for noise in jax.random.normal(noise_key, (5, *state_init.shape))
]
print(
f"{eps:>6.1f}{eps / seed_norm:>15.0%}{adversarial:>13.4f}"
f"{np.mean(random_masses):>12.4f}"
)
eps % of seed norm adversarial random (5)
1.0 15% 0.0000 0.9465
2.0 30% 0.0000 0.9464
4.0 61% 0.0000 0.9456
8.0 122% 0.0000 0.7498
key, subkey = jax.random.split(key)
delta = attack(2.0, subkey)
fig, axes = plt.subplots(1, 3, figsize=(10, 3.5))
limit = float(jnp.max(jnp.abs(delta)))
axes[0].imshow(state_init[..., 0], cmap="viridis")
axes[0].set_title("Orbium seed")
axes[1].imshow(delta[..., 0], cmap="RdBu_r", vmin=-limit, vmax=limit)
axes[1].set_title("lethal perturbation")
axes[2].imshow(jnp.clip(state_init + delta, 0.0, 1.0)[..., 0], cmap="viridis")
axes[2].set_title("perturbed seed")
for axis in axes:
axis.axis("off")
plt.tight_layout()
plt.show()
_ = show_run(
make_lenia(orbium_rule_params, kernel_fn=exponential_kernel_fn),
jnp.clip(state_init + delta, 0.0, 1.0),
)
Orbium is robust to noise and fragile to intent. Random perturbations carrying as much energy as the creature itself still leave it alive; the adversarial direction kills it with a nudge a tenth that size, one the eye can barely find on the seed. The perturbed creature runs for a while, then unravels.
Takeaways¶
- Official Lenia is differentiable end to end — through the convolution, the growth mapping, the clip, and hundreds of steps of dynamics — as long as the kernel core vanishes at its support boundary. The exponential core does; the Gaussian core is the one exception.
- The gradient cannot reach everything: the channel wiring is discrete, and normalized parameters have one flat scale direction each.
- Death is an absorbing state with no gradient out of it. Every optimization above had to be built around that fact: start from a living creature, keep the best surviving checkpoint, verify far beyond the training window — because an objective enforced on a window is enforced on that window only.