Overview
ThePhoenixRetrievalModel implements a two-tower architecture for efficient retrieval of relevant posts. It uses the Phoenix transformer to encode user representations and a separate candidate tower for post embeddings, enabling fast approximate nearest neighbor (ANN) search.
Classes
PhoenixRetrievalModel
A two-tower retrieval model using the Phoenix transformer for user encoding.Transformer
required
The underlying Grok transformer model for user encoding
PhoenixRetrievalModelConfig
required
Model configuration including embedding size and sequence lengths
Any
default:"jnp.bfloat16"
Forward propagation data type for computation
Optional[str]
Optional name for the model
PhoenixRetrievalModelConfig
Configuration for the Phoenix Retrieval Model. Uses the same transformer architecture as the Phoenix ranker for encoding user representations.TransformerConfig
required
Transformer architecture configuration
int
required
Embedding dimension size
int
default:"128"
Maximum length of user history sequence
int
default:"32"
Maximum number of candidate posts (used for batch processing)
Optional[str]
Optional model name
Any
default:"jnp.bfloat16"
Forward propagation data type
HashConfig
Hash configuration for multi-hash embeddings
int
default:"16"
Size of product surface vocabulary
CandidateTower
Candidate tower that projects post+author embeddings to a shared embedding space.int
required
Output embedding dimension
Optional[str]
Optional name for the tower
Named Tuples
RetrievalOutput
Output of the retrieval model.jax.Array
L2-normalized user embedding [B, D]
jax.Array
Indices of top-k retrieved candidates [B, K]
jax.Array
Similarity scores for top-k candidates [B, K]
Methods
__call__
RecsysBatch
required
Batch containing hashes, actions, and product surfaces
RecsysEmbeddings
required
Pre-looked-up embeddings from embedding tables
jax.Array
required
Normalized corpus candidate embeddings [N, D]
int
required
Number of candidates to retrieve
Optional[jax.Array]
Optional mask for valid corpus entries [N]
RetrievalOutput
Retrieval output containing user representation and top-k results
build_user_representation
RecsysBatch
required
Batch containing hashes, actions, and product surfaces
RecsysEmbeddings
required
Pre-looked-up embeddings from embedding tables
jax.Array
L2-normalized user embedding [B, D]
jax.Array
Pre-normalization L2 norm [B, 1]
build_candidate_representation
RecsysBatch
required
Batch containing candidate hashes
RecsysEmbeddings
required
Pre-looked-up embeddings containing candidate embeddings
jax.Array
L2-normalized candidate embeddings [B, C, D]
jax.Array
Valid candidate mask [B, C]
_retrieve_top_k
jax.Array
required
Normalized user embeddings [B, D]
jax.Array
required
Normalized corpus candidate embeddings [N, D]
int
required
Number of candidates to retrieve
Optional[jax.Array]
Optional mask for valid corpus entries [N]
jax.Array
Indices of top-k candidates [B, K]
jax.Array
Similarity scores of top-k candidates [B, K]
CandidateTower Methods
__call__
jax.Array
required
Concatenated post and author embeddings. Shape: [B, C, num_hashes, D] or [B, num_hashes, D]
jax.Array
Normalized candidate representation. Shape: [B, C, D] or [B, D]
Architecture
The two-tower architecture consists of:-
User Tower: Encodes user features + history using the Phoenix transformer
- Processes user embeddings and interaction history
- Outputs a single L2-normalized user representation vector
-
Candidate Tower: Projects candidate embeddings to a shared space
- Two-layer MLP with SiLU activation
- Projects concatenated post+author embeddings
- Outputs L2-normalized candidate representations
-
Retrieval: Uses dot product similarity between normalized embeddings
- Enables efficient ANN search in production
- Returns top-k candidates based on similarity scores
Usage Notes
- Both user and candidate representations are L2-normalized
- Dot product similarity is used for retrieval (equivalent to cosine similarity for normalized vectors)
- The model is designed for efficient ANN search using libraries like FAISS or ScaNN
- Constants: EPS = 1e-12, INF = 1e12 for numerical stability