▲ 153 points
back
43 comments
It's interesting that it's the middle layers of the Transformer that are affected most by RL post-training, but it perhaps makes some intuitive sense given that RL is being used to shape high level planning-type direction of the output.
It seems that the input layers to a Transformer are necessarily going to be doing the most low level work of syntax -> semantic augmentation starting with things like tagging parts of speech etc. Similarly the output layers are by necessity going to be concerned with mapping high level representations back into surface level word sequence form. This leaves the middle layers to do the work of first recognizing deep enough patterns to support good quality prediction, then do the high level predication itself which is what RL is typically going to be trying to shape.
> the middle layers of the Transformer that are affected most by RL post-training
this is explicitly why qwen opted to use recurrent layers as the middle layers of their hybrid model.
AFAIK Alibaba's Qwen3-Next architecture doesn't do anything different specifically for the middle layers - it uses a 3:1 mix of recurrent and regular attention blocks throughout the full depth of the transformer.
The motivation for this mostly recurrent hybrid attention is to efficiently support long context lengths.
This result feels very intuitive. The early layers of a transformer can be thought of as understanding surface level things like syntax, how tokens group, which groups are entities and how to disambiguate them, etc. The last layers are in a sense decoding ideas into a selection of words, ensuring the grammar makes sense, that the text flows and is structured correctly, etc. The middle layers are where the abstract thought and manipulation of concepts is happening.
But for the tasks this paper uses for RL training, it's all about improving the way the net is manipulating concepts. So the middle layers are where the focus should be.
Note: RL is also used for tasks that aren't about conceptual manipulation, like instruct training. I bet that their result doesn't hold for that because the delta vs the foundation model is all about the selection of words and flow of the text, not the core understanding.
I keep thinking of the RYS (Repeat Yourself) experiment of simply looping some of the inner layers of LLMs for better results and wonder if any progress was made on it.
https://dnhkng.github.io/posts/rys/
Feels it should be straightforward to integrate in LLMs a network to control the looping. Or just duplicate entire blocks of layers after the initial training.
Yes, computing in latent space is a big thing now.
Is that in effect doing the same as “chain of thought” or “think through your steps” aka “Reasoning Models”?
It seems related. The primary benefit of reasoning tokens seems to be just giving the model more chances to loop through all its layers. There were experiments at one point with "pause tokens" that were meaningless but seemed to improve performance. Reasoning is a lot more interpretable though.
there certainly are experiments in keeping the reasoning in latent space.
i dont think this is quite the same though, since you arent picking tokens for the chain of thought. inatead, its staying on trying to pick the immediate next token.
as an alternative, maybe you could stack these to produce most likely token lists instead by stacking these?
but i think youd end up with the similar blurriness that llm video generators get where theyre returning an average of all the likely combinations rather than collapsing that wave function
I've often thought about how rich a machine language could be that communicates machine to machine on an interface that is really really close to those middle layers. I imagine a standardize meta interface that each model 'grows' a connection to with RL.
This idea is called neuralese and the labs are reluctant to do it because of the interpretability and control issues it poses.
The authors have some inconsistencies with training token length…
Most errors are probably responses that didn’t finish before their 3K token limit. They’ve measured how well RL is able to shorten the response to their limit.
If you think about it for some time then you’ll come to realise transformers are autoencoders on steroids. A small input space is expanded onto a big manifold and contracted again.
Now, suppose you want to impose a function to regulate the output of an autoencoder. It’s actually pretty obvious that you need exactly one layer to do so… f(manifold).
Everything can be represented as f(), a full scale SotA transformer model is also just f(context). That does not mean one layer is sufficient. It all depends on the level of expressivity required by this f to be a good model.
What you're suggesting seems to go implausibly far beyond what the paper says.
RL post-training alters the parameters of the transformer, while your f(manifold) idea seems to suggest that a new layer on top would suffice, no need to alter the transformer itself at all.
It would be extremely handy if that were so, but I'm guessing it isn't, or it would be the prevailing approach.
The manifold is in the middle (“small input space is expanded onto a big manifold and contracted again”) so f(manifold) would need to be in the middle too.
The problem really is one has to find out what constitutes the embedding and output if you want to apply f(manifold) but as others have pointed out that’s out of scope for this paper. My only insight here is that it is not surprising at all that one layer, or even a single function pass, suffices to get the desired reinforcement. I don’t actually know if there’s any research that quantifies the “manifold threshold” of each layer to try out this approach but it would be interesting for sure
I might be misunderstanding your point but this conflates the distinguishing features of each. you mention expansion but autoencoders canonically compress their inputs. autoencoders have an explicit encoder and decoder. most transformers we interact with these days (LLMs) are decoder only. the manifold isn't typically something the model is applied to directly. we apply the function/model to the latent representations. those are what live on the manifold.
Now that’s interesting.. what exactly distinguishes latent representations and the manifold? IMHO, those are the same, and you’re constructing a piecewise function of the manifold itself. Decoders also produce manifolds much in the same way, with the distinction being that the encoder isn’t learned but static after initialisation. So fundamentally it is still DOING the same operation.
The latent representations of the data are like points on a surface. That surface is the manifold. We don't typically have the full manifold and can only sample points from it by embedding data into it.
Worth noting a different manifold "exists" after each transformation (e.g. layer). You only sample from the same manifold when you apply the same transformation(s).
Also worth noting that in reality manifolds will be "spiky" in very high dimension, so the idea of a "surface" is best understood through patterns of distance between samples in embedding space and way they collapse in low D.
Took me a short time to understand what you mean with "autoencoders on steroids", but I believe you mean they are autoencoders with an inverse bottleneck - an intermediate representation that isn't smaller, but that's much larger than the input space. Is my understanding of your comment correct?
Kind of. Autoencoders don’t need to have an embedding that’s smaller than the input. Their only requirement is that they compress information and thus create reconstruction loss. Typically however they are not trained this way because they don’t converge.. transformers do the same thing, but they can squeeze much more bits of information through one pass because the way they are designed. This holds true even for decoder only networks because they’re still doing the same thing
If the embedding isn’t smaller than the input, how is it compressing information? It might lose information in its mapping to the embedding space, but in my understanding, the definition of compression means it has to use less bits than the original to hold the same information. As such, the embedding space must be smaller.
You’re thinking information content in the neurons but not distributions. D(world of text) is much much much larger than D(transformer). Just because you’re passing information through a smaller channel doesn’t mean the original distribution necessarily isn’t smaller. Tokenisers are the numbers but then you have V^n where V is your vocabulary and n is the sequence length as a subset of the whole distribution of possible text. Now you pass that through something with sufficient capacity, which compresses the input distribution (you loose bits), and then recreate from your lossy distribution. AutoEncoders are not trained this way because they simply don’t converge if you try to pass smaller input to produce larger embeddings, but they’re doing the same thing.
Another way is to think of it is that LLM layers act like binary classifiers via soft max for the input distribution (there were a slew of papers comparing LLMs to SVMs). Essentially during your training you decide what part of the input distribution is going to be part of the embedding.
In my humble opinion, while the training diverges, these models produce the same manifold induced by the original distribution.
This is interesting theoretically, but in practical terms it’s hard to apply.
RL is already hard. There are many things which can go wrong. You have all of the problems with regular LLM SFT, plus now you have a reward model which can be hacked or too hard. Or KL collapse because the outputs are repetitive. Or maybe your groups in GRPO aren’t producing advantages. Or the rollouts are OOD for your reward model. Or maybe you’re running the rollout at a different precision as the trained weights. Or maybe your importance sampling should be clipping when it’s not, or should be clipping at the token level rather than sequence level.
Maybe after reading the above you think that the above are not problems because smart people wouldn’t make those mistakes. Fair enough. But I would prefer RL people like myself who are not geniuses.
Now, this is adding another variable into the mix: choosing a single layer to train. If it doesn’t work is it because there’s a problem with your RL setup? Or did you just choose the wrong layer? Or maybe there’s no problem with your setup but you chose a suboptimal layer to train.
Also note that we already have LoRA, which is a more established method for low memory parameter updates.
I guess the practical implication is probably just train on middle layers or increase the corresponding learning rate
Good work! I wonder if meta-learning can play a better role here compared to heuristics or hindsight. MAML requires hessians, but first-order MAML or Reptile variants could help apply layer-wise adjustments to learning rates.
I'm reminded of this dude who was sitting at or near the top of some kaggle leaderboard simply[0] splicing together some duplicated middle layers and applying a bit of fine tuning
[0] not simply
Was it this guy? https://news.ycombinator.com/item?id=47322887
May have been, yes!
His articles blow me away. Like the one where he plumbed a pair of Grace-Hopper superchips with liquid cooling for a home rig.
I'm wondering if the big problem is just the lack of recurrent connections in the standard Transformer design, and selective layer duplication is just a weird way to fix the same problem. I have to wonder if it would be possible to deliberately architecture a model to discover and exploit layers worth duplicating at training time.
The current model architectures we use have a fixed routing of residuals per layer, from the first to the last. I'm imagining replacing this with a matrix of routing weights[0] that determines how "strong" the connection is between each Transformer layer. We still evaluate each layer "in order", but now instead of just giving the layer the last layer's residuals, it gets the sum of all prior layers times their weight in the routing matrix. Recurrent connections (i.e. output of layer 9 to input of layer 3) could be handled by doing a second pass and using the first pass's recurrent residuals as inputs. You could then "loop" the model as many times as desired per token, or even have it do parallel decoding with each token communicating with the others while also recurring on itself.
You'd probably need some kind of normalization akin to what Deepseek did with Manifold Hyper Connections (mHC). Hell, mHC might also be useful in combination with this kind of layer routing, so the model could grow different recurrent loops for various bits of it's thought-space.
EDIT: if anyone uses it please call it "neuralese recurrence" just to scare the AI safety bros
[0] I'm not sure how you'd initialize these weights. Maybe each row/column is a narrow gaussian centered around the prior layer, with some random or constant weighting everywhere else?
Moe style router?
No. Mixture-of-Experts is for dynamically loading different sets of weights[0] per-token. The idea with MoE is that you can have more weights than fit in RAM without having to actually load everything in RAM.
My idea is more about creating recurrent paths in an otherwise all-forward network. The same number of weights would be loaded, we'd just be routing residuals differently.
[0] Usually multi-layer perceptron / linear portion weights - although maybe someone's tried attention head MOE?
Makes sense - This is very similar to fine tuning a down stream task in encoder-decoder architecture (~Bert style)
Hi I am the author. I am very happy to discuss with you about my paper. Feel free to contact me via email or just add comments here
If most of the performance gains are hidden in a few middle layers, you can save a massive amount of compute by freezing the rest
Really good work here, bravo
I still can't believe that LLM encoders aren't unsupervised learned.
So much left on the table
They are using Qwen, so this is decoder only.
Yes, my comment was kind of a non sequitur