amazon-science/muss

amazon-science/muss

Language: Python

License: NOASSERTION

Stars: 1

Forks: 0

Open issues: 2

Created: 2026-06-24T23:54:08Z

Pushed: 2026-06-29T00:39:31Z

Default branch: main

Fork: no

Archived: no

README:

---

Overview

MUSS uses a *multilevel* (partition → select-clusters → select-within) strategy that:

• 📈 improves recommendation precision by up to 4 points,
• ⚡ runs 20–80× faster than MMR,
• 🧠 boosts RAG question-answering accuracy, and
• 🧮 comes with a constant-factor approximation guarantee (and a 2× tighter bound for DGDS).

It supports two applications out of the box:

ApplicationWhat MUSS picks

Retrieval‑Augmented Generation (RAG)Diverse, relevant context passages for an LLM to read Candidate RetrievalA diverse, high-quality subset from a large corpus (e.g. a product catalog)

How It Works

MUSS runs in three stages (src/methods/muss.py):

1. Partition — split the corpus into n_partitions groups (FAISS k-means or random). 2. Select clusters — pick m partitions using MMR on cluster centroids + median quality. 3. Select within + final pass — run MMR in parallel inside each selected partition, union with the top-k highest-quality items, then a final MMR pass over the union.

Parameter reference

| Parameter | Description | |---|---| | k | Number of items to select | | m | Number of partitions selected in stage 2 | | n_partitions | Total partitions created in stage 1 (must be > m) | | lamb | Item-level quality/diversity trade-off (0 = pure diversity, 1 = pure quality) | | lamb_c | Partition-level quality/diversity trade-off | | k_within | Items selected per partition in stage 3 (defaults to k) | | diversity_type | "sum" (average pairwise) or "min" (nearest-neighbour) |

Installation

pip install -r requirements.txt

> The RAG experiments additionally require an AWS account with access to Claude Haiku via Bedrock.

Quick Start

The fastest way to see MUSS in action — no local setup required — is the Colab demo:


It generates synthetic data (N = 100k, d = 50), runs MUSS to pick k = 50 items, and compares quality and diversity against random selection — with a PCA visualization of the result. Prefer to run locally? The same walkthrough lives in [muss_example.ipynb](muss_example.ipynb).

Minimal Python usage:

import numpy as np
from muss.src.methods.muss import muss

x = np.random.rand(100_000, 50) # item embeddings
q = np.random.rand(100_000) # quality scores

selected_idx, diagnostics = muss(x, q, k=50, m=100, n_partitions=200)

Experiments

> [!IMPORTANT] > All experiment scripts use package-style imports (from muss.src.*) and must be run > from the parent directory of `muss/` as modules.

RAG

The RAG runner reads a JSON config, embeds the corpus with HuggingFace, calls the chosen retrieval method, queries Claude Haiku via Bedrock, and evaluates accuracy.

# from the parent of muss/:
python muss/src/generate_settings_rag.py # writes example configs to muss/outputs/
python -m muss.run_experiment_rag -i muss/outputs/ours-000.json -d 1000 -n 50

| Flag | Meaning | |---|---| | -i / --input | config file (required) | | -d / --data_size | cap corpus size (optional) | | -n / --num_questions | cap test-question count (optional) |

Results are written to muss/outputs/results~*.csv and muss/outputs/*~answers.json. LLM responses are cached in muss/outputs/_cache.json to avoid redundant API calls.

Available `method_key` values:

| method_key | Description | |---|---| | muss | Full MUSS algorithm (requires pre_proc_params) | | random | Random retrieval baseline | | dpp | k-DPP baseline |

For muss, pre_proc_params takes n_clusters and algorithm — one of packing_k_means, packing_random, packing_balanced, packing_balanced_fast. MMR is recovered by setting m=1; DGDS by setting algorithm=packing_random and add_max_q_to_union=False.

Candidate Retrieval

# from the parent of muss/:
python -m muss.run_experiment_candidate_retrieval # MUSS, DGDS, MMR, Clustering, DPP, Random
python -m muss.run_experiment_candidate_retrieval_ablation # MUSS ablation variants

Results are written to temp/{dataset}.csv (relative to the working directory you launched from).

The input DataFrame must have columns: emb (embedding vectors), quality_score, label, subcategory_code, ctr. Optional categorical columns cat_a, cat_b, cat_c are used for per-feature entropy reporting if present.

Datasets

| Dataset | Path | Used in | |---|---|---| | StackExchange QA | data/StackExchange/ExamData/llamav2_2023091223/exam.json | RAG | | DevOps QA | data/DevOps/ExamData/html_llamav2_2023091421/exam.json | RAG |

For candidate retrieval, supply your own parquet file with the columns listed above. The Colab demo (or local [muss_example.ipynb](muss_example.ipynb)) demonstrates the algorithm on synthetic data and is a good starting point.

Repository Layout

muss/
├── data/ # RAG eval datasets (JSON)
├── assets/ # figures used in this README
├── outputs/ # example RAG configs + experiment results
├── __init__.py # makes the repo importable as the `muss` package
├── requirements.txt
├── muss_colab_demo.ipynb # one-click Colab demo
├── muss_example.ipynb # quick-start notebook (local)
├── run_experiment_rag.py # RAG experiment runner
├── run_experiment_candidate_retrieval.py
├── run_experiment_candidate_retrieval_ablation.py
├── src/
│ ├── methods/
│ │ ├── muss.py # core MUSS algorithm
│ │ ├── dpp.py, random.py # baselines
│ │ └── objectives.py # global quality+diversity objective (RAG only)
│ ├── candidate_retrieval/
│ │ ├── helpers.py # run_MUSS, evaluation, save_result
│ │ ├── baselines.py # MMR, clustering, DPP, random, DGDS
│ │ └── ablation_muss.py # MUSS ablation variants
│ ├── data_utils/
│ │ ├── read_data.py # StackExchange / DevOps loaders
│ │ └── embeddings.py # cached embedding + FAISS index
│ ├── llm/bedrock_language_model.py # AWS Bedrock wrapper
│ └── generate_settings_rag.py # build RAG configs
├── baselines.md # baseline methods MUSS is compared against
├── LICENSE
└── README.md

See...