Toy Models of Superposition: Feature Reconstruction

MediumMachine Learning

In the toy models of superposition setup, a neural network compresses a high-dimensional sparse feature vector into a lower-dimensional hidden space and then attempts to reconstruct the original features. The encoder-decoder uses tied weights: the same matrix W is used to encode and its transpose to decode, followed by a bias and a ReLU nonlinearity.

Given:

W: a 2D array of shape (n_hidden, n_features) representing the tied encoder/decoder weight matrix.
b: a 1D array of shape (n_features,) representing the output bias.
X: a 2D array of shape (batch_size, n_features) representing a batch of input feature vectors.

Implement the function superposition_reconstruct(W, b, X) that computes the reconstructed features for each input row using:

$\hat{x} = \mathrm{ReLU}(W^\top W x + b)$

Return the reconstruction as a Python list of lists (one inner list per input row). Do not round inside the function.

Example:

Input:

W = [[1.0, 1.0]]
b = [-0.5, -0.5]
X = [[1.0, 0.0]]

Output:

[[0.5, 0.5]]

Reasoning:

Both features share the single hidden direction. The hidden activation is W @ x = [1.0]. Decoding gives W^T @ [1.0] = [1.0, 1.0]. Adding the bias [-0.5, -0.5] yields [0.5, 0.5], and ReLU leaves both values unchanged. The model recovers the active feature (index 0) but also produces interference on feature 1, which the bias only partially suppresses.

Learn About topic

Contributors:

Deep-ML

Contribute