Attention vs MLP — who does what

Background

A transformer block has two sub-layers with complementary jobs:

• Attention mixes information across token positions — each output is a weighted sum of other tokens. (cross-token)
• MLP applies a non-linear transformation within each token, the same weights at every position, no communication between tokens. (per-token)

The block is "route, then transform", repeated n_layer times.

Problem statement

Implement sublayer_role(sublayer) returning whether it mixes across tokens or transforms per token.

Input

• sublayer — "attention" (or "attn"), or "mlp" (or "ffn").

Output

Returns "cross-token" for attention, "per-token" for the MLP.

Examples

Example 1

Input:  sublayer = "attention"
Output: "cross-token"

Example 2

Input:  sublayer = "mlp"
Output: "per-token"

Constraints

• "attention"/"attn" → "cross-token".
• "mlp"/"ffn" → "per-token".

Notes

• This split is why both are needed: attention alone can route but not richly transform; the MLP alone can transform but not move information between positions.