UBC-SLIME/splade-lexicon-large-c8000

SPLADE-Lexicon, ModernBERT-large, 8000 clusters

Final saved checkpoint for this run of the lexicon model accompanying Learned Sparse Lexical Representations for Large-scale Authorship Attribution.

Use the code in lingjzhu/neural_lexical. The output has one vector per document with 8000 dimensions. Token cluster assignments are learned during training and included in cluster_ids.pt. Load them along with the weights; do not run k-means again for inference. This checkpoint uses ReLU, fused mean pooling and one post-pooling log1p operation, matching its training implementation. Padding is excluded from pooling, as in the shared release code. The PyTorch fallback preserves this activation ordering; set module.use_triton=False for CPU inference.

Usage

Clone the code repository and install its dependencies, then run from its root:

from src.models.ClusteredTransformer import ClusteredTransformer
from src.models.SparseEncoder import SparseEncoder
import torch

module = ClusteredTransformer.load("UBC-SLIME/splade-lexicon-large-c8000")
model = SparseEncoder(modules=[module]).to("cuda").eval()
features = model.tokenize(["A sample document.", "Another document."])
features = {k: v.to("cuda") for k, v in features.items()}
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    embeddings = model(features)["sparse_embeddings"]
    # Optional: retain the top 1,024 dimensions for the paper's evaluation setting.
    values, indices = embeddings.topk(1024, dim=-1)
    embeddings = torch.zeros_like(embeddings).scatter(-1, indices, values)
    scores = embeddings @ embeddings.T

Checkpoint provenance

• Source: exp2/outputs_clustered_modernbert_large_8000_relu_reg5e-3/checkpoint-1582
• Training step: 1582
• Epoch: 2.0
• Backbone: answerdotai/ModernBERT-large
• Maximum sequence length: 512
• FLOPS regularizer weight: 5e-3 for queries and documents
• Use the top 1,024 dimensions per document for the paper's SPLADE-Lexicon evaluation.

Weights and cluster assignments are preserved from the source checkpoint. Packaging adds portable module configuration and removes training logs and training-data examples. Full HRS evaluation has not been rerun for this release; no new benchmark results are claimed.