Run this notebook

Download it with the View page source link at the top right, or from the repository at docs/source/reproducing-sasrec-results-on-ml1m.ipynb.

Reproducing SASRec Results on ML1M

This notebook reproduces the MovieLens 1M data preparation and evaluation protocols for Compresso RecSys’s modernized SASRec variant. It is not a line-for-line recreation of the original TensorFlow model: the attention block uses contemporary pre-norm transformer conventions, including one normalized stream for queries, keys, and values, a residual around the full attention module, and a standard output projection.

The trained model is evaluated under two protocols:

  1. the paper’s held-out target against 100 sampled negatives at cutoff 10;

  2. Compresso RecSys’s full-catalog protocol at cutoff 20.

Keeping the protocols separate matters: sampled ranking is much easier than ranking the entire catalog, so their numbers are not interchangeable. The paper’s Table III values are included only as historical context for the shared setup, not as a like-for-like architecture comparison. The final section also evaluates validation data and stores the trained model in the dataset checkpoint.

Configuration and run modes

A full run keeps SASRecConfig’s MovieLens 1M hyperparameter defaults, including 201 epochs. Set SMOKE_EPOCHS to a small integer to check the workflow quickly. Smoke runs deliberately hide the published comparison column so a shortened run cannot be mistaken for a full experiment.

[ ]:

from dataclasses import replace
from pathlib import Path
import time

import numpy as np
import torch
from compresso import SRPTensor

import compresso_recsys as cr
from compresso_recsys.evaluation import (
    evaluate_ranked_predictions,
    evaluate_recommender,
)
from compresso_recsys.metrics import MRR, NDCG, HitRate, Recall
from compresso_recsys.models import (
    ItemTokenizer,
    SASRecConfig,
    SASRecTrainer,
    SequenceBatcher,
)
from compresso_recsys.sequences import ItemSequences

PUBLISHED_HR10 = 0.8245
PUBLISHED_NDCG10 = 0.5905
N_NEGATIVES = 100
SAMPLED_CUTOFF = 10
NEGATIVE_SEED = 0
FULL_CUTOFF = 20
CHECKPOINT = Path("artifacts/ml1m/sasrec_ml1m.zip")

if torch.cuda.is_available():
    DEVICE = "cuda"
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
    DEVICE = "mps"
else:
    DEVICE = "cpu"

SMOKE_EPOCHS: int | None = None  # Set to 1 for a quick workflow check.
cfg = replace(SASRecConfig(), device=DEVICE)
if SMOKE_EPOCHS is not None:
    cfg = replace(cfg, epochs=SMOKE_EPOCHS)

started = time.time()

def elapsed() -> str:
    return f"[{time.time() - started:7.1f}s]"

print(
    f"device: {DEVICE}   epochs: {cfg.epochs}   "
    f"window: {cfg.max_history_length}"
)
print("model: Compresso SASRec (modernized pre-norm attention block)")
if SMOKE_EPOCHS is not None:
    print("SMOKE_EPOCHS is set: this is a smoke run, not a full experiment")

Build the MovieLens 1M checkpoint

The paper uses chronological leave-last-out evaluation and treats every rating as implicit feedback. min_value_to_keep=1.0 therefore keeps all ratings. SASRec reads item IDs rather than text, so entity-text filtering is disabled as well.

[ ]:

print(f"{elapsed()} checkpoint")
if CHECKPOINT.exists():
    print(f"reusing {CHECKPOINT}")
else:
    print(f"building {CHECKPOINT} from MovieLens 1M")
    cr.build_recsys_checkpoint(
        dataset="ml1m",
        data_dir="data",
        checkpoint_path=str(CHECKPOINT),
        split_mode="leave_last_out",
        min_user_support=5,
        item_min_support=5,
        min_value_to_keep=1.0,
        min_entity_text_words=0,
        seed=0,
        show_progress=False,
    )

Inspect the split

Leave-last-out reserves each user’s last interaction for test and the second-to-last for validation. The paper’s dataset statistics match this prepared training split, with those two held-out interactions excluded.

[ ]:

with cr.read_checkpoint(CHECKPOINT) as root:
    split = cr.load_recsys_split(root)

train_sequences = split["x_train_sequences"]
test_source = split["test_source_sequences"]
test_targets = split["test_target_matrix"]
item_ids = split["train_item_ids"]

lengths = np.diff(train_sequences.indptr)
users = train_sequences.n_rows
items = train_sequences.n_items
actions = int(train_sequences.values.size)
print(f"{users} users, {items} items, {actions} actions")
print(
    f"{actions / users:.1f} actions/user, "
    f"{actions / items:.1f} actions/item, "
    f"median history {int(np.median(lengths))}"
)
print("paper Table 2: 6040 users, 3416 items, 163.5 and 289.1")

Train SASRec

The tokenizer retains MovieLens item IDs for stable-ID recommendation. The batcher leaves its window unset so SASRecConfig.max_history_length remains the single source of truth.

[ ]:

print(f"{elapsed()} train")
batcher = SequenceBatcher(
    ItemTokenizer(train_sequences.n_items, item_ids=item_ids),
)
trainer = SASRecTrainer(cfg, batcher).fit(
    train_sequences,
    item_ids=item_ids,
)

every = max(1, len(trainer.history) // 10)
print(f"{'epoch':>6}  {'loss':>9}")
for index, entry in enumerate(trainer.history):
    if index % every == 0 or index == len(trainer.history) - 1:
        print(f"{int(entry['epoch']):>6}  {entry['loss']:>9.4f}")

Score the catalog

The sampled protocol needs scores for specific items, including candidates outside the model’s top-k result. Compute the full score matrix once in batches and reuse it below.

[ ]:

def catalog_scores(
    model: SASRecTrainer,
    source: ItemSequences,
    *,
    rows: int = 512,
) -> np.ndarray:
    out = np.empty((source.n_rows, source.n_items), dtype=np.float32)
    model.model.eval()
    for start in range(0, source.n_rows, rows):
        stop = min(start + rows, source.n_rows)
        chunk = ItemSequences(
            values=source.values[source.indptr[start] : source.indptr[stop]],
            indptr=source.indptr[start : stop + 1] - source.indptr[start],
            n_items=source.n_items,
        )
        with torch.no_grad():
            tokens, mask = model.batcher.encode(chunk, device=model.device)
            final = model.batcher.gather_final(model.model(tokens), mask)
            out[start:stop] = model.model.score(final).float().cpu().numpy()
    return out

print(f"{elapsed()} scoring {test_source.n_rows} test users")
scores = catalog_scores(trainer, test_source)

Published protocol: 100 sampled negatives at 10

For each user, rank the held-out item against 100 items that user never interacted with. Ties go to the negative; otherwise an untrained model that assigns every item the same score would report a perfect hit rate. Matching the evaluation protocol makes the published values useful context, but the modernized attention block means this is not an exact architecture reproduction.

[ ]:

def sampled_protocol(
    scores: np.ndarray,
    interacted: np.ndarray,
    targets,
    *,
    n_negatives: int,
    cutoff: int,
    seed: int,
) -> tuple[float, float, int]:
    rng = np.random.default_rng(seed)
    hits: list[float] = []
    gains: list[float] = []
    for row in range(scores.shape[0]):
        positives = targets[row].indices
        if positives.size == 0:
            continue
        pool = np.flatnonzero(~interacted[row])
        if pool.size < n_negatives:
            continue
        negatives = rng.choice(pool, size=n_negatives, replace=False)
        target = int(positives[0])
        above = int((scores[row, negatives] >= scores[row, target]).sum())
        hits.append(float(above < cutoff))
        gains.append(1.0 / np.log2(above + 2) if above < cutoff else 0.0)
    return float(np.mean(hits)), float(np.mean(gains)), len(hits)

interacted = (
    split["test_source_matrix"] + test_targets
).astype(bool).toarray()
hr, ndcg, scored = sampled_protocol(
    scores,
    interacted,
    test_targets,
    n_negatives=N_NEGATIVES,
    cutoff=SAMPLED_CUTOFF,
    seed=NEGATIVE_SEED,
)

print(
    f"{elapsed()} published sampled protocol -- "
    f"{N_NEGATIVES} sampled negatives, @{SAMPLED_CUTOFF}"
)
print(f"{'metric':<12}{'this run':>10}{'published*':>12}{'delta':>10}")
if SMOKE_EPOCHS is not None:
    print(f"{'HR@10':<12}{hr:>10.4f}{'--':>12}{'--':>10}")
    print(f"{'NDCG@10':<12}{ndcg:>10.4f}{'--':>12}{'--':>10}")
    print("published column withheld: SMOKE_EPOCHS shortened this run")
else:
    print(
        f"{'HR@10':<12}{hr:>10.4f}{PUBLISHED_HR10:>12.4f}"
        f"{hr - PUBLISHED_HR10:>+10.4f}"
    )
    print(
        f"{'NDCG@10':<12}{ndcg:>10.4f}{PUBLISHED_NDCG10:>12.4f}"
        f"{ndcg - PUBLISHED_NDCG10:>+10.4f}"
    )
print("* Original SASRec Table III; context only, not architecture parity.")
print(f"scored {scored} of {test_source.n_rows} users")

Package protocol: full catalog at 20

Now rank every unseen catalog item and compare SASRec with a most-popular baseline evaluated identically. This is the protocol used by the package’s benchmark tables.

[ ]:

metrics = [
    NDCG(FULL_CUTOFF),
    Recall(FULL_CUTOFF),
    HitRate(FULL_CUTOFF),
    MRR(FULL_CUTOFF),
]
result = evaluate_recommender(
    trainer,
    source=test_source,
    targets=test_targets,
    metrics=metrics,
    sample_ids=split["test_eval_user_ids"],
)

popularity = np.asarray(split["x_train"].sum(axis=0)).ravel()
by_popularity = np.argsort(-popularity, kind="stable")
seen = split["test_source_matrix"].astype(bool).toarray()[:, by_popularity]
unseen_first = np.argsort(seen, axis=1, kind="stable")[:, :FULL_CUTOFF]
baseline = evaluate_ranked_predictions(
    predictions=SRPTensor(
        cols=torch.from_numpy(by_popularity[unseen_first]).long(),
        vals=torch.arange(
            FULL_CUTOFF,
            0,
            -1,
            dtype=torch.float32,
        ).expand(test_source.n_rows, FULL_CUTOFF),
        shape=(test_source.n_rows, train_sequences.n_items),
    ),
    targets=test_targets,
    metrics=metrics,
    sample_ids=split["test_eval_user_ids"],
)

print(f"{elapsed()} benchmark protocol -- full catalog, @{FULL_CUTOFF}")
print(f"{'metric':<16}{'SASRec':>10}{'popular':>10}{'lift':>9}")
for name in result.metrics:
    ours, theirs = result[name], baseline[name]
    lift = f"{ours / theirs:.1f}x" if theirs else "--"
    print(f"{name:<16}{ours:>10.4f}{theirs:>10.4f}{lift:>9}")
print(f"scored {result.n_scored_rows} of {result.n_rows} users")

Check the epoch budget on validation data

SASRec uses a fixed epoch budget. Validation and test nDCG tracking each other suggests that budget is reasonable; validation far above test is a warning that the run overfit.

[ ]:

validation = evaluate_recommender(
    trainer,
    source=split["val_source_sequences"],
    targets=split["val_target_matrix"],
    metrics=[NDCG(FULL_CUTOFF)],
    sample_ids=split["val_eval_user_ids"],
)
val_ndcg = validation[f"ndcg@{FULL_CUTOFF}"]
test_ndcg = result[f"ndcg@{FULL_CUTOFF}"]
print(
    f"validation ndcg@{FULL_CUTOFF}: {val_ndcg:.4f}   "
    f"test ndcg@{FULL_CUTOFF}: {test_ndcg:.4f}"
)

Save the trained model

Store the fitted recommender under models/sasrec.zip inside the same data checkpoint. The model entry carries its configuration, tokenizer, stable item IDs, training history, and learned Torch state.

[ ]:

trainer.save_to_checkpoint(CHECKPOINT, "sasrec")
print(f"{elapsed()} saved into {CHECKPOINT.name} as 'sasrec'")