Five features, two dimensions
A small ReLU model learns to squeeze five features through two numbers. Make the features rarer and watch them start sharing directions.
How much features overlap (WTW) and bias
| Feature | 1 | 2 | 3 | 4 | 5 | bias |
|---|---|---|---|---|---|---|
| 1 | 0.6 | −0.1 | 0.0 | −0.1 | 0.3 | 0.00 |
| 2 | −0.1 | 0.0 | 0.0 | 0.0 | 0.0 | 0.00 |
| 3 | 0.0 | 0.0 | 0.9 | −0.9 | 0.4 | 0.00 |
| 4 | −0.1 | 0.0 | −0.9 | 1.0 | −0.5 | 0.00 |
| 5 | 0.3 | 0.0 | 0.4 | −0.5 | 0.3 | 0.00 |
Switch features on, read them back
0.00
0.00 dropped
0.00
0.00
0.00
input output
A dropped feature always reads back as the same number, its bias: the model's best guess without any information.
seed 1
StatusTraining, step 0
Features stored4 of 5
ArrangementTraining...
Stored features per dimension1.47
Challenges
0 of 4 complete
- With the ReLU on, get the model to store more features than it has dimensions: at least three in two. (not yet)
- Find a model that stores two features in exactly opposite directions. (not yet)
- With the ReLU on, fit all five features into the two dimensions. (not yet)
- With the ReLU on and three or more features stored, switch on two stored features at once so the model gets a stored feature wrong by more than 0.3. (not yet)
Try this
- Lower how often each feature is active, one step at a time, and count the stored arrows.
- Find two arrows that point in opposite directions, then switch both of those features on at once.
- At 5%, switch on two features that sit 144 degrees apart and look at the feature between them.
- Compare Steep importance with Equal, and try training again with a new seed.
- Switch off the ReLU. With falling or steep importance, can you get more than two features stored now? With equal importance, watch the stored features per dimension.
Train the toy model and measure its geometryPython with NumPy
import numpy as np
def train(density, n=5, m=2, steps=3000, batch=256, lr=0.02, seed=0):
"""Train x' = ReLU(W^T W x + b) on sparse features with Adam."""
rng = np.random.default_rng(seed)
W = rng.normal(0, np.sqrt(2 / (n + m)), (m, n)) # column i: feature i's direction
b = np.zeros(n)
importance = 0.9 ** np.arange(n)
params, moments = [W, b], [[np.zeros_like(W), np.zeros_like(W)], [np.zeros_like(b), np.zeros_like(b)]]
for t in range(1, steps + 1):
# Each feature is active with probability `density`, then uniform on [0, 1].
x = rng.random((batch, n)) * (rng.random((batch, n)) < density)
h = x @ W.T # squeeze 5 numbers into 2
pre = h @ W + b # read 5 numbers back out
out = np.maximum(pre, 0)
g = 2 * importance * (out - x) * (pre > 0) / (batch * n) # d loss / d pre
grads = [h.T @ g + (g @ W.T).T @ x, g.sum(0)]
step_lr = lr * (1 - (t - 1) / steps)
for p, grad, (mo, v) in zip(params, grads, moments):
mo[:] = 0.9 * mo + 0.1 * grad
v[:] = 0.999 * v + 0.001 * grad**2
p -= step_lr * (mo / (1 - 0.9**t)) / (np.sqrt(v / (1 - 0.999**t)) + 1e-8)
return W, b
for density in [1.0, 0.3, 0.05]:
W, b = train(density)
lengths = np.linalg.norm(W, axis=0)
stored = np.flatnonzero(lengths > 0.5)
angles = np.sort(np.degrees(np.arctan2(W[1], W[0]))[stored])
gaps = np.diff(np.append(angles, angles[0] + 360))
print(f"active {density:.0%}: stores features {(stored + 1).tolist()}")
print(f" angles between neighbors: {np.round(gaps).astype(int).tolist()}")active 100%: stores features [1, 2]
angles between neighbors: [90, 270]
active 30%: stores features [1, 2, 4, 5]
angles between neighbors: [90, 90, 90, 90]
active 5%: stores features [1, 2, 3, 4, 5]
angles between neighbors: [70, 71, 73, 73, 72]