# Candidate isolation

> make_recsys_attn_mask rules: candidates attend to user and history plus self, never to other candidates, so scores stay batch-independent.

- Repository: xai-org/x-algorithm
- GitHub: https://github.com/xai-org/x-algorithm
- Human docs: https://grok-wiki.com/public/docs/xai-org-x-algorithm-23c09c39074c
- Complete Markdown: https://grok-wiki.com/public/docs/xai-org-x-algorithm-23c09c39074c/llms-full.txt

## Source Files

- `phoenix/grok.py`
- `phoenix/recsys_model.py`
- `phoenix/test_recsys_model.py`
- `phoenix/README.md`
- `phoenix/recsys_retrieval_model.py`

---

---
title: "Candidate isolation"
description: "make_recsys_attn_mask rules: candidates attend to user and history plus self, never to other candidates, so scores stay batch-independent."
---

`make_recsys_attn_mask` in `phoenix/grok.py` is the ranking attention contract. `RecsysModel.build_inputs` concatenates one user token, the history window, and the candidate window; `Transformer.__call__` then multiplies that mask into the padding mask so each candidate can attend to user plus history and itself, and never to another candidate. Blocked keys are written to `-1e30` before the fp32 softmax, so a candidate's logits do not depend on which other candidates share the sequence.

<Info>
Isolation is ranking-only. The retrieval user tower calls the same `Transformer` with `candidate_start_offset=None` and keeps a standard causal mask over `[user | history]`.
</Info>

## Sequence layout

`RecsysModel.build_inputs` concatenates three blocks and returns `candidate_start_offset` as the first candidate index:

```text
[ user (1) | history (S) | candidates (C) ]
             ^             ^
             1             candidate_start_offset = 1 + S
```

| Block | Source | Token count | Padding (`True` = valid) |
| --- | --- | --- | --- |
| User | `block_user_reduce` | 1 | `user_hashes[:, 0] != 0` |
| History | `block_history_reduce` | `history_seq_len` (`S`) | `history_post_hashes[:, :, 0] != 0` |
| Candidates | `block_candidate_reduce` | `candidate_seq_len` (`C`) | `candidate_post_hashes[:, :, 0] != 0` |

`candidate_start_offset` is `user_padding_mask.shape[1] + history_padding_mask.shape[1]`, not a config field. Hash `0` is reserved padding on user, history, and candidate first-hash slots.

<ParamField body="history_seq_len" type="int" default="128">
`PhoenixModelConfig` history window. Live ranker length comes from published `config.json` via `build_model_config` in `run_pipeline.py`.
</ParamField>

<ParamField body="candidate_seq_len" type="int" default="32">
`PhoenixModelConfig` candidate window packed into one forward. Isolation is what makes packing safe.
</ParamField>

## Mask rules

`make_recsys_attn_mask(seq_len, candidate_start_offset, dtype=jnp.float32)` returns `[1, 1, seq_len, seq_len]` with `1` = can attend and `0` = blocked.

| Query positions | Allowed keys | Blocked keys |
| --- | --- | --- |
| `0 .. offset-1` (user + history) | Causal: key `j` only if `j <= i` | Future user/history tokens; every candidate |
| `offset .. seq_len-1` (candidates) | All user + history keys, plus the query's own index | Every other candidate |

User and history use a **causal** lower triangle (`jnp.tril`), not bidirectional attention. Candidates attend to the entire prefix, including history tokens that later history tokens cannot yet see.

The `TestMakeRecsysAttnMask` fixture `[user, h1, h2, c1, c2, c3]` (`seq_len=6`, `offset=3`) is the contract:

```text
         U   h1  h2  c1  c2  c3
    U    1   0   0   0   0   0
    h1   1   1   0   0   0   0
    h2   1   1   1   0   0   0
    c1   1   1   1   1   0   0
    c2   1   1   1   0   1   0
    c3   1   1   1   0   0   1
```

Edge cases covered by the same tests: a single candidate is prefix-plus-self; `offset=1` (no history) is user-plus-self per candidate.

## Construction

The function does not start from a custom sparse pattern. It edits a full causal mask in three steps:

1. `causal_mask = jnp.tril(ones([1, 1, seq_len, seq_len]))`
2. Zero the candidate–candidate block: `[:, :, offset:, offset:] = 0`
3. Restore the candidate diagonal: `attn_mask[:, :, i, i] = 1` for `i` in `range(offset, seq_len)`

Step 2 is what removes later-candidate → earlier-candidate edges that a plain causal mask would allow.

<ParamField body="seq_len" type="int" required>
Total length `1 + S + C`.
</ParamField>

<ParamField body="candidate_start_offset" type="int" required>
First candidate index. Must match the concatenated layout from `build_inputs`.
</ParamField>

<ParamField body="dtype" type="jnp.dtype" default="jnp.float32">
Mask dtype. `Transformer` passes `embeddings.dtype` (`PhoenixModelConfig.fprop_dtype` defaults to `jnp.bfloat16`).
</ParamField>

## How the transformer applies it

```mermaid
flowchart TD
  subgraph ranker ["RecsysModel.__call__"]
    BI["build_inputs → embeddings, padding_mask, offset"]
    ROPE{"right_anchored_rope?"}
    POS["right_anchored_rope_positions"]
    TF["Transformer(..., candidate_start_offset=offset)"]
    CUT["out_embeddings[:, offset:, :]"]
    HEAD["action logits + continuous sigmoid"]
    BI --> ROPE
    ROPE -->|true| POS --> TF
    ROPE -->|false, default| TF
    TF --> CUT --> HEAD
  end

  subgraph transformer ["Transformer.__call__"]
    PAD["padding_mask[:, None, None, :]"]
    OFF{"candidate_start_offset is not None?"}
    ISO["make_recsys_attn_mask"]
    CAU["jnp.tril causal"]
    MUL["mask = padding * attn"]
    MHA["MHA: where(mask, logits, -1e30) then softmax"]
    PAD --> OFF
    OFF -->|ranker| ISO --> MUL
    OFF -->|retrieval user tower| CAU --> MUL
    MUL --> MHA
  end

  TF --> PAD
```

The padding mask is a **key** mask: invalid (hash-`0`) positions cannot be attended to, even when the isolation mask would allow the index. Decoder layers do not consume `padding_mask` again after this multiply.

`MultiHeadAttention` keeps softmax in fp32, tanh-clips logits at 30, then replaces masked locations with `-1e30`. Isolation is attention-only; the pointwise FFN does not mix candidates.

## Ranker vs retrieval

| Caller | `candidate_start_offset` | Sequence | Effect |
| --- | --- | --- | --- |
| `RecsysModel.__call__` | `1 + S` | `[user \| history \| candidates]` | Isolation on |
| Retrieval user tower | `None` | `[user \| history]` | Standard causal only |
| Retrieval `CandidateTower` | n/a | No transformer over the candidate set | Candidates encoded independently by construction |

Omitting `candidate_start_offset` on a packed ranker sequence is a silent regression: later candidates regain causal access to earlier candidates.

## RoPE companion

When `PhoenixModelConfig.right_anchored_rope` is `True` (dataclass default is `False`), `right_anchored_rope_positions` assigns every candidate the same position `history_end = num_user_prefix_tokens + history_seq_len`. `RecsysModel` hard-codes `num_user_prefix_tokens=1`. Padded tokens get position `0`. Shared candidate positions keep rotary encodings from encoding batch order; they do not replace the attention mask.

## Outputs after isolation

Only candidate positions are unembedded:

| Field | Shape | Source |
| --- | --- | --- |
| `RecsysModelOutput.logits` | `[B, C, num_actions]` | `dot(candidate_embeddings, unembedding)` |
| `RecsysModelOutput.continuous_preds` | `[B, C, num_continuous]` | `sigmoid(dot(candidate_embeddings, continuous_head))` |

A candidate's vector is a function of user, unpadded history, and that candidate. Changing neighbors in the `C` window does not change those logits.

The repo states this keeps scores consistent across batch composition and therefore cacheable. Isolation itself does not implement a cache.

## Verify

From `phoenix/`:

```bash
uv run pytest test_recsys_model.py::TestMakeRecsysAttnMask -v
```

`TestMakeRecsysAttnMask` asserts shape `[1, 1, T, T]`, causal prefix, full prefix visibility from every candidate, self-attend, zero off-diagonal candidate block, `float32`/`float16` dtypes, and the two edge layouts above. `TestRightAnchoredRopePositions.test_candidates_share_position` covers the optional RoPE companion.

<Warning>
Do not treat user/history as fully visible to each other. The implemented prefix is causal. A bidirectional prefix would be a behavior change and would fail `test_user_history_has_causal_attention`.
</Warning>

## Contributor constraints

- Always pass `candidate_start_offset` from `build_inputs` into `Transformer`. `None` is only correct when the sequence has no candidate block.
- Keep `offset` aligned with `[user | history | candidates]`. Shifting the cut mixes history into the isolated block or candidates into the causal prefix.
- Do not add candidate–candidate edges for diversity or calibration inside this mask. Diversity and weighted scoring happen after logits (`WeightedScorer`, `AuthorDiversityScorer`).
- Isolation does not hide padded keys. Preserve hash-`0` padding so the key mask stays correct.

## Related pages

<CardGroup>
  <Card title="Phoenix model configuration" href="/phoenix-model-configuration">
    PhoenixModelConfig fields, TransformerConfig, RecsysBatch, and published config.json keys that set S and C.
  </Card>
  <Card title="Hash embeddings" href="/hash-embeddings">
    How RecsysBatch hashes are looked up before the isolated transformer.
  </Card>
  <Card title="Multi-action scoring" href="/multi-action-scoring">
    Per-action logits taken from isolated candidate embeddings.
  </Card>
  <Card title="Test Phoenix" href="/test-phoenix">
    Attention-mask and retrieval pytest targets.
  </Card>
  <Card title="Run the inference pipeline" href="/run-inference-pipeline">
    Retrieve then rank with the published checkpoints.
  </Card>
</CardGroup>