Architectures / Recommendation

📈 HSTU (Generative Recommender)

Meta 2024 - recommendation as sequential transduction over one interleaved item+action stream, with pointwise aggregated attention instead of softmax.

From Zhai et al. (2024). Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations. ICML 2024. This page is the graph, not the PDF: open it, edit it, verify it, export it.

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

10 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 43 structural checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class HSTUBlock(nn.Module):
    """Hierarchical Sequential Transduction Unit (Zhai et al., ICML 2024).
    Pointwise aggregated attention: NO softmax over the sequence, and four
    projections (u, v, q, k) rather than a transformer block's attention plus
    a separate FFN. It is not a transformer block and does not cost like one."""

    def __init__(self, embed_dim: int, num_heads: int = 4,
                 linear_dim: int = 128, attn_dim: int = 128):
        super().__init__()
        self.h, self.dv, self.dqk = num_heads, linear_dim, attn_dim
        self.uvqk = nn.Linear(embed_dim, num_heads * (2 * linear_dim + 2 * attn_dim))
        self.out = nn.Linear(num_heads * linear_dim, embed_dim)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, t, _ = x.shape
        h, dv, dqk = self.h, self.dv, self.dqk
        u, v, q, k = self.uvqk(x).split([h * dv, h * dv, h * dqk, h * dqk], dim=-1)
        q = q.view(b, t, h, dqk).transpose(1, 2)
        k = k.view(b, t, h, dqk).transpose(1, 2)
        vh = v.view(b, t, h, dv).transpose(1, 2)
        attn = (F.silu(q @ k.transpose(-1, -2)) / t).tril()
        o = (attn @ vh).transpose(1, 2).reshape(b, t, h * dv)
        return self.norm(x + self.out(o * F.silu(u)))


class HSTUGenerativeRecommender(nn.Module):
    def __init__(self):
        super().__init__()

        self.embedding_1 = nn.Embedding(200000, 512)
        self.hstuBlock_1 = HSTUBlock(embed_dim=512, num_heads=4, linear_dim=128, attn_dim=128)
        self.hstuBlock_2 = HSTUBlock(embed_dim=512, num_heads=4, linear_dim=128, attn_dim=128)
        self.hstuBlock_3 = HSTUBlock(embed_dim=512, num_heads=4, linear_dim=128, attn_dim=128)
        self.hstuBlock_4 = HSTUBlock(embed_dim=512, num_heads=4, linear_dim=128, attn_dim=128)
        self.rmsNorm_1 = nn.RMSNorm(512)
        self.lmHead_1 = nn.Linear(512, 200000, bias=False)

For agents

Also in Recommendation