[Edit on 17th Mar] After conducting experiments on more data points (5000 texts) on the Pile dataset (more sample sources), we are confident that the experimental results described earlier are reliable. Therefore, we have opened the code.

Recently, we conducted several experiments focused on the evolution and migration of token embeddings across different layers during the forward processing of LLMs. Specifically, our research targeted open-source, decoder-only architecture models, such as GPT-2 and Llama-2^[1].

Our experiments are initiated from the perspective of an older research topic known as the logit lens. Utilizing the unembedding matrix on embeddings from various layers is an innovative yet thoughtless approach. Despite yielding several intriguing observations through the logit lens, this method lacks a solid theoretical foundation, rendering subsequent research built upon it potentially questionable. Moreover, we observed that some published studies have adopted a similar method without recognizing it in community^[2]. A primary point of contention is that the unembedding matrix is tailored to the embeddings of the final layer. If we view the unembedding matrix as a solver, its solvable question space is inherently designed for embeddings from the last layer. Applying it to embeddings from other layers introduces a generalization challenge. The absence of a definitive ground truth for this extrapolation means we cannot conclusively determine the validity of the decoder's outputs.

**Fig 1. Hypothesis**: Embeddings output from different layers, along with input token embeddings, each occupy distinct subspaces. The unembedding matrix (or decoder) is only effective within the subspace where it was trained. While applying it to other layers' embeddings is mathematically feasible, the results may not always make sense.

Next, we'll introduce several experiments about this hypothesis.

Do embeddings form different spaces?

It's crucial to acknowledge that the question of reasonableness may be inherently unanswerable due to the absence of a definitive ground truth. Our hypothesis is merely a side thought inspired by this question. Therefore, related studies such as Tuned Lens^[3], which seek to address this question of reasonableness, fall outside the scope of our discussion.

Intuitively, to identify whether there exist subspaces, we can do some visualization about e...

Posts

Wikitag Contributions

Comments