Notes on replicating J-lens on RWKVv7

Anthropic's J-lens was one of the rare tech in mech interp that somehow is scalable (compute once and reuse forever) and not labor intensive (doesn't try to manually understand one action at a time through raw sciencing). Computing J-lens was really expensive. So much so that creating one lens for layer 18 for RWKV v7 1.5B took me about 2.5 hours on my Radeon R9700 (the workstation version of Radeon 9070XT).

And... that utterly failed. The resulting lens produces utter garbage. But upon realizing how RWKV is constructed (vs Transformers that J-lens is developed on), I thought to myself, I am too dumb to reason with raw model internals.. what if I just apply Genetic Algorithm and make it search all possible coordinate transformations and find one that is simple yet reduces the post-lens entropy?

My computer spitted out:

tanh((layer_norm(resid_in) − resid_in)/2)

I genuinely do not know why, this expression is insane. Algebraically it is doing some scale control and then clipping, but offers little insight. Yet.. hell it does work wonders. tanh((layer_norm(resid_in) − resid_in)/2) @ J (where J is the computed J-lens) does produce coherent and topic related tokens. And this thing is hella stable, working across all new text I manage to throw at it (news, blog posts, conversations, LLM generated stories, etc..).

It is magic, but working magic is just undiscovered science.

Genetic Algorithm to recover the J-lens action

That got me thinking - if a trained J-lens is just a NxN matrix that transforms the (N-wide) residual.. could there be a way to recover part of the lens? Might be very unintuitive, but not like GA cares about something being intuitive or not.

Going into this, I already do not expect it to work well. If language modeling is that easy, we'd already have AGI a decade ago and doesn't need giant LLMs trained over the entirety of the internet. However it does serve a dual purpose - a. if it even works for a very little bit, the GAed expression would serve as a very easy target for mechanistic interpretability to reason about and look into, and b. if it doesn't work.. we know what to not look into in the future.

Of course, the first obvious thing to try is to attempt to make GA directly recover the J-lens matrix itself from model weights and activation statistics. Some quick data collection and experiments later, the computer spat out the expression:

I

That doesn't work. It is just the identity matrix. Obviously it ain't gonna be that easy. What if instead we aim to replicate the effect of the J-lens on the (rescaled and clamped) residuals across a wide corpus of text? That is, we ask GA to construct a matrix from model weights and activation statistics that, when applied to the modified residuals, produces the same output as the J-lens would. This way, the problem is much easier to solve and does not demand an exact solution.

And I realized that the matrix is likely a mix of an identity plus some other stuff. Thus the actual optimization target is to find a matrix that when v @ (alpha * I + beta * M).T mimics the effect of v @ J (where v is the modified residual, alpha and beta are scalars fitted to get close to J itself) plus some terms to discourage the GA from overfitting.

1000 documents, storing the residual at layer 18, limiting the GA to only look at Channel Mix weights/stats (to avoid all the time-related noise, I assume it makes the initial debugging easier) and some waiting later, I got:

alpha = 0.04201892
beta = 0.23095226
W_v @ (mean_relu_sq_derrivatives * W_k)

With a training set cosine similarity of 0.20858 and testing cosine at 0.18156 (both with the identity part of the matrices removed, else it dominates the results) -- it works but barely so. mean_relu_sq_derrivatives being the vector mean of the squared relu (part of RWKV's construction) derivatives, it itself is a vector of size N, and W_v and W_k are the channel mix weights.

At this point I doubt if this is a real result or some fluke that happened to work for the dataset I used. If it's truly something, shuffling mean_relu_sq_derrivatives should break any structure and collapse the results. Using W_v @ (shuffled(mean_relu_sq_derrivatives) * W_k) and refitting alpha and beta, the testing set cosine dropped to 0.10836.

At least it is not nothing....?

Widening the search space to include all weights at the layer (including time-mix now) and allowing the GA to search for more complex expressions, adding more tensors for statistics and updating the scoring be to cosine_similarity(sum(v, axis=0) @ estimated_matrix, sum(v, axis=0) @ J).mean() instead of cosine_similarity(v @ estimated_matrix, v @ J).mean() to weight better, I got:

 0.02920150 * (
     diag(mean(tmix.r))
     @ diag(mean(resid_out))
 )
 + 0.16724125 * (
     cmix_value_weight
     @ (mean_relu_sq_derivative[:, None] * cmix_key_weight)
     @ center
 )
 + 0.00081642 * diag(var(ln2))
 - 0.00055360 * diag(var(ln1))

Where LN1 and LN2 being the 2 layernorms in channel and time mix and center = I - ones(N, N) / N. That got us the new trace removed (testing set) cosine of 0.32292, Almost double the previous result! Hooray!

After a lot of trying to falsify and making sure this is not a dumb coincidence or GA discovering some trivial solution. I suddenly thought to myself - what if, what if I can push the performance even further adding constant terms that the GA can use to shift or scale the output? Right now the constants are lived outside the GA search space, being fitted after GA came up with the matrix (GA cannot insert constants into the matrix constructions). But what should I do? Add a new node for constants and let it mutate according to some standard deviation? That's a huge axis for GA to search over. I ended up with a funny idea. Instead, fit:

internal = (
    I
    + a * diag(mean_tensor_a)
    + b * diag(mean_tensor_b)
) @ diag(mean_tensor_c) @ diagmean_tensor_d)

fit a, b to the target J, and make multiple of them so GA can hierarchically discover the best combination with automatically fitted constants. This way the search space is much more constrained and won't take an entire day to search over. After not that long.. the computer says:

T = (
    - 0.00230497 * ln1_full_derivative
    - 0.23572776 * diag(var(cmix_out))
    + 0.16379395 * centered_cmix
    + 0.03869630 * diag(var(resid_out))
    + 0.00485161 * tmix_affine_atom
    + 0.00348761 * ln2_isotropic
)

where (the same expression as before)

centered_cmix = (
    value_weight
    @ (mean_relu_sq_derivative[:, None] * key_weight)
    @ center
)

and

tmix_affine_atom = normalize(
    (I
        + 0.8640476 * diag(mean(tmix_xa))
        - 9.7268698 * diag(mean(tmix_v_raw))

    ) @ diag(mean(tmix_xg))@ diag(var(tmix_r))
)

with training cosine of 0.32226 and testing 0.34316. Marginal but better. At this point I think complexity growing faster than the performance gain.

Falsifying the function

I still have doubt if what GA discovered is a real result or artifact from the dataset. The testing set results indicates it's fine. But being honest, I expected to get nothing honestly. Quickly, I grabbed several datasets and let GA fit some expression to replicate the effect of the J-lens on modified residuals (note that the original J was created on Pile10k dataset). The results are as follows:

Not optimal, but I pooled the diagonal (trained on Pile, test on Pile) cosine scores and plotted the histogram. It looks pretty normal and is not an outlier cluster pulling up the average. More importantly, there is no samples ever getting close or below 0. Which random projections should induce. The expressions have definitely learned something.

Obviously the dataset here is still biased - it's English text from the internet. But given the diversity I think it's safe to say that GA discovered real structure in the J-lens (trained on Pile too).

Until I realized that GA might be approximating the mean-field jacobian. Which I got GPT to analytically compute the mean-field jacobian of layer 18 and approximate the derrivatives from activation stats. That has a cosine of 0.76 on vector action and 0.6 on the measured J-lens according to Anthropic's method.

That is, GA was approximating the mean field jacobian, poorly. It does work but there's a much easier way to get the same result. Analytically.

I know, the story died suddenly and I hope I had something better to share. The real world being the real world and all. At least I learned something.