Phoenix ML モデル仕様書

[User Embedding] + [History Embeddings (S tokens)]
     1 token              S tokens (default: 128)

# 1. ユーザー埋め込みの生成
user_embeddings, user_padding_mask = block_user_reduce(
    user_hashes,           # [B, num_user_hashes]
    user_embeddings,       # [B, num_user_hashes, D]
    num_user_hashes,       # 2
    emb_size,             # 128
)
# Output: [B, 1, D]

# 2. 履歴埋め込みの生成
history_embeddings, history_padding_mask = block_history_reduce(
    history_post_hashes,        # [B, S, num_item_hashes]
    history_post_embeddings,    # [B, S, num_item_hashes, D]
    history_author_embeddings,  # [B, S, num_author_hashes, D]
    history_product_surface_embeddings,  # [B, S, D]
    history_actions_embeddings,          # [B, S, D]
    num_item_hashes,   # 2
    num_author_hashes, # 2
)
# Output: [B, S, D]

# 3. シーケンス連結
embeddings = concat([user_embeddings, history_embeddings], axis=1)
# Shape: [B, 1+S, D]

# 4. Transformer エンコーディング（因果的アテンション）
model_output = transformer(embeddings, padding_mask)

# 5. Mean Pooling + L2 正規化
user_representation = mean_pool(model_output, padding_mask)
user_representation = l2_normalize(user_representation)
# Output: [B, D]

def _retrieve_top_k(
    self,
    user_representation: jax.Array,   # [B, D]
    corpus_embeddings: jax.Array,     # [N, D]
    top_k: int,
    corpus_mask: Optional[jax.Array] = None,
) -> Tuple[jax.Array, jax.Array]:
    """Top-K候補の検索

    Returns:
        top_k_indices: [B, K] 上位K件のインデックス
        top_k_scores: [B, K] 類似度スコア
    """
    # 内積による類似度計算
    scores = jnp.matmul(user_representation, corpus_embeddings.T)

    # マスク適用（無効な候補を除外）
    if corpus_mask is not None:
        scores = jnp.where(corpus_mask[None, :], scores, -INF)

    # Top-K選択
    top_k_scores, top_k_indices = jax.lax.top_k(scores, top_k)

    return top_k_indices, top_k_scores

product_surface_vocab_size: int = 16  # サーフェスの種類数

@dataclass
class HashConfig:
    """ハッシュベース埋め込みの設定"""

    num_user_hashes: int = 2     # ユーザーIDに使用するハッシュ関数の数
    num_item_hashes: int = 2     # 投稿IDに使用するハッシュ関数の数
    num_author_hashes: int = 2   # 著者IDに使用するハッシュ関数の数

@dataclass
class TransformerConfig:
    emb_size: int           # 埋め込み次元
    key_size: int           # アテンションキーの次元
    num_q_heads: int        # クエリヘッド数
    num_kv_heads: int       # キー/バリューヘッド数（GQA対応）
    num_layers: int         # Transformer レイヤー数
    widening_factor: float = 4.0          # FFN拡張係数
    attn_output_multiplier: float = 1.0   # アテンション出力スケール

    name: Optional[str] = None

    def make(self) -> "Transformer":
        return Transformer(
            num_q_heads=self.num_q_heads,
            num_kv_heads=self.num_kv_heads,
            widening_factor=self.widening_factor,
            key_size=self.key_size,
            attn_output_multiplier=self.attn_output_multiplier,
            num_layers=self.num_layers,
        )

emb_size = 128
num_actions = 19
history_seq_len = 32
candidate_seq_len = 8

hash_config = HashConfig(
    num_user_hashes=2,
    num_item_hashes=2,
    num_author_hashes=2,
)

recsys_model = PhoenixModelConfig(
    emb_size=emb_size,
    num_actions=num_actions,
    history_seq_len=history_seq_len,
    candidate_seq_len=candidate_seq_len,
    hash_config=hash_config,
    product_surface_vocab_size=16,
    model=TransformerConfig(
        emb_size=emb_size,
        widening_factor=2,
        key_size=64,
        num_q_heads=2,
        num_kv_heads=2,
        num_layers=2,
        attn_output_multiplier=0.125,
    ),
)

# Q/K/V 射影
query_heads = self._linear_projection(query, self.key_size, self.num_q_heads)
key_heads = self._linear_projection(key, self.key_size, self.num_kv_heads)
value_heads = self._linear_projection(value, self.value_size, self.num_kv_heads)

# Rotary Position Embedding (RoPE) 適用
rotate = RotaryEmbedding(dim=self.key_size)
key_heads = rotate(key_heads, seq_dim=1, offset=0)
query_heads = rotate(query_heads, seq_dim=1, offset=0)

# GQA: query heads をグループ化
query_heads = query_heads.reshape((b, t, kv_h, h // kv_h, d))

# アテンションスコア計算
attn_logits = jnp.einsum("...thHd,...Thd->...hHtT", query_heads, key_heads)
attn_logits *= self.attn_output_multiplier

# Soft-capping（数値安定性のため）
max_attn_val = 30.0
attn_logits = max_attn_val * jnp.tanh(attn_logits / max_attn_val)

# マスク適用 + Softmax
attn_logits = jnp.where(mask, attn_logits, -1e30)
attn_weights = jax.nn.softmax(attn_logits)

@dataclass
class DenseBlock(hk.Module):
    """SwiGLU 活性化を使用した FFN"""

    widening_factor: float = 4.0

    def __call__(self, inputs: jax.Array) -> jax.Array:
        _, _, model_size = inputs.shape

        # SwiGLU: gate * GELU(x)
        h_v = Linear(ffn_size)(inputs)           # Value branch
        h_w1 = jax.nn.gelu(Linear(ffn_size)(inputs))  # Gate branch

        h_dense = Linear(model_size)(h_w1 * h_v)  # Output projection

        return h_dense

class RotaryEmbedding(hk.Module):
    """回転位置埋め込み（RoPE）
    参考: https://arxiv.org/abs/2104.09864
    """

    def __init__(self, dim: int, base_exponent: int = 10000):
        self.dim = dim
        self.base_exponent = base_exponent

    def __call__(self, x: jax.Array, seq_dim: int, offset: jax.Array) -> jax.Array:
        # 周波数の計算
        exponents = jnp.arange(0, self.dim, 2, dtype=jnp.float32)
        inv_freq = 1.0 / (self.base_exponent ** (exponents / self.dim))

        # 位置インデックス
        t = jnp.arange(x.shape[seq_dim]) + offset

        # 位相角の計算
        phase = jnp.einsum("bi,j->bij", t, inv_freq)
        phase = jnp.tile(phase, reps=(1, 2))

        # 回転の適用
        x = x * jnp.cos(phase) + rotate_half(x) * jnp.sin(phase)

        return x

# 1. モデル設定
recsys_model = PhoenixModelConfig(...)

# 2. 推論ランナーの作成と初期化
inference_runner = RecsysInferenceRunner(
    runner=ModelRunner(model=recsys_model, bs_per_device=0.125),
    name="recsys_local",
)
inference_runner.initialize()

# 3. バッチの作成
batch, embeddings = create_example_batch(
    batch_size=1,
    emb_size=emb_size,
    history_len=history_seq_len,
    num_candidates=candidate_seq_len,
    num_actions=num_actions,
    ...
)

# 4. ランキング実行
ranking_output = inference_runner.rank(batch, embeddings)

# 5. 結果の取得
scores = ranking_output.scores        # [B, C, num_actions]
ranked_indices = ranking_output.ranked_indices  # [B, C]
p_favorite = ranking_output.p_favorite_score    # [B, C]

# 1. モデル設定
retrieval_model_config = PhoenixRetrievalModelConfig(...)

# 2. 推論ランナーの作成と初期化
inference_runner = RecsysRetrievalInferenceRunner(
    runner=RetrievalModelRunner(model=retrieval_model_config, bs_per_device=0.125),
    name="retrieval_local",
)
inference_runner.initialize()

# 3. コーパスの設定
corpus_embeddings, corpus_post_ids = create_example_corpus(
    corpus_size=1000,
    emb_size=emb_size,
)
inference_runner.set_corpus(corpus_embeddings, corpus_post_ids)

# 4. バッチの作成
batch, embeddings = create_example_batch(...)

# 5. 検索実行
retrieval_output = inference_runner.retrieve(
    batch,
    embeddings,
    top_k=10,
)

# 6. 結果の取得
top_k_indices = retrieval_output.top_k_indices  # [B, K]
top_k_scores = retrieval_output.top_k_scores    # [B, K]
user_representation = retrieval_output.user_representation  # [B, D]