Supplement · The X algorithm · page 5 of 11

Ranking: the transformer that forbids candidates to look at each other

Retrieval has selected hundreds of posts — now they have to be ordered. Let us go through the Phoenix ranking model: how the history and the candidates are laid into one sequence, what exactly the attention mask does, why there are many predictions and how they are trained when the data for different heads comes from different places.

1. How everything is laid into one sequence

The first thing to get straight: this is not a model of the kind «feed in a (user, post) pair — get a number». One pass processes the whole set of candidates at once, and it is laid out like this:

Such a layout is a direct continuation of what we discussed in the chapter «Transformers over history»: the user's history as a sequence, a candidate as one more token, attention instead of hand-made aggregates. The difference is that here there are 64 candidates at once, and that is exactly what gives rise to the main question of this page.

2. The isolation mask: one line that determines the whole system

On the overview page we mentioned that candidates do not see each other. Now let us look at how that is written down. Every token is assigned a segment:

HISTORY_SEGMENT_ID = 1
CANDIDATE_SEGMENT_ID = -1
PADDING_SEGMENT_ID = 0

segment_ids = jnp.where(padding_mask,
                        jnp.where(idx >= candidate_start_offset, -1, 1),
                        0)

And inside the attention kernel stands this:

mask = jnp.logical_or(seg_k == HISTORY_SEGMENT_ID, span_q[:, None] == span_k[None, :])

A fragment of ranker_attention_v2.py · code by X, Apache 2.0, commit 28e414f

Let us read it slowly. A query at position \(q\) may look at a key at position \(k\) if and only if at least one of two things holds:

• the key belongs to the history or the profile — seg_k == 1;
• the key is the query itself — span_q == span_k.

That is all. No other options. Which means:

Note: the history is bidirectional. There is no causal mask by default, and an event from the middle of the history looks calmly at a later one. This is not a language model predicting the next token; it is an encoder, and there is no reason to forbid it to look forward along the history.

acc, m_i, l_i = lax.fori_loop(0, history_upper_bound, body, (acc, m_i, l_i))

candidate_lower_bound = jnp.maximum(history_upper_bound, offset_q // block_k)
candidate_upper_bound = (offset_q + block_q - 1) // block_k + 1
acc, m_i, l_i = lax.fori_loop(candidate_lower_bound, candidate_upper_bound, body, (acc, m_i, l_i))

What is paid for this

The trade-off is worth stating honestly, because in an interview it is precisely the flip side that is usually asked about.

A model that sees the whole slate is fundamentally stronger. It can notice that three posts in a row are about the same thing, that the second post duplicates the first, that after a long video a short text should be put. That is the listwise formulation, and it gives better quality all else being equal.

Having given it up, the system gets four things:

1. Cacheability. The score of a (user, post) pair can be saved and reused on the next request. In the code of the feed there is a separate source of cached ranked posts — it exists precisely thanks to this property.
2. Reproducibility. Two identical requests give the same order. Without the isolation the order would depend on how the posts were split into batches.
3. Explainability. The question «why did this post end up lower» has an answer that does not contain the words «because another post happened to be next to it».
4. Speed. The quadratic term in the candidates disappears.

And the listwise effects are brought back later and separately — by re-ranking through a determinantal process, which we discussed on the scoring page. The result is a division of labour: the model answers «how good is this post for this person», and a separate service «how do these posts look together».

3. The heads: what exactly is predicted

Two sets of numbers are taken from every candidate position.

The discrete heads are the probabilities of actions: a like, a reply, a repost, a quote, a click, expanding a photo, opening a video, following the author, a report, a block, hiding, «not interested» and so on. The full list and the weights they are combined with are on the scoring page.

The continuous heads are quantities that have no «happened or not»: the dwell time on a post, the time after a click, the active seconds. That is regression. There are eight of them in the config.

4. The most interesting part: the heads learn on different subsamples

Here begins what makes someone else's production code worth reading. Naively it seems that all the heads are trained on the same examples. In reality it is decided separately for every head whether this example is informative.

Conversions — only on clicked posts

if self.config.condition_conversion_on_click:
    has_click = targets[:, :, CLICK_ACTION_INDEX]
    no_click = 1 - has_click
    conv_head_mask = jnp.zeros(num_actions).at[jnp.array(CLICK_CONDITIONED_ACTION_INDICES)].set(1.0)
    conv_zero_mask = no_click[:, :, None] * conv_head_mask
    loss_mask = loss_mask * (1 - conv_zero_mask)

A fragment of recsys_model.py · code by X, Apache 2.0, commit 28e414f

It reads like this: if there was no click on the post, then the loss of the conversion heads on this example is zeroed out. They get no gradient at all.

The logic is ironclad. What does «the conversion did not happen» mean for a post that was not clicked? Nothing. The user did not refuse the action — they simply had no opportunity to perform it. Training a head on such an example means teaching it that «no opportunity» is «a refusal».

This is exactly the sample selection bias for the sake of which ESMM was invented in the chapter «ESMM and sample selection bias»: there it was shown that a conversion model trained only on clicks makes systematic errors when rolled out over all impressions. There are two solutions — either model the whole chain through a product of probabilities, or honestly restrict the training sample. Here the second path is chosen, but with an important addition: the training is joint, in one model, so the shared body still sees all the examples and only the head specialises.

Negative feedback — only on real impressions

if self.config.mask_neg_feedback_on_negatives:
    neg_head_mask = jnp.zeros(num_actions).at[jnp.array(NEGATIVE_FEEDBACK_HEAD_INDICES)].set(1.0)
    zero_mask = negative_sample_mask[:, :, None] * neg_head_mask
    loss_mask = loss_mask * (1 - zero_mask)

A fragment of recsys_model.py · code by X, Apache 2.0, commit 28e414f

Here the same idea from the other side. Sampled negatives are added to the training set — posts that were never shown to the user. For heads like «will they like it» that is fine: it was not shown, so they almost certainly would not have liked it.

But for the head «will they report it» such an example is poisonous. The absence of a report on an unshown post is not evidence that there would have been no report. Having trained on that, we get a model convinced that reports are rare — and it will be statistically right and practically useless.

So the negative-feedback heads are not trained on sampled negatives at all. We discussed exactly this trap in the chapter «The problems of implicit data»: the absence of a signal and a negative signal are different things, and confusing them is expensive.

zero_mask = has_negative_feedback_action_expanded & non_negative_feedback_action_mask
targets_for_loss = jnp.where(zero_mask, False, targets_for_loss)

5. Features: what the model knows about a post and about the moment

The FeaturePrepConfig block in the config is essentially a list of everything fed to the model besides the history and the identifiers.

Two rows deserve a separate look.

hour_of_day_dither_fraction: 0.1

post_age_granularity_mins: 60

6. Training

There are several decisions here, and each deserves an explanation.

Two optimisers for one model

optim_config=RecsysDenseOptimConfig(
    optim="muon",
    muon_consistent_rms=0.2,
    muon_matrix_weight_decay=0.014,
    b1=0.95, b2=0.98,
),
emb_optim_config=RecsysEmbeddingOptimConfig(
    rowwise_adagrad=RecsysRowwiseAdagradConfig(
        learning_rate=0.28, half_life_steps=2500, lazy_decay=True, weight_decay=2.8e-4,
    ),
),

A fragment of xrecsys.py · code by X, Apache 2.0, commit 28e414f

The dense layers are trained with Muon, the embedding tables with row-wise Adagrad. That is not eclecticism but a consequence of their fundamentally different gradient structure.

• A dense matrix gets a gradient on every example of the batch. Muon is a comparatively new optimiser that works with a weight matrix as a matrix rather than as a flat vector of parameters: it evens out the spectrum of the update, not letting one direction dominate. The learning rate is 7.1e-4.
• A row of the embedding table gets a gradient only when its item lands in the batch — that is, thousands of times less often. It needs a large step (0.28, four hundred times bigger!) and individual adaptation: Adagrad accumulates the history of gradients for every row separately.

The parameter half_life_steps: 2500 sets the half-life of Adagrad's accumulated statistics: without forgetting, the denominator grows monotonically, the step goes to zero and old rows stop learning. With forgetting the model stays able to adapt to changed behaviour — a direct answer to data drift.

And lazy_decay: True means the forgetting is applied to a row lazily, at the moment of its next update, rather than to the whole table on every step. For a table of millions of rows that is the difference between «expensive» and «free».

Packing the sequences

People's histories are of different lengths, while a tensor is rectangular. The naive solution — pad everyone up to the maximum — means that at an average length of 510 and a maximum of 1022, half the computation goes on padding.

In the config use_seqpack is on, with a Beta distribution of lengths from 126 to 1022 with a mean of 510. Several short histories are laid into one physical row, and they are separated by those same segment identifiers we looked at in the mask. Attention does not let them flow into each other.

An elegant detail: the mechanism of segments, introduced for the sake of isolating the candidates, solved the packing task for free as well. One and the same field serves two purposes.

The LogQ correction here too

The ranking config has log_q_correction: True. That may look strange: the correction was discussed in the context of retrieval, where the negatives are sampled from the batch. But the training of the ranker also has sampled negatives — posts that were not shown to the user, added so that the model sees more than what was served.

They are not sampled uniformly, and without a correction the ranker would inherit the same bias towards the popular. The mechanics are the same as those discussed in detail on the retrieval page and in the LogQ widget.

7. The ranker's configuration against the retriever's

The row about the normalisation of queries and keys deserves a comment. qk_norm normalises the vectors before the dot product is computed, bounding the magnitude of the attention logits from above. The alternative is to clip the logits hard, which is what the retriever does with the parameter 80.0. Both tricks cure one disease: diverging attention, where one logit runs away, the softmax collapses into a delta function and the gradient dies. In the ranker the softer of the two was chosen.

Common mistakes and hidden rocks

Interview questions

One-screen cheat sheet

Primary sources

• recsys_model.py — the model, the segment identifiers, the per-head loss masks.
• ranker_attention_v2.py — the attention kernel and that line of the mask.
• xrecsys.py — the production config: sizes, features, optimisers.
• ranker_attention_fa4.py — the block-sparse attention layout for newer GPUs.
• The course: the chapter «Transformers over history» on transformers over history, the chapter «ESMM and sample selection bias» on ESMM and sample selection bias, the chapter «Listwise losses» on listwise formulations.