torch-rechub

English

中文

English

中文

Appearance

Reproducing the HSTU Model in torch-rechub ​

This document describes the current implementation status of HSTU (Hierarchical Sequential Transduction Units) in torch-rechub. The core layer follows the main path of Eqs. 2-4 in the paper and Meta's reference implementation: the joint UVQK projection is passed through SiLU, rab^{p,t} is added to the attention scores, the output path retains only f_2(Norm(AV) * U), and inter-layer residual connections are placed outside HSTUBlock layers.

1. Module Structure ​

The main HSTU-related modules are:

  • Main model: torch_rechub/models/generative/hstu.py
    • HSTUModel: item token embeddings, position/time input embeddings, HSTUBlock, and output projection.
  • Core layer and block: torch_rechub/basic/layers.py
    • HSTULayer: a single sequential transduction unit implementing HSTU Eqs. 2-4.
    • HSTUBlock: stacks layers using x = x + HSTULayer(x).
  • HSTU/HLLM utilities: torch_rechub/utils/hstu_utils.py
    • RelativeBucketedTimeAndPositionBias: the rab^{p,t} attention-score bias used by HSTU.
    • RelPosBias: a legacy relative-position bias retained for HLLM and compatibility experiments; the current HSTUModel does not use it.
    • VocabMask: masks PAD or invalid items during inference/ranking.
  • Data and training:
    • torch_rechub/utils/data.py: SeqDataset and SequenceDataGenerator.
    • torch_rechub/trainers/seq_trainer.py: SeqTrainer.
    • examples/generative/run_hstu_movielens.py and examples/generative/run_hstu_amazon_books.py: training and ranking-evaluation examples.

Minimal Reproduction Commands ​

Run all commands below from the repository root. The processed/ directory is generated by preprocessing; it is not a built-in sample dataset:

bash
pip install -e .

# Missing MovieLens raw files are downloaded automatically; add --no_download when using existing files
python examples/generative/data/ml-1m/preprocess_ml_hstu.py

# CPU smoke run; formal experiments normally use --device cuda and more epochs
mkdir -p outputs/hstu_ml
python examples/generative/run_hstu_movielens.py \
    --device cpu --epoch 1 --batch_size 16 \
    --save_dir outputs/hstu_ml

For Amazon Books, use:

bash
python examples/generative/data/amazon-books/preprocess_amazon_books.py \
    --data_source bytedance
mkdir -p outputs/hstu_amazon
python examples/generative/run_hstu_amazon_books.py \
    --device cuda --epoch 3 --batch_size 64 \
    --save_dir outputs/hstu_amazon

The preprocessors rewrite vocab.pkl and the three split files in the output directory. To preserve existing artifacts, back them up first or use --output_dir to write to a new directory. Amazon Books training generates full-vocabulary logits of shape [B, L, V]. The default script parameters assume approximately a 24 GB GPU; if memory is insufficient, reduce --batch_size and --max_seq_len first.

2. HSTULayer: Eqs. 2-4 ​

2.1 Eq. 2: Joint UVQK Projection ​

HSTULayer first applies LayerNorm to the input, then uses one linear layer to generate Q/K/U/V together:

python
proj_out = F.silu(self.proj1(self.norm_in(x)))

Here, SiLU is applied to the complete UVQK projection before it is split. All four paths, U, V, Q, and K, therefore pass through the same nonlinearity, instead of activating only the gate path.

2.2 Eq. 3: Adding rab^{p,t} to Attention Scores ​

The current attention path is:

python
scores = (Q @ K.transpose(-2, -1)) * (1.0 / sqrt(dqk))
scores = scores + rab(time_diffs, seq_len)
attn_weights = silu(scores) / max_seq_len
AV = attn_weights @ V

Key points:

  • rab^{p,t} is a learnable per-head bias derived by bucketing relative position and relative time differences.
  • The bias is added to the attention scores, not to the input token embeddings.
  • HSTU uses silu(scores) / N, rather than standard Transformer softmax.
  • The causal mask prevents each position from attending to future tokens; the padding mask prevents PAD positions from being used as keys.

RelativeBucketedTimeAndPositionBias expects time_diffs to mean query_time - timestamp[i]. When pairwise differences are calculated, the anchor cancels out and yields the relative time differences between events.

2.3 Eq. 4: Gated Output and Projection ​

The current output path is:

python
gated = LayerNorm(AV) * U
output = f2(gated)

Notes:

  • U has already passed through SiLU in the joint projection from Eq. 2, so SiLU(U) is not applied again here.
  • proj2 receives only the gated attention output.
  • [U, x, gated] is no longer concatenated, and there is no separate position-wise FFN.

2.4 External Residual Connections ​

HSTUBlock applies residual connections outside each layer:

python
for layer in self.layers:
    x = x + layer(x, padding_mask=padding_mask, time_diffs=time_diffs)

This avoids requiring proj2 to learn an identity mapping and matches the inter-layer residual form in the HSTU paper/reference implementation.

3. rab^{p,t} and Time Features ​

3.1 RelativeBucketedTimeAndPositionBias ​

RelativeBucketedTimeAndPositionBias maintains two sets of learnable parameters:

  • pos_w: per-head bias for the relative position difference i - j, with a table size of 2 * max_seq_len - 1.
  • ts_w: per-head bias for buckets of relative time differences, with a table size of num_time_buckets + 1.

Time differences are bucketed as follows:

text
dt = abs(time_diffs[i] - time_diffs[j]) / 60.0
bucket = sqrt(dt) or log(dt)
bucket = clamp(bucket / time_bucket_divisor, 0, num_time_buckets)

When time_diffs=None, rab falls back to a position-only bias with shape (1, H, L, L). When time_diffs is supplied, the output shape is (B, H, L, L).

3.2 Role of RelPosBias ​

RelPosBias remains in hstu_utils.py, but it is a legacy relative-position bias:

  • HLLM still reuses it.
  • It can be used independently for compatibility with older experiments.
  • The current HSTU main path no longer uses it. Relative position/time modeling on HSTU attention scores is handled by RelativeBucketedTimeAndPositionBias.

4. HSTUModel Wrapper ​

HSTUModel converts item token sequences into hidden states, invokes HSTUBlock, and outputs vocabulary logits.

The current wrapper contains:

  • token_embedding(vocab_size, d_model, padding_idx=0).
  • Absolute position_embedding(max_seq_len, d_model).
  • Optional input-side time_embedding(num_time_buckets, d_model), controlled by use_time_embedding.
  • HSTUBlock(..., num_time_buckets, time_bucket_fn, time_bucket_divisor).
  • A tied-embedding output projection by default: F.linear(hstu_output, token_embedding.weight, output_bias).

Two time/position paths must be distinguished:

  • Relative position/time modeling for core HSTU Eq. 3 occurs on the attention scores through rab^{p,t}.
  • The current HSTUModel wrapper retains absolute position embeddings and optional input-side time embeddings for compatibility with the existing data interface and experiment settings.

PAD token is defined as 0. The model explicitly zeroes PAD rows after input embedding and after HSTU output, preventing position/time embeddings from leaking signals through PAD positions.

5. Data and Training Conventions ​

5.1 Time-Difference Semantics ​

The preprocessing scripts generate seq_time_diffs with the following semantics:

text
time_diffs[i] = query_time - timestamp[i]

Here, query_time is typically the timestamp of the last behavior in the current history. For example:

text
timestamps  = [100, 200, 300, 400]
query_time  = 400
time_diffs  = [300, 200, 100, 0]

This format allows rab to recover the relative time differences between events through pairwise subtraction.

5.2 Dataset Format ​

SequenceDataGenerator uses four-tuples:

text
(seq_tokens, seq_positions, seq_time_diffs, targets)

seq_positions is retained for data-interface compatibility. The current HSTUModel generates position indices internally from the sequence length during forward, and SeqTrainer does not use the batch's seq_positions.

5.3 Training Objective ​

SeqTrainer uses full-sequence next-token cross-entropy:

text
logits[:, i, :] -> seq_tokens[:, i + 1]
logits[:, -1, :] -> targets

PAD token 0 is excluded from the loss through ignore_index=0. During validation/testing, evaluate returns the average loss and top-1 accuracy for the final held-out target. Ranking evaluation in the example scripts additionally computes HR@K and NDCG@K.

6. Alignment and Differences Relative to Meta's Reference ​

Main points of alignment:

  • Eq. 2: the complete UVQK projection passes through SiLU before splitting.
  • Eq. 3: per-head rab^{p,t} is added to attention scores, followed by silu(scores) / N.
  • Eq. 4: f_2(LayerNorm(AV) * U), with no concat-u/x bypass and no additional FFN.
  • Inter-layer residuals: applied outside each layer by HSTUBlock.
  • Time differences: the recommended form is the anchor-delta representation query_time - timestamp[i].

Remaining engineering differences:

  • DLRM, multi-task heads, and complex feature crossing are not included; this implementation focuses on single-task next-item prediction.
  • The HSTUModel wrapper retains absolute position embeddings and optional input-side time embeddings; the core HSTU layer injects relative position/time bias through rab^{p,t}.
  • Multi-step autoregressive decoding and generation interfaces such as beam search are not wrapped.
  • Initialization and some training settings do not target bit-level reproduction.

7. Summary of Recent HSTU Changes ​

  • Corrected Eq. 2: SiLU now applies to the complete UVQK projection instead of only the gate path.
  • Corrected Eq. 3: added RelativeBucketedTimeAndPositionBias and connected rab^{p,t} to the attention scores.
  • Corrected Eq. 4: removed the [U, x, gated] concatenation bypass, and made proj2 project only the gated attention output.
  • Corrected the inter-layer structure: HSTUBlock now uses the external residual form x = x + layer(x).
  • Updated documentation and API notes: clarified that RelPosBias belongs to the legacy/HLLM path, while the HSTU main path uses rab^{p,t}.