Reproducing the HLLM Model in torch-rechub

This document describes the runnable path for the current HLLM (Hierarchical Large Language Model for Recommendation) example in this repository. It implements an "offline precomputed item text embeddings + lightweight User Transformer" variant. It is not ByteDance's official end-to-end HLLM training stack, and its results should not be treated as a reproduction of the paper's metrics.

1. Code Entry Points

• Model: torch_rechub/models/generative/hllm.py
  • HLLMTransformerBlock
  • HLLMModel
• General sequence trainer: torch_rechub/trainers/seq_trainer.py
• MovieLens-1M:
  • examples/generative/data/ml-1m/preprocess_ml_hstu.py
  • examples/generative/data/ml-1m/preprocess_hllm_data.py
  • examples/generative/run_hllm_movielens.py
• Amazon Books:
  • examples/generative/data/amazon-books/preprocess_amazon_books.py
  • examples/generative/data/amazon-books/preprocess_amazon_books_hllm.py
  • examples/generative/run_hllm_amazon_books.py

The examples/generative/data/*/processed/ directories are generated by the preprocessing scripts. The repository does not commit their pkl or embedding files.

2. What the Current Model Does

2.1 Item Side (Offline and Frozen)

The preprocessors support two --model_type values:

Parameter	HuggingFace Model	Hidden Dimension
tinyllama	TinyLlama/TinyLlama-1.1B-Chat-v1.0	2048
baichuan2	baichuan-inc/Baichuan2-7B-Chat	4096

The text format is:

Compress the following sentence into embedding: title: {title}genres: {genres}

Amazon Books uses title and description. The scripts take the hidden state of the last token and write it to item_embeddings_<model_type>.pt, using the token ID as the row index; row 0 is reserved for PAD. HLLMModel checks that the number of rows equals vocab_size and the number of columns equals d_model, then L2-normalizes the matrix and registers it as a non-trainable buffer.

Tokens without text remain zero vectors. The preprocessing log prints the number of covered items. For formal experiments, check this count rather than only checking whether the file exists.

2.2 User Side (Trainable)

The forward path is:

seq_tokens [B, L]
  -> frozen item embedding lookup
  + learnable absolute position embeddings
  + optional time-bucket embeddings
  -> causal Transformer blocks
  -> L2-normalize hidden states
  -> hidden @ normalized_item_embeddings.T / 0.07
  -> logits [B, L, V]

Each block uses pre-norm multi-head self-attention, a feed-forward network, and residual connections. The current attention path creates only a causal mask, with no additional padding attention mask. Left-padded positions still receive position/time embeddings; this is an important boundary to verify against a complete production implementation.

2.3 Training Objective and the Actual Semantics of NCELoss

SeqTrainer trains next-token targets across the whole sequence, rather than only training the final position:

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

PAD label 0 is ignored. --loss_type cross_entropy uses CrossEntropyLoss, while --loss_type nce uses the project's NCELoss.

Important: the current NCELoss applies temperature scaling, log_softmax, and negative log-likelihood for the target class over full-vocabulary logits. It does not sample noise or automatically construct in-batch negatives. Therefore, it does not justify claims of sampled-NCE speedups or additional metric gains. Both HLLM training scripts pass temperature=1.0 to NCELoss because the model output is already scaled by 0.07, avoiding duplicate temperature scaling.

3. Installation and Model Cache

Install from the repository root:

pip install -e ".[generative]"

The generative extra provides transformers and accelerate. If the tokenizer for the selected model reports that SentencePiece is missing, which is common with Baichuan2 environments, also run pip install sentencepiece; this package is not currently included in the extra.

The MovieLens HLLM preprocessing script first checks the HuggingFace cache with local_files_only=True and exits if the target LLM is not cached. Download the corresponding model in a networked environment first, or copy a prepared cache into the runtime environment. --no_download only controls dataset files; it cannot make an uncached LLM available automatically.

4. MovieLens-1M Reproduction Commands

Run all commands below from the repository root:

# 1. Generate sequence data; missing ratings.dat/movies.dat/users.dat files are downloaded automatically
python examples/generative/data/ml-1m/preprocess_ml_hstu.py

# 2. Generate the text mapping and token-ID-aligned item embeddings
python examples/generative/data/ml-1m/preprocess_hllm_data.py \
    --model_type tinyllama \
    --device cuda

# 3. Train and evaluate
mkdir -p outputs/hllm_ml
python examples/generative/run_hllm_movielens.py \
    --model_type tinyllama \
    --epoch 5 \
    --batch_size 64 \
    --learning_rate 1e-3 \
    --weight_decay 1e-5 \
    --max_seq_len 200 \
    --loss_type nce \
    --device cuda \
    --save_dir outputs/hllm_ml \
    --seed 2022

The default data directory is examples/generative/data/ml-1m/processed/, which should contain:

vocab.pkl
train_data.pkl
val_data.pkl
test_data.pkl
movie_text_map.pkl
item_embeddings_tinyllama.pt

MovieLens sequence preprocessing uses per-user leave-last-out: the final interaction is the test target, the penultimate interaction is the validation target, and earlier prefixes generate training examples. It is not a 70/10/20 random user split.

When using a custom directory, both preprocessing scripts must receive the same --output_dir, and training must point --dataset_path to that directory. The training script resolves an explicitly relative --dataset_path in a script-directory-dependent way, so automated jobs should use an absolute path.

5. Amazon Books Reproduction Commands

The sequence data and item metadata must use the same data source. The default bytedance source downloads ByteDance's processed files; raw uses the original Stanford SNAP files.

# 1. Generate sequence data
python examples/generative/data/amazon-books/preprocess_amazon_books.py \
    --data_source bytedance \
    --max_seq_len 200 \
    --min_seq_len 5

# 2. Generate the text mapping and item embeddings
python examples/generative/data/amazon-books/preprocess_amazon_books_hllm.py \
    --data_source bytedance \
    --model_type tinyllama \
    --device cuda

# 3. Train and evaluate
python examples/generative/run_hllm_amazon_books.py \
    --data_dir examples/generative/data/amazon-books/processed \
    --model_type tinyllama \
    --batch_size 64 \
    --epochs 5 \
    --learning_rate 1e-3 \
    --n_layers 2 \
    --dropout 0.1 \
    --max_seq_len 200 \
    --loss_type nce \
    --device cuda

Amazon preprocessing produces:

vocab.pkl
train_data.pkl
val_data.pkl
test_data.pkl
item_text_map.pkl
item_embeddings_tinyllama.pt

Rerunning preprocessing rewrites the mappings, splits, and embedding files in the output directory. --overwrite only describes whether downloaded files are replaced. To preserve existing experiment artifacts, back them up first or use a different --output_dir.

6. Resource and Evaluation Notes

• The HLLM forward pass creates full-vocabulary logits of shape [B, L, V]. Amazon Books has a large vocabulary, so memory usage is usually much higher than "storing item embeddings only." If an OOM occurs, reduce --batch_size and --max_seq_len first.
• TinyLlama/Baichuan2 runs only during offline embedding generation. Training uses precomputed embeddings, but the 2048/4096-dimensional User Transformer is still large.
• The training scripts report the full-sequence loss and held-out top-1 accuracy from SeqTrainer.evaluate(), and additionally compute HR/NDCG@10/50/200.
• Unlike the HSTU example, the current HLLM ranking evaluation does not mask PAD or previously seen items. Before comparing paper or online metrics, align the candidate set and filtering protocol.
• Time and memory usage depend on the hardware, vocabulary size, sequence length, and cache state. This document does not provide fixed runtime estimates or percentage improvements that have not been verified by a benchmark script.

7. Boundaries Relative to Official End-to-End HLLM

Capabilities that can be verified in the current repository:

• Item text embeddings are aligned by token ID and frozen.
• A causal User Transformer produces full-vocabulary cosine logits.
• MovieLens-1M and Amazon Books have preprocessing, training, and top-k evaluation examples.
• A single-machine SeqTrainer can select either full-vocabulary CE or the full-vocabulary classification loss currently named NCELoss.

Parts that are not implemented or not aligned:

• End-to-end joint training of the Item LLM and User LLM.
• A component-by-component reproduction of the official large-model architecture, learnable item embedding tokens, and training configuration.
• Sampled NCE or hard negatives.
• DeepSpeed or distributed training.
• A padding attention mask, unified candidate filtering, and experiments that reproduce paper metrics.
• Multi-step autoregressive decoding and a production inference service.

Existing tests cover item embedding row/dimension validation and the range of cosine logits, but do not prove metric equivalence with the official implementation. The accurate positioning is therefore a lightweight research example, not a "97% aligned" or production-ready official reproduction.