Amazon
Amazon Reviews is the Multilingual Amazon Reviews Corpus. The English split contains 200,000 training and 5,000 test product reviews, each rated from 1 to 5 stars. Reviews combine a title and a body; the 5-class ordinal nature of the task makes it harder than binary sentiment classification. It serves here as a first benchmark on a dataset with few classes before tackling the larger-class-count CLINC150 and NAF settings.
1 Hyperparameter sweep
This section explores the results across five complementary angles: model capacity and robustness; the effect of vocabulary size and embedding dimension; the effect of the number of attention heads; training time; and the interaction between embedding dimension and vocabulary size.
All these angles are read off a single grid search covering 960 configurations in total. n_layers is the number of transformer encoder blocks and n_head the number of attention heads per block. For FastText (n_layers=0), n_head is not applicable; without label attention this yields \(4 \times 4 \times 3 = 48\) runs, and with label attention another \(4 \times 4 \times 3 = 48\) runs, for a FastText subtotal of 96. For transformer architectures (n_layers ∈ {1, 2, 4}), the full grid gives \(4 \times 3 \times 3 \times 2 \times 4 \times 3 = 864\) runs.
| Hyperparameter | Values | Applies to |
|---|---|---|
| Embedding dimension | 64, 128, 256, 512 | All |
Number of layers (n_layers) |
0, 1, 2, 4 | All |
| Vocabulary size | 2,000 / 5,000 / 10,000 / 20,000 | All |
| Learning rate | 0.0001, 0.0005, 0.001 | All |
Attention heads (n_head) |
2, 4, 8 | n_layers ≥ 1 |
Label attention (n_heads_label_attention) |
none, 4 | All |
1.1 Capacity & robustness
This section examines two complementary dimensions of the hyperparameter sweep. The first tab plots validation loss against model size, revealing the performance frontier each architecture can reach for a given parameter budget. The second tab looks at the spread of validation loss across all runs for a fixed architecture, measuring how sensitive each model is to the choice of hyperparameters.
A key observation is that n_head and lr do not affect num_params. In a standard transformer, the total size of the attention matrices is \(4 \times d^2\) regardless of how many heads are used, since \(d_{\text{head}} = d / n_{\text{head}}\). Similarly, lr is a training hyperparameter with no effect on architecture. As a result, for any fixed combination of (emb_dim, n_layers, vocab_size), FATE architectures have exactly \(3 \times 3 = 9\) runs sharing the same num_params (3 values of n_head × 3 values of lr), while FastText has only \(3\) runs (3 values of lr, since n_head does not apply). This motivates two complementary views:
- Scatter plot: all 960 runs, each point being that run’s
min_val_loss(the minimum over its training epochs), showing the full spread of achievable val_loss for a given parameter budget. - Curve plot: for each
(architecture, emb_dim, vocab_size)group, take themin(val_loss)overn_headandlr, collapsing the 9 (FATE) or 3 (FastText) runs into a single best-case point, tracing the performance frontier as model capacity grows.
Both FastText variants plateau regardless of model size — FastText around 0.93–0.94 and FastText + Label Attention slightly lower at ~0.92 — confirming that the mean-pooling bottleneck sets a hard performance ceiling that extra parameters cannot overcome. FATE architectures improve steadily with capacity: FATE (L=4) reaches the lowest validation loss (~0.845) and dominates across all parameter budgets, followed by FATE (L=2) and then FATE (L=1) at ~0.88–0.90. Adding label attention to FATE does not systematically change the frontier, suggesting that the transformer blocks already capture most of the relevant context.
The violin plot shows the distribution of validation loss across all runs for each architecture. A narrow violin indicates that the architecture is robust to hyperparameter choice, while a wide one signals high sensitivity.
The two FastText variants are by far the most stable: their violins are narrow and tightly concentrated (FastText ~0.94, FastText + Label Attention ~0.93), meaning hyperparameter choice has almost no impact on the outcome. FATE (L=1) shows a moderate spread centered around 0.90. FATE (L=2) exhibits wider distributions (~0.86–0.92). FATE (L=4) achieves the best minimum loss (~0.845) but also shows a spread, with the plain (non-label-attention) variant reaching the lowest tails. The trade-off is clear: the most capable architectures achieve the best performance but require more careful tuning.
1.2 Effect of vocabulary size and embedding dimension
The previous analysis treats model capacity (number of parameters) as a single axis, but two hyperparameters drive that capacity in very different ways: vocabulary size controls the size of the embedding look-up table - the dominant parameter cost for all architectures - while embedding dimension determines the width of every layer. Varying them independently allows us to disentangle their respective contributions to performance and identify practical sweet spots.
The main gains occur between 2,000 and 5,000 tokens; beyond that, all curves flatten. FATE architectures show a noticeable drop from 2k to 5k, after which performance stabilises: FATE (L=1) variants around 0.895–0.900, FATE (L=2) around 0.88, and FATE (L=4) around 0.873. This suggests that a vocabulary of 5,000 tokens already captures most of the useful subword coverage for this dataset.
FastText variants are insensitive throughout, their curves remaining essentially flat from 64 to 512. FATE (L=1) architectures improve from 64 to 128–256 then plateau. FATE (L=2) and FATE (L=4) show a similar pattern with validation loss decreasing up to 256 and slightly increasing at 512, making 256 the clear sweet spot for deeper transformer architectures. Larger embeddings add parameters without improving and can marginally hurt the most expressive models.
1.3 Effect of number of attention heads
The number of attention heads has virtually no effect on validation loss across all six FATE variants. All mean curves are essentially flat from 2 to 8 heads, and the ranking between architectures — FATE (L=4) best, then L=2, then L=1 — is fully preserved regardless of head count. This is consistent with the theoretical expectation: the total attention matrix size is fixed at \(4 \times d^2\), so increasing the number of heads simply redistributes the same capacity into smaller per-head subspaces without adding parameters.
1.4 Training time
The graph shows the typical training time ranges per architecture, but no clear correlation between training time and validation loss within a given family. FastText variants cluster on the left (under 2,000 s) at high loss (~0.93–0.96); FATE (L=1) variants are scattered across 500–2,000 s with losses in the 0.88–0.92 range; FATE (L=2) architectures span ~500–3,500 s and reach ~0.86–0.90; FATE (L=4) is the most expensive (1,500–60,000 s) and achieves the lowest losses (~0.86). Within each family, a longer training run does not reliably predict a lower validation loss — the spread in both axes is driven by hyperparameter configurations (embedding size, vocabulary size) rather than a causal time–performance relationship.
1.5 Interaction: embedding dimension × vocabulary size
Across most matrices, the top-right corner (high emb_dim, large vocab) is consistently lighter, indicating that both axes contribute to lower validation loss. FATE (L=1), (L=2) and (L=4) all show essentially the same pattern: the emb_dim=64 row stands out as noticeably darker, and loss decreases as vocabulary grows, with no sign of this gradient fading as the architecture gains depth. For FastText the effect is subtler but still present. Label attention leaves the color range roughly unchanged for FATE architectures and without changing which corner performs best; for FastText, however, it noticeably lightens the whole matrix, improving loss across the board.
2 Prediction analysis
For this section, we take the best-performing model of each architecture on the test set and look at its predictions.
Across all eight architectures, the vast majority of predictions land within one star of the truth, and off-by-3+ errors stay under 8%. Correctness rises modestly with capacity (~59% for FastText vs. ~62–63% for the deeper FATE variants), but the gap is small next to the shared ±1-star error bucket. No architecture shows a strong systematic bias either: net bias stays within ±1.4% at ±2 stars and beyond, and even the largest ±1-star biases (FATE (L=1) at +1.8%, FATE (L=4) + Label Attention at −2.5%) are modest relative to the ~30% of examples that fall in that bucket.
Rather than asking which model is wrong, this view asks a deeper question: when the eight architectures disagree, is that because some of them are simply weaker or because the example itself is ambiguous (e.g. a “mixed” review that a human would also struggle to rate)? For each test example, we count how many distinct star ratings the eight models collectively predict, then look at (a) how often the majority vote is actually correct at each agreement level, and (b) which true ratings dominate the disagreement zone.
The eight architectures agree unanimously on a portion of the test set, and there the majority vote is right most of the time. Agreement collapses quickly once predictions split into 2 distinct labels and majority-vote accuracy drops significantly. The right panel explains why by looking, within each disagreement group, at the true rating of its examples (not what the models predicted): the share of extreme ratings (1★ and 5★) shrinks significantly as disagreement grows, replaced by middle ratings (2★–4★): exactly the boundary cases where a human annotator would also hesitate between adjacent stars. This backs the label-ambiguity reading: a meaningful share of the remaining error sits in middle-rating reviews, a ceiling that no architecture change is likely to break.
3 Explainability
This section turns the spotlight on why the models predict what they predict, not just how often they’re right. The analysis is built around a focused comparison: our best-performing model : FATE (L=4) without Label Attention is set against its direct counterpart FATE (L=4) with label attention (identical architecture and hyperparameters, mean pooling as the sole difference). This pairing isolates the contribution of the label attention mechanism to how the model reasons about text. Both models are probed with an identical battery of analyses on the same 200 sampled test reviews, so every figure below is a direct, paired comparison rather than two separate reports glued together.
Four complementary instruments are used throughout:
- Layer Integrated Gradients (Captum) : a gradient-based, post-hoc attribution method that assigns every input word a signed score for every class: positive = pushes the prediction toward that star rating, negative = pushes away. It only needs a forward+backward pass through the model, so it can be computed identically for both architectures which is exactly what makes the comparison fair.
- Label-attention weights : for the Label Attention model only: the raw cross-attention weights between each class query and the sentence’s tokens — an explanation built into the architecture, rather than reconstructed after the fact.
- Self-attention matrices : both models are built on the same transformer backbone, so we can also peek at how the encoder itself distributes attention across the sentence, layer by layer, before any classification happens.
- Class direction vectors : we compare two 5×emb_dim matrices, one per model: the label-attention queries from the Label Attention model, and the weight matrix of the linear classification head from the Mean Pooling model. We examine how the five star rating class vectors are arranged relative to each other in each matrix, to see what each architecture’s internal geometry reveals about the rating space.
3.1 Word attributions
Method: Captum (Integrated Gradients) vs. label-attention weights
The fastest way to trust or distrust an explanation is to look at it on a real sentence. The figure below renders a handful of test reviews as highlighted text: each word’s background is shaded in proportion to its attribution score for the selected class : use the star buttons to switch class. Green = pushes toward the selected rating, red = pushes away from it, with colour intensity proportional to |score| / max|score| within that review. By default the selected class is the model’s predicted one, but you can inspect any class, including those not predicted, to see what evidence the model associates with each star rating. Both architectures see the same review, so you can directly compare which words each one leans on.
Label Attention also produces its own built-in explanation — cross-attention weights, one query per class over the sentence — shown as a third row (blue intensity = share of attention received). Comparing it against Label Attention’s own Captum row tells you whether this built-in mechanism agrees with the post-hoc explanation of itself: a meaningful check here since, unlike on CLINC150 (150 classes, 100 training examples each), Amazon’s 5 query vectors are trained on 200,000 reviews and have ample signal to converge to discriminative directions. Each row keeps its own star selector, so you can compare different classes side by side if you like.
The selection below covers four informative cases: reviews where the two architectures disagree on the predicted rating (the richest source of qualitative insight: what did one model see that the other didn’t?); a couple where they agree and are correct (a sanity check on what a “good” explanation looks like); cases where both are wrong (what misleads them?); and cases where Mean Pooling is correct but Label Attention is not (when does the simpler architecture do better?).
3.2 Corpus-level word importance
Method: Captum (Integrated Gradients)
Zooming out from individual reviews : which words does each architecture rely on in general to recognise a 1★ vs. a 5★ review? For every one of the 200 sampled reviews, Captum produces a signed attribution for each word and each class; here we average these scores across all occurrences of each word, keeping only words seen at least 5 times, to avoid one-off noise and display the twelve most positively-attributed words per star rating.
3.3 Distinctiveness of per-class explanations
Method: Captum (Integrated Gradients)
A genuinely useful per-class explanation should look different depending on which class it explains, otherwise the model (or Captum) is essentially saying “this word is generically important”, which is far less actionable than “this word specifically signals a 5★ review”.
To measure this, we compute a 5×5 Pearson correlation matrix between the classes’ word-attribution profiles. The construction is:
- Reuse the per-class mean IG scores from the previous section. For each word seen ≥ 5 times in the 200 reviews, we already have its mean signed IG score for each of the 5 classes : call this a vocabulary profile per class.
- This gives 5 vectors of length |vocabulary|, one per star rating, where entry \(j\) reads “how much does word \(j\) push toward class \(c\) on average?”
- We compute the Pearson \(r\) between every pair of these class vectors, yielding the 5×5 symmetric matrix shown below.
How to read it: the diagonal is 1 by construction. Near +1 means two classes share the same vocabulary, e.g. 4★ and 5★ are told apart by degree, not by different words. Near −1 means mirror-image vocabularies: what pushes toward 1★ pushes away from 5★. Near 0 means the two classes rely on unrelated evidence.
Both architectures show the expected pattern along the diagonal’s neighbours: adjacent ratings correlate positively for Label Attention (1★–2★ \(r \approx 0.53\), 4★–5★ \(r \approx 0.67\)) and Mean Pooling (1★–2★ \(r \approx 0.59\), 4★–5★ \(r \approx 0.60\)), confirming they’re told apart mostly by intensity rather than by different words. The extremes don’t show a clean architectural split (\(r \approx -0.65/-0.73\) and \(-0.72/-0.58\), roughly cancelling out) — the real difference shows up around 3★, the pivot class. Label Attention’s 3★ correlates positively with 2★ (\(r \approx 0.61\)) and barely with 4★ (\(r \approx 0.14\)): its vocabulary leans negative. Mean Pooling does the opposite — 3★ is uncorrelated with 2★ (\(r \approx -0.04\)) but strongly anti-correlated with neighbour 4★ (\(r \approx -0.72\)). This fits each architecture’s design: Mean Pooling derives all 5 classes from a single shared representation, pushing any pair toward mirror images along some polarity axis — here that axis falls between 3★ and 4★; Label Attention’s per-class learned queries let 3★ partially overlap with 2★ instead of being forced into opposition with its neighbours.
3.4 Self-attention and label attention: how each mechanism reads the sentence
3.4.1 Self-attention
Method: raw self-attention weights
Both architectures are built on the same transformer backbone : they only differ in how its token-level outputs get turned into a single classification decision. So before looking at how each model aggregates information, it’s worth checking whether they end up encoding the sentence the same way internally: for the same short review, we extract every layer’s per-head self-attention matrices, for both models, so individual heads can be compared side by side.
Reading an attention heatmap: each cell (i, j) shows how much token i attends to token j when building its contextual representation (rows sum to 1). A strong diagonal means tokens mostly look at themselves and their immediate neighbours (local, syntax-like processing); vertical stripes : a column lighting up across many rows mean many tokens converge onto a few “hub” tokens, often punctuation or the very first token, a well-documented attention-sink phenomenon; a diffuse, low-contrast matrix means attention is spread thinly across the whole sentence (more global, semantic mixing).
In Layer 1, both architectures stay diffuse on both heads — attention spread thinly across the sentence, no clear focus yet.
From Layer 2 onward the two architectures start reading the sentence differently, and each one’s two heads split into different jobs. For Label Attention, one head tends to lock onto a strong, tight diagonal (each token mostly attending to itself and its close neighbours, a local, syntax-like pattern), while the other stays comparatively diffuse. For Mean Pooling, the signature pattern is different throughout: vertical bands, i.e. most tokens funnelling attention onto a handful of the same hub words, rather than onto themselves or their neighbours. The clearest instance of this is in layer 2, where one head turns almost the entire sentence’s attention onto the very last token — a textbook attention-sink.
Layer 3 continues the hub pattern for Mean Pooling — still funnelling attention onto a handful of fixed words rather than tracking the sentence’s structure. It’s also, perhaps surprisingly, where Label Attention leans hardest into hub behaviour itself: the same head that held a tight diagonal in layer 2 abandons it here and turns sharply onto a single word instead — the single most concentrated head anywhere in the model, sharper even than Mean Pooling’s own hub heads — while its partner head stays comparatively diffuse. Unlike a diagonal, this head’s target doesn’t move with the query: it’s a genuine hub, the same handful of words pulling in attention from across the whole sentence.
By Layer 4, Label Attention’s heads swing back toward locality — both heads tighten onto the diagonal, more so than at any point since layer 2’s specialist head. Mean Pooling shows its one moment of local structure here, and only in one head: instead of a fixed hub, that head’s attention band narrows and starts tracking the query position — each word looking toward a spot further along the sentence, shifting in lockstep as the query moves, more like a displaced diagonal than a hub. The other head doesn’t follow suit, staying anchored on a fixed hub regardless of where the query sits. So even at the very last layer, before the sentence gets boiled down to an unweighted average, only half of Mean Pooling’s attention machinery ever tracks position at all — the rest keeps funnelling toward fixed hub words right up to the end.
A plausible explanation: Mean Pooling has no learned mechanism to decide, at the very end, which tokens matter — the pooling step is an unweighted average — so most of its heads keep funnelling information onto fixed hub tokens all the way through, ready to dominate that average; only one head, at the last possible moment, allows itself to track the sentence’s structure instead. Label Attention doesn’t carry that burden: the actual class-specific token selection happens downstream, in its dedicated cross-attention module, so its transformer layers are free to spend their last layer on local, syntax-level refinement instead of pre-selecting which tokens will matter.
3.4.2 Label attention
Label attention heads. The label-attention module has 4 heads, each producing its own (class × word) attention matrix; these are then averaged into the single aggregated matrix used as the architecture’s built-in explanation. Note that the averaging is done purely for the explanation: the model’s actual sentence embedding concatenates the 4 heads’ outputs rather than averaging them, so the mean-over-heads matrix shown here never feeds back into the network — it is a post-hoc summary for human reading, not a computation the model itself performs.
The 4 heads seem to barely agree with each other — they don’t look like redundant copies of the same explanation, more like each reading different words. The split doesn’t look random either: Head 2 seems to lean toward negative reviews, Head 3 appears to dominate sharply on the neutral 3★, Head 4 seems to lean positive, and Head 1 rarely stands out for any class in particular, acting more like a generalist backup. The “mean over heads” row at the bottom of the figure — the one used elsewhere as the label-attention explanation — seems to quietly blend these differing, class-specialised views into one blurrier average.
3.5 Class representations: label queries vs. linear-head directions
Method: weight/embedding geometry
At the very last step, both architectures reduce a review to a decision by comparing some representation of the sentence against one reference vector per star rating:
- Label Attention learns five query embeddings (
label_embeds, one per class) that directly drive the cross-attention : they are the model’s notion of “what a 1★ review looks like”, “what a 5★ review looks like”, and so on. - Mean Pooling instead learns a plain linear classification head; each of its five weight rows (
linear_weight) plays the same role : the direction in embedding space onto which the pooled sentence vector is projected for that class.
label_embeds and linear_weight happen to be matrices of the exact same shape (5 classes × emb_dim), which makes a side-by-side comparison tempting. But same shape does not mean same role, and it’s worth being explicit about how different these two objects are before reading anything into their geometry:
linear_weightrows are decision-boundary normals: the predicted class isargmax_c (w_c · x), so the five rows live in a competitive space: pushing the sentence representation toward one row’s direction mechanically pushes it away from the others. Geometric relationships between rows (e.g. opposite directions for opposite ratings) directly shape the decision boundary.label_embedsrows are attention queries: each is projected throughW_Q, dotted with the token keys (not with each other), and softmax-normalised independently per class. Class c’s query never directly competes with class c′’s query: only their attention outputs do, several steps downstream, after going through their own value projections and the post-attention RMSNorm. The raw embedding directions are free to take whatever shape is convenient for retrieval, with no pressure to encode inter-class geometry.
Put differently, the two architectures place the competition between classes at opposite ends of the pipeline: Mean Pooling has one shared representation and five competing judges (the five linear_weight rows, which must spread apart in embedding space since they all score the same pooled vector x), whereas Label Attention has five different representations and one shared judge (the single classification-head vector, applied identically to each y_c — it needs no per-class specialisation, since the classes are already told apart upstream, by which tokens each query attended to).
So while comparing their geometry is a legitimate, mechanism-agnostic probe of “how does this architecture represent the five star ratings internally?”, a difference in pattern between the two should not be read as one architecture being more or less “structured” in some absolute sense: it mostly reflects that one matrix’s geometry is the decision rule, while the other’s geometry is several non-linear steps removed from it.
(a) Cosine similarity matrix: how aligned are the five class vectors with one another, within each model? Absolute cosine values are not comparable across the two models (the two embedding spaces are unrelated and have different dimensionalities of meaning), what matters is the pattern: are adjacent ratings (1★–2★, 4★–5★) closer to each other than distant ones (1★–5★)? A model whose class vectors trace out the ordinal structure of star ratings has, in some sense, “discovered” that this is a regression-like problem dressed up as classification.
Mean Pooling’s directions don’t quite trace a single smooth polarity axis so much as two opposing clusters: 1★-2★ and 4★-5★ each sit mildly positive with one another (\(\cos \approx 0.15\)-\(0.19\)), while any pair straddling the two — 1★/2★ against 4★/5★ — is strongly negative (\(\cos \approx -0.39\) to \(-0.50\)). The 3★ direction doesn’t sit peacefully in between: it’s negatively aligned with every other class (\(\cos \approx -0.19\) to \(-0.29\)), more like an odd one out than a true midpoint. Label Attention’s query vectors read differently: adjacent ratings are consistently positively aligned (\(\cos \approx 0.17\)-\(0.33\), weakest at 3★-4★), while anything beyond immediate neighbours sits close to 0, suggesting each class carves out its own largely independent direction rather than sitting on one shared scale.
This difference likely traces back to how each architecture is built. Mean Pooling’s 5 classes all compete for the same shared pooled representation via a linear softmax head, which mechanically pushes opposite ratings into opposite directions — a true bipolar axis. Label Attention has no such shared vector to compete over: each class gets its own independent query, so its near-orthogonal geometry beyond neighbours isn’t surprising — there’s no mechanism forcing distant classes apart or together.
(b) 2D PCA projection: a visual complement to (a). We project the five emb_dim-dimensional vectors onto their first two principal components and connect them in rating order (1★→2★→…→5★). Again, only the shape matters, not the absolute coordinates: a smooth arc or line suggests an ordinal layout; a tight, undifferentiated cluster suggests the model leans more on magnitude (how far the sentence representation sits from a boundary) than on direction to separate ratings; an erratic zig-zag would be the most surprising and most interesting finding of all.
Both models trace the same qualitative shape: a “V” rather than a smooth 1★→5★ arc, with 1★–2★ clustered on one arm, 3★ at the vertex, and 4★–5★ forming the other arm — a sentiment-polarity layout more than a strict ordinal gradient, consistent with the axis structure seen in the cosine matrix above. The two models differ, however, in how clean each arm is: Mean Pooling’s distance from the vertex grows monotonically on both arms, while Label Attention gets this right on the negative side but not the positive one, where 5★ actually sits closer to the vertex than 4★ — suggesting it separates the top two ratings less cleanly than Mean Pooling does.
3.6 Are these explanations faithful? A word-deletion stress test
Method: Captum (Integrated Gradients)
All the analyses so far take attribution scores at face value but a high attribution score is only meaningful if the model would actually change its mind were that word removed. This is the faithfulness question, and the standard way to test it is the “comprehensiveness” test from the ERASER benchmark : progressively delete the words Captum ranked as most important for the predicted class, and watch how fast that class’s predicted probability collapses. As a control, we also delete the same number of randomly chosen words, averaged over 3 draws, isolating the effect of which words are removed from the trivial effect of mutilating the sentence at all.
How to read it:
- The solid lines (guided removal) should drop faster than the dotted lines (random removal) of the same colour : the larger that gap, the more the explanation reflects words the model genuinely depends on, rather than scores that merely correlate with importance without causing it.
- If a model’s guided and random curves are nearly superimposed, its Captum attributions are not very faithful: deleting the words it called “important” hurts the prediction no more than deleting random ones — the explanation is, in a meaningful sense, decorative.
- Comparing the gap between the two curves across architectures tells you which one’s attributions are more trustworthy as a debugging/auditing tool — a property that matters in production independently of raw accuracy.
Label Attention’s two guided curves sit almost on top of each other, both well below the random control — its built-in attention weights are about as faithful a guide as the independent Captum scores, a reassuring sign it isn’t just decorative.