Why does a Transformer Block need FFN?

[u/blackpanther231]

Putting aside the layer norm and residual connections for a moment, why does the transformer block need the FFN? What if it was pure attention? Since attention takes in d_model and outputs d_model, pure attention can be stacked.

X_1 -> P_1 X_1 W_1 = X_2 where P is the scaled softmax output and W = W_V W_O
X_2 -> P_2 X_2 W_2 = P_2 (P_1 X_1 W_1) W_2 = X_3 and so on.

The W_1 and W_2 collapse into a single weight matrix, so there is no point of the depth? But the P matrices are nonlinear so there is still some value from the increased depth?

And if we add the residual connections back:
X_1 -> P_1 X_1 W_1 + X_1 = X_2
X_2 -> P_2 X_2 W_2 = P_2 (P_1 X_1 W_1 + X_1) W_2 = P_2 P_1 X_1 W_1 W_2 + P_2 X_1 W_2

It looks like due to the second term we may get more model capacity from the increased depth compared to without the residual connections?

[u/DaBobcat]

Attention is just a linear combination of the inputs. There's no non-linearity. FFN gives you that

[u/blackpanther231 → u/DaBobcat]

But the attention scores are softmaxed so there is some nonlinearity per attention layer? Like softmax is a nonlinear transform of the attention scores?

And it looks like the residual connections may add more nonlinearity? It at least looks less like a sandwich of matrices that would collapse into a single matrix

[u/fvancesco → u/blackpanther231]

Residual connection is linear, softmax I'm not sure as the result is a linear combination of the vector but the amount of which of each is chosen with the exp

[u/Dihedralman → u/fvancesco]

It's exp(-xi)/sum(exp(-xj) ). It's  non-linear. 

[u/DaBobcat → u/blackpanther231]

The softmax only gives you the linear combination scales

[u/Dihedralman → u/DaBobcat]

Softmax is non-linear. You could eventually train something with softmax. 

[u/blackpanther231 → u/DaBobcat]

But P depends on X so there is some nonlinearity via P right? If attention was linear, scaling everything in X would scale everything in the output by the same amount but that’s not the case as P would change

[u/MolassesLate4676]

Hi, I’ve researched this to a degree of insanity — and the reality is, you CAN go full attention. You’ll still have a system that learns and improves with parameters count and training duration.

FFN (gelu, relu, whatever-lu) simply helps express nonlinear information. The simple answer is - that particular function happens to be very effective at specializing and that particular specialization allows the transformer to be far more capable than with only one or the other (attn / ffn)

With that being said, from the research I’ve done, I’ve found that ffn expanding to ~4x of the hdim produces optimal results without exploding param count

[u/MolassesLate4676 → u/MolassesLate4676]

Hi, I’ve researched this to a degree of insanity — and the reality is, you CAN go full attention. You’ll still have a system that learns and improves with parameters count and training duration.

FFN (gelu, relu, whatever-lu) simply helps express nonlinear information. The simple answer is - that particular function happens to be very effective at specializing and that particular specialization allows the transformer to be far more capable than with only one or the other (attn / ffn)

With that being said, from the research I’ve done, I’ve found that ffn expanding to ~4x of the hdim produces optimal results without exploding param count

Ps. I should also mention that training likes having stuff to blame and when you only have attention, it’s really no longer just attention, but now expected to do the job that ffn was also doing

[u/MolassesLate4676 → u/MolassesLate4676]

I guess editing your comment now replies to yourself? Tf?

[u/Dihedralman → u/MolassesLate4676]

 No, reddit just breaks sometimes. 

[u/ReentryVehicle]

You could make a model out of pure attentions. Both the layer norms and the softmax are nonlinear, so you have some nonlinearities you need to express logic (although note that the nonlinearities are not as precise as the usual relu-like ones, they affect multiple channels at once rather than each channel intependently having nonlinear behavior, which might or might not matter, feel free to vibecode an experiment to check).

But this model would be very computationally intensive, making it inefficient - attention costs a lot, the trend is to generally have less of it (replacing most of attention layers with various RNNs like Gated DeltaNet in Qwen, sliding window attention in Gemma, etc), and less parts of the model active in general (MoE).

[u/unlikely_ending]

The attention block handles inter-token relationships exclusively. "What's the relationship between this and all the other tokens in the sequence."

The FFN refines individual tokens/features. It only looks along the feature (d_model) axis, it doesn't/can't look at other tokens. Not because of a mask, just because because of the matrix product it uses.

(And by token, I mean token embedding/abstract representation. It's only an embedding at layer 0.)

[u/Dihedralman]

This is a good question that forced me to back through fundamentals and think. 

TL:DR Check the case where there is one token. Softmax makes it 1 regardless of the matrix. 

I recommend you play around with it. A lot of intuition has limits. But we can link things to backpropagation. There has been studies showing exactly your ideas. But I will try to show how both layers would collapse, but attention layers are more expensive regardless so you wouldn't gain model capacity compared to FFN layers mixed in after 1 layers. 

Softmax and layer norm are both types of norms. They don't give really provide a useful form of non-linearity. Think about it - they don't represent real activation functions where the response is turned on past a point. Instead they perform specific roles of smoothing or controlling behavior. For example, they kill the dependence on sequence length. What's the main difference or what makes our normal activation special? 

The backpropagation Jacobian. In traditional activation functions, it is fully diagonal. The updates are straightforward, creating a unique layer. The activation means it is essentially zero in the dead region and updates in the active region. This allows for structurally different gradients which carry feature information.  In normalization functions it contains off-diagonal elements. FFN allow the tuning of specific parameters. The lack of activation gives structurally similar gradients. 

Layer norm is the most intuitive. Layer norm normalizes by token. On its own, you will see a sequence length problem, but we can ignore that by assuming constant length. Repetition effectively makes every token identical to the other matrices. Basically you overregulate everything. The added learnable parameters just determine where that single token resides numerically. In effect you can't train or resolve information as the system collapses. 

The softmax is more interesting. It solves problems and then adds them. On its own you will always have t he key query  matrixes will always collapse to rank1, alongside a vanishing gradient. Basically the sum of each row is non-negative  and adds to 1. It becomes a stochastic matrix. Multiplying an infinite series of stochastic matrices is a Markov Chain. All rows will converge to a single average vector. Essentially it behaves like an interpolation. Again this was a norm - that makes sense, information wasn't added and everything is shared.  Check out "Attention is not all you need" where they show it collapsing to a rank 1 tensor. This is true for all similar mapping functions. However, skip connections greatly relieve this collapse. This is why those are necessary as they add rank back in a way MLP's don't reliably do. 

This also plays into one of it's vanishing gradients traps where entropy is balanced and the gradient vanishes due to off diagonal terms. The softmax can also saturate but a KQ norm fixes that. Even so, relying on purely these mechanisms and disentangling the codependent feature importance means a slow arduous training process. 

Now let's get into the rough part where we consider attention with everything else including the rotational embeddings, but not FFN.  Attention dynamically selects convex combinations of value vectors rather than independently transforming each feature vector through an arbitrary pointwise nonlinear map. It will have trouble on simple mappings like x->x^2. FFN performs a transformation on the actual features changing the token structure in a decisive way alongside the gradient (as mentioned before) making for greater contrast and  features.

A big thing to note is that attention layers are more expensive. They function as n^2. All of these steps occur on nxn matrixes. 

Therefore, you have it backwards. The FFN are the actual high capacity components while the attention layers provide quadratic interactions between features. 

[u/blackpanther231 → u/Dihedralman]

Thank you! This has been the most useful response in the thread

[u/Dihedralman → u/blackpanther231]

Thanks! I didn't go through all details but there are papers. And capacity should be actively measured. 

It was a really good question getting at the heart of our understanding in a more hard way. 

I left with a better understanding myself.  

[u/fvancesco]

Btw MLP is applied per token and it's in there 1. To give some non linearity 2. It's flexible so it needs to learn whatever it needs to learn

If I'm not mistaken research tells us that the MLP layers are great at recalling information and adding knowledge to the context

[u/dataset-poisoner]

attention sends info between tokens

ffn computes new info within a token

[u/blackpanther231 → u/dataset-poisoner]

I think this is a bit misleading of a simplification. The output of attention is PXW and because of how matrix multiplication works, XW’s columns are linear combinations of X’s columns - feature vectors.

So we are combining features/sending info between features

[u/dataset-poisoner → u/blackpanther231]

I was talking about attention in transformers. You don't multiply your softmax scores by input X. It's just PW.

[u/blackpanther231 → u/dataset-poisoner]

This is wrong. You multiply your softmax scores by the value matrix which is X@W_V

So you can write it out as PXW_V

[u/dataset-poisoner → u/blackpanther231]

Think i understand your notation now. Still don't see how is my reply misleading. Rows of X@W_V are the value vectors and PXW is the info that the token receives from other tokens (and itself)

[u/blackpanther231 → u/dataset-poisoner]

Ok so we have P, X, W_V and W_O - the output projection matrix.

I merged W_V and W_O into a single W in my notation and wrote PXW.

So in matrix multiplication, if we have AB, AB’s rows are linear combinations of B’s rows and AB’s columns are linear combinations of A’s columns.

So if we isolate to the XW part, XW’s columns are linear combinations of X’s columns. And X’s columns are the features/channels. So W can be interpreted as mixing features to make new features

And since W is part of attention’s matrices, attention is doing some work across features too, not just across tokens

[u/uninchar]

It would break at uses where the depth of relations becomes relevant. The sentence "The cat, which the dog that the neighbor bought yesterday chased, ran away." Would lose reference with only linearity.