Skip to main content

​
Overview

The Grok transformer provides the core neural architecture for both ranking and retrieval models. It implements a decoder-only transformer with grouped-query attention (GQA), RoPE positional encodings, and RMS normalization.

​
Classes

​
Transformer

A transformer stack implementing the decoder architecture.
​
int
required
Number of query attention heads
​
int
required
Number of key/value attention heads (for grouped-query attention)
​
int
required
Dimension of attention keys and queries
​
float
required
Factor for expanding FFN hidden dimension
​
float
required
Multiplier applied to attention logits
​
int
required
Number of transformer layers
​
Optional[str]
Optional name for the transformer

​
TransformerConfig

Configuration dataclass for the transformer architecture.
​
int
required
Embedding dimension size
​
int
required
Dimension of attention keys
​
int
required
Number of query heads
​
int
required
Number of key/value heads
​
int
required
Number of transformer layers
​
float
default:"4.0"
FFN widening factor
​
float
default:"1.0"
Attention output multiplier
​
Optional[str]
Optional configuration name

​
DecoderLayer

A single transformer decoder layer.
​
int
required
Number of query attention heads
​
int
required
Number of key/value attention heads
​
int
required
Dimension of attention keys
​
int
required
Total number of layers in the stack
​
Optional[int]
Index of this layer in the stack
​
float
default:"4.0"
FFN widening factor
​
Optional[str]
Optional layer name
​
float
default:"1.0"
Attention output multiplier

​
MultiHeadAttention

Multi-head attention with grouped-query attention and RoPE.
​
int
required
Number of query heads
​
int
required
Number of key/value heads
​
int
required
Dimension of keys and queries
​
bool
default:"True"
Whether to use bias in projections
​
Optional[int]
Dimension of values (defaults to key_size)
​
Optional[int]
Model dimension (defaults to key_size * num_q_heads)
​
float
default:"1.0"
Multiplier for attention logits
​
Optional[str]
Optional module name

​
RotaryEmbedding

Applies rotary positional embeddings (RoPE) as described in RoFormer.
​
int
required
Dimensionality of the feature vectors (must be even)
​
Optional[str]
Optional module name
​
int
default:"10000"
Base exponent for computing frequencies

​
Named Tuples

​
TransformerOutput

Output of the transformer.
​
jax.Array
Output embeddings from the transformer [B, T, D]

​
DecoderOutput

Output of a decoder layer.
​
jax.Array
Output embeddings from the layer [B, T, D]

​
MHAOutput

Output of multi-head attention.
​
jax.Array
Output embeddings from attention [B, T, D]

​
TrainingState

Container for training state.
​
hk.Params
Model parameters

​
Methods

​
Transformer.__call__

Transforms input embedding sequences to output embedding sequences.
​
jax.Array
required
Input embeddings [B, T, D]
​
jax.Array
required
Padding mask [B, T], True for valid positions
​
Optional[int]
If provided, positions >= this offset are treated as candidates that can only attend to positions before the offset (user+history) and themselves (self-attention), but not to other candidates. Used for recommendation system inference.
​
TransformerOutput
Transformer output containing embeddings [B, T, D]

​
RotaryEmbedding.__call__

Apply rotary embeddings to input tensor.
​
jax.Array
required
Input tensor to apply RoPE to
​
int
required
Dimension index corresponding to the sequence
​
jax.Array
required
Position offset (scalar or per-batch element)
​
Optional[int]
Use constant position for all tokens if provided
​
Optional[jax.Array]
Custom position indices [B, T]
​
jax.Array
Tensor with rotary embeddings applied

​
Utility Functions

​
make_recsys_attn_mask

Create attention mask for recommendation system inference. Creates a mask where:
  • Positions 0 to candidate_start_offset-1 (user+history): causal attention
  • Positions candidate_start_offset onwards (candidates): can attend to user+history and themselves (self-attention), but NOT to other candidates
This ensures each candidate is scored independently based on user+history context.
​
int
required
Total sequence length (user + history + candidates)
​
int
required
Position where candidates start in the sequence
​
jnp.dtype
default:"jnp.float32"
Data type for the mask
​
jax.Array
Attention mask [1, 1, seq_len, seq_len] where 1 means “can attend”

​
ffn_size

Calculate FFN hidden dimension from embedding size and widening factor.
​
int
required
Embedding dimension
​
float
required
Widening factor for FFN
​
int
FFN hidden size (adjusted to be multiple of 8)

​
layer_norm

Apply RMS normalization to input.
​
jax.Array
required
Input tensor
​
jax.Array
RMS-normalized tensor

​
rotate_half

Obtain the rotated counterpart of each feature for RoPE.
​
jax.Array
required
Input tensor
​
jax.Array
Rotated tensor

​
Architecture Details

​
Attention Mechanism

  • Grouped-Query Attention (GQA): Reduces KV cache size by sharing key/value heads across multiple query heads
  • Rotary Positional Embeddings (RoPE): Encodes positional information directly into attention keys and queries
  • Attention Clipping: Logits are clipped using tanh to prevent overflow: 30.0 * tanh(logits / 30.0)

​
Normalization

  • Uses RMS normalization instead of LayerNorm for efficiency
  • Applied before attention and FFN blocks (pre-norm architecture)

​
Feed-Forward Network

  • Uses GeGLU activation: GELU(W1 * x) * (W2 * x)
  • Hidden dimension calculated as: int(widening_factor * emb_size) * 2 // 3
  • Adjusted to be a multiple of 8 for hardware efficiency

​
Special Attention Mask

The make_recsys_attn_mask function creates a specialized attention pattern for ranking:
  • User and history tokens use causal attention (can attend to previous tokens)
  • Candidate tokens can attend to all user+history tokens and themselves
  • Candidates cannot attend to other candidates (ensures independent scoring)