Finally gonna start properly experimenting on stuff. Just writing up what I’m doing to force myself to do something, not claiming this is necessarily particularly important.
Llama (and many other models, but I’m doing experiments on Llama) has a piece of code that looks like this:
h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) out = h + self.feed_forward(self.ffn_norm(h))
Here, out is the result of the transformer layer (aka the residual stream), and the vectors self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) and self.feed_forward(self.ffn_norm(h)) are basically where all the computation happens. So basically the transformer proceeds as a series of “writes” to the residual stream using these two vectors.
I took all the residual vectors for some queries to Llama-8b and stacked them into a big matrix M with 4096 columns (the internal hidden dimensionality of the model). Then using SVD, I can express M=∑isi(ui⊗vi), where the u‘s and v’s are independent units vectors. This basically decomposes the “writes” into some independent locations in the residual stream (u’s), some latent directions that are written to (v’s) and the strength of those writes (s’s, aka the singular values).
To get a feel for the complexity of the writes, I then plotted the s’s in descending order. For the prompt “I believe the meaning of life is”, Llama generated the continuation “to be happy. It is a simple concept, but it is very difficult to achieve. The only way to achieve it is to follow your heart. If you follow your heart, you will find happiness. If you don’t follow your heart, you will never find happiness. I believe that the meaning of life is to”. During this continuation, there were 2272 writes to the residual stream, and the singular values for these writes were as follows:
The first diagram shows that there were 2 directions that were much larger than all the others. The second diagram shows that most of the singular values are nonnegligible, which indicates to me that almost all of the writes transfer nontrivial information. This can also be seen in the last diagram, where the cumulative size of the singular values increases approximately logaritmically with their count.
This is kind of unfortunate, because if almost all of the s was concentrated in a relatively small number of dimensions (e.g. 100), then we could simplify the network a lot by projecting down to these dimensions. Still, this was relatively expected because others had found the singular values of the neural networks to be very complex.
Since variance explained is likely nonlinearly related to quality, my next step will likely be to clip the writes to the first k singular vectors and see how that impacts the performance of the network.
Ok, so I’ve got the clipping working. First, some uninterpretable diagrams:
In the bottom six diagrams, I try taking varying number (x-axis) of right singular vectors (v’s) and projecting down the “writes” to the residual stream to the space spanned by those vectors.
The obvious criterion to care about is whether the projected network reproduces the outputs of the original network, which here I operationalize based on the log probability the projected network gives to the continuation of the prompt (shown in the “generation probability” diagrams). This appears to be fairly chaotic (and low) in the 1-300ish range, and then stabilizes while still being pretty low in the 300ish-1500ish range, and then finally converges to normal in the 1500ish to 2000ish range, and is ~perfect afterwards.
The remaining diagrams show something about how/why we have this pattern. “orig_delta” concerns the magnitude of the attempted writes for a given projection (which is not constant because projecting in earlier layers will change the writes by later layers), and “kept_delta” concerns the remaining magnitude after the discarded dimensions have been projected away.
In the low end, “kept_delta” is small (and even “orig_delta” is a bit smaller than it ends up being at the high end), indicating that the network fails to reproduce the probabilities because the projection is so aggressive that it simply suppresses the network too much.
Then in the middle range, “orig_delta” and “kept_delta” explodes, indicating that the network has some internal runaway dynamics which normally would be suppressed, but where the suppression system is broken by the projection.
Finally, in the high range, we get a sudden improvement in loss, and a sudden drop in residual stream “write” size, indicating that it has managed to suppress this runaway stuff and now it works fine.
An implicit assumption I’m making when I clip off from the end with the smallest singular values is that the importance of a dimension is proportional to its singular values. This seemed intuitively sensible to me (“bigger = more important”), but I thought I should test it, so I tried clipping off only one dimension at a time, and plotting how that affected the probabilities:
Clearly there is a correlation, but also clearly there’s some deviations from that correlation. Not sure whether I should try to exploit these deviations in order to do further dimension reduction. It’s tempting, but it also feels like it starts entering sketchy territories, e.g. overfitting and arbitrary basis picking. Probably gonna do it just to check what happens, but am on the lookout for something more principled.
Back to clipping away an entire range, rather than a single dimension. Here’s ordering it by the importance computed by clipping away a single dimension:
Less chaotic maybe, but also much slower at reaching a reasonable performance, so I tried a compromise ordering that takes both size and performance into account:
Doesn’t seem like it works super great tbh.
Edit: for completeness’ sake, here’s the initial graph with log-surprise-based plotting.
To quickly find the subspace that the model is using, I can use a binary search to find the number of singular vectors needed before the probability when clipping exceeds the probability when not clipping.
A relevant followup is what happens to other samples in response to the prompt when clipping. When I extrapolate “I believe the meaning of life is” using the 1886-dimensional subspace from
[I believe the meaning of life is] to be happy. It is a simple concept, but it is very difficult to achieve. The only way to achieve it is to follow your heart. It is the only way to live a happy life. It is the only way to be happy. It is the only way to be happy. The meaning of life is
, I get:
[I believe the meaning of life is] to find happy. We is the meaning of life. to find a happy. And to live a happy and. If to be a a happy. . to be happy. . to be happy. . to be a happy.. to be happy. . to be happy.
Which seems sort of vaguely related, but idk.
Another test is just generating without any prompt, in which case these vectors give me:
Question is a single thing to find. to be in the best to be happy. I is the only way to be happy. I is the only way to be happy. I is the only way to be happy. It is the only way to be happy.. to be happy.. to be happy. to
Using a different prompt:
[Simply put, the theory of relativity states that ]1) the laws of physics are the same for all non-accelerating observers, and 2) the speed of light in a vacuum is the same for all observers, regardless of their relative motion or of the motion of the source of the light. Special relativity is a theory of the structure of spacetime
I can get a 3329-dimensional subspace which generates:
[Simply put, the theory of relativity states that ] 1) time is relative and 2) the speed of light in a vacuum is constant for all observers. 1) Time is relative, meaning that if two observers are moving relative to each other, the speed of light is the same for all observers, regardless of their motion. For example, if you are moving relative
or
Question: In a simple harmonic motion, the speed of an object is A) constant B) constant C) constant D) constant In the physics of simple harmonic motion, the speed of an object is constant. The speed of the object can be constant, but the speed of an object can be
Another example:
[A brief message congratulating the team on the launch:
Hi everyone,
I just ] wanted to congratulate you all on the launch. I hope that the launch went well. I know that it was a bit of a challenge, but I think that you all did a great job. I am proud to be a part of the team.
Thank you for your
can yield 2696 dimensions with
[A brief message congratulating the team on the launch:
Hi everyone,
I just ] wanted to say you for the launch of the launch of the team.
The launch was successful and I am so happy to be a part of the team and I am sure you are all doing a great job.
I am very looking to be a part of the team.
Thank you all for your hard work,
or
def measure and is the definition of the new, but the the is a great, but the The is the The is a The is a The is a The The is a The The The is a The The is a
And finally,
[Translate English to French:
sea otter ⇒ loutre de mer peppermint ⇒ menthe poivrée plush girafe ⇒ girafe peluche cheese =>] fromage pink ⇒ rose blue ⇒ bleu red ⇒ rouge yellow ⇒ jaune purple ⇒ violet brown ⇒ brun green ⇒ vert orange ⇒ orange black ⇒ noir white ⇒ blanc gold ⇒ or silver ⇒ argent
can yield the 2518-dimensional subspace:
[Translate English to French:
sea otter ⇒ loutre de mer peppermint ⇒ menthe poivrée plush girafe ⇒ girafe peluche cheese =>] fromage cheese ⇒ fromage cheese ⇒ fromage f cheese ⇒ fromage butter ⇒ fromage apple ⇒ orange yellow ⇒ orange green ⇒ vert black ⇒ noir blue ⇒ ble purple ⇒ violet white ⇒ blanc
or
Question: A 201 The sum of a The following the sum the time the sum the the the The The The The The The The The The The The The The The The The The The
Given the large number of dimensions that are kept in each case, there must be considerable overlap in which dimensions they make use of. But how much?
I concatenated the dimensions found in each of the prompts, and performed an SVD of it. It yielded this plot:
… unfortunately this seems close to the worst-case scenario. I had hoped for some split between general and task-specific dimensions, yet this seems like an extremely uniform mixture.
If I look at the pairwise overlap between the dimensions needed for each generation:
… then this is predictable down to ~1% error simply by assuming that they pick a random subset of the dimensions for each, so their overlap is proportional to each of their individual sizes.
Oops, my code had a bug so only self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) and not self.feed_forward(self.ffn_norm(h)) was in the SVD. So the diagram isn’t 100% accurate.
Finally gonna start properly experimenting on stuff. Just writing up what I’m doing to force myself to do something, not claiming this is necessarily particularly important.
Llama (and many other models, but I’m doing experiments on Llama) has a piece of code that looks like this:
h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
out = h + self.feed_forward(self.ffn_norm(h))
Here, out is the result of the transformer layer (aka the residual stream), and the vectors self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) and self.feed_forward(self.ffn_norm(h)) are basically where all the computation happens. So basically the transformer proceeds as a series of “writes” to the residual stream using these two vectors.
I took all the residual vectors for some queries to Llama-8b and stacked them into a big matrix M with 4096 columns (the internal hidden dimensionality of the model). Then using SVD, I can express M=∑isi(ui⊗vi), where the u‘s and v’s are independent units vectors. This basically decomposes the “writes” into some independent locations in the residual stream (u’s), some latent directions that are written to (v’s) and the strength of those writes (s’s, aka the singular values).
To get a feel for the complexity of the writes, I then plotted the s’s in descending order. For the prompt “I believe the meaning of life is”, Llama generated the continuation “to be happy. It is a simple concept, but it is very difficult to achieve. The only way to achieve it is to follow your heart. If you follow your heart, you will find happiness. If you don’t follow your heart, you will never find happiness. I believe that the meaning of life is to”. During this continuation, there were 2272 writes to the residual stream, and the singular values for these writes were as follows:
The first diagram shows that there were 2 directions that were much larger than all the others. The second diagram shows that most of the singular values are nonnegligible, which indicates to me that almost all of the writes transfer nontrivial information. This can also be seen in the last diagram, where the cumulative size of the singular values increases approximately logaritmically with their count.
This is kind of unfortunate, because if almost all of the s was concentrated in a relatively small number of dimensions (e.g. 100), then we could simplify the network a lot by projecting down to these dimensions. Still, this was relatively expected because others had found the singular values of the neural networks to be very complex.
Since variance explained is likely nonlinearly related to quality, my next step will likely be to clip the writes to the first k singular vectors and see how that impacts the performance of the network.
Ok, so I’ve got the clipping working. First, some uninterpretable diagrams:
In the bottom six diagrams, I try taking varying number (x-axis) of right singular vectors (v’s) and projecting down the “writes” to the residual stream to the space spanned by those vectors.
The obvious criterion to care about is whether the projected network reproduces the outputs of the original network, which here I operationalize based on the log probability the projected network gives to the continuation of the prompt (shown in the “generation probability” diagrams). This appears to be fairly chaotic (and low) in the 1-300ish range, and then stabilizes while still being pretty low in the 300ish-1500ish range, and then finally converges to normal in the 1500ish to 2000ish range, and is ~perfect afterwards.
The remaining diagrams show something about how/why we have this pattern. “orig_delta” concerns the magnitude of the attempted writes for a given projection (which is not constant because projecting in earlier layers will change the writes by later layers), and “kept_delta” concerns the remaining magnitude after the discarded dimensions have been projected away.
In the low end, “kept_delta” is small (and even “orig_delta” is a bit smaller than it ends up being at the high end), indicating that the network fails to reproduce the probabilities because the projection is so aggressive that it simply suppresses the network too much.
Then in the middle range, “orig_delta” and “kept_delta” explodes, indicating that the network has some internal runaway dynamics which normally would be suppressed, but where the suppression system is broken by the projection.
Finally, in the high range, we get a sudden improvement in loss, and a sudden drop in residual stream “write” size, indicating that it has managed to suppress this runaway stuff and now it works fine.
An implicit assumption I’m making when I clip off from the end with the smallest singular values is that the importance of a dimension is proportional to its singular values. This seemed intuitively sensible to me (“bigger = more important”), but I thought I should test it, so I tried clipping off only one dimension at a time, and plotting how that affected the probabilities:
Clearly there is a correlation, but also clearly there’s some deviations from that correlation. Not sure whether I should try to exploit these deviations in order to do further dimension reduction. It’s tempting, but it also feels like it starts entering sketchy territories, e.g. overfitting and arbitrary basis picking. Probably gonna do it just to check what happens, but am on the lookout for something more principled.
Back to clipping away an entire range, rather than a single dimension. Here’s ordering it by the importance computed by clipping away a single dimension:
Less chaotic maybe, but also much slower at reaching a reasonable performance, so I tried a compromise ordering that takes both size and performance into account:
Doesn’t seem like it works super great tbh.
Edit: for completeness’ sake, here’s the initial graph with log-surprise-based plotting.
To quickly find the subspace that the model is using, I can use a binary search to find the number of singular vectors needed before the probability when clipping exceeds the probability when not clipping.
A relevant followup is what happens to other samples in response to the prompt when clipping. When I extrapolate “I believe the meaning of life is” using the 1886-dimensional subspace from
, I get:
Which seems sort of vaguely related, but idk.
Another test is just generating without any prompt, in which case these vectors give me:
Using a different prompt:
I can get a 3329-dimensional subspace which generates:
or
Another example:
can yield 2696 dimensions with
or
And finally,
can yield the 2518-dimensional subspace:
or
Given the large number of dimensions that are kept in each case, there must be considerable overlap in which dimensions they make use of. But how much?
I concatenated the dimensions found in each of the prompts, and performed an SVD of it. It yielded this plot:
… unfortunately this seems close to the worst-case scenario. I had hoped for some split between general and task-specific dimensions, yet this seems like an extremely uniform mixture.
If I look at the pairwise overlap between the dimensions needed for each generation:
… then this is predictable down to ~1% error simply by assuming that they pick a random subset of the dimensions for each, so their overlap is proportional to each of their individual sizes.
Oops, my code had a bug so only self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) and not self.feed_forward(self.ffn_norm(h)) was in the SVD. So the diagram isn’t 100% accurate.