Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/facebookresearch/LoRe/llms.txt

Use this file to discover all available pages before exploring further.

Before training a reward basis, raw preference data must be converted into feature difference tensors — vectors of the form embedding(chosen) - embedding(rejected). This page documents LoRe’s data preparation helpers: simulation utilities for synthetic populations built on multi-reward datasets (PersonalLLM), sparse sampling for data-limited regimes, and PRISM-specific loaders that parse embedding dictionaries into the list-of-tensors format expected by all training functions.

​
simulate_user()

Generates preference feature differences for a single synthetic user given their reward weight vector. 
from utils import simulate_user

feature_diff = simulate_user(
    reward_tensor=reward_tensor[0],  # shape [num_responses, num_reward_models]
    features=features[0],            # list of num_responses embedding tensors
    w=W[0],                          # shape [num_reward_models]
)
# feature_diff: Tensor [num_prompts, F]
For each prompt, the function scores all responses under the user’s personalized reward w · scores, then returns the difference between the highest-scoring (chosen) and lowest-scoring (rejected) response embeddings.

​
Parameters

​
reward_tensor
ndarray
required
Reward scores for all responses to a single prompt, shape [num_responses, num_reward_models]. Typically a row of the 3D reward tensor indexed by prompt.
​
features
list[Tensor]
required
Embedding tensors for all responses per prompt. features[i] is a list (or indexable) of num_responses embedding tensors, each shape [F].
​
w
ndarray
required
User weight vector over reward models, shape [num_reward_models]. Typically drawn from a Dirichlet distribution via generate_popupulation().

​
Return value

​
feature_diff
torch.Tensor
Stack of chosen-minus-rejected embedding differences, shape [num_prompts, F].

​
simulate_population()

Calls simulate_user() for every user in W and stacks the results into a 3D tensor. 
from utils import simulate_population

all_feature_diff = simulate_population(
    reward_tensor=reward_tensor,   # shape [num_prompts, num_responses, num_models]
    features=features,             # list of num_prompts × list of num_responses tensors
    W=W,                           # ndarray [N, num_reward_models]
)
# all_feature_diff: Tensor [N, num_prompts, F]

​
Parameters

​
reward_tensor
ndarray
required
3D array of reward scores, shape [num_prompts, num_responses, num_reward_models].
​
features
list
required
Nested list of embeddings. features[prompt_idx][response_idx] is a tensor of shape [F].
​
W
ndarray
required
Population weight matrix, shape [N, num_reward_models]. Each row is one user’s simplex vector over reward models.

​
Return value

​
all_feature_diff
torch.Tensor
Stacked feature differences, shape [N, num_prompts, F]. Pass directly to create_sparse_tensor() to produce the list format expected by training functions.

​
Full PersonalLLM workflow

import numpy as np
from utils import generate_popupulation, simulate_population, create_sparse_tensor

# 1. Sample a synthetic user population
alpha = 0.1 * np.ones(10)   # 10 reward models
N = 1000
W = generate_popupulation(alpha, N)

# 2. Build feature difference tensors for all users × all prompts
all_feature_diff = simulate_population(reward_tensor, features, W)
# shape: [1000, num_prompts, 4096]

# 3. Sub-sample prompts to simulate limited interaction data
train_features = create_sparse_tensor(all_feature_diff, sample_percentage=0.005)
# train_features: list of 1000 tensors, each [~num_prompts * 0.005, 4096]

​
generate_popupulation()

Samples a population of user weight vectors from a Dirichlet distribution.
The function name contains a typo in the source: two p’s (generate_popupulation). Import and call it with the double-p spelling to match the source.
from utils import generate_popupulation

alpha = 0.1 * np.ones(10)   # concentration parameter
W = generate_popupulation(alpha, N=1000)
# W: ndarray [1000, 10], each row sums to 1

​
Parameters

​
alpha
ndarray
required
Dirichlet concentration parameter array, shape [num_reward_models]. Smaller values (e.g., 0.1) produce more peaked / specialized users; larger values (e.g., 1.0) produce more uniform users. In PersonalLLM experiments, alpha = 0.1 * np.ones(10) is used.
​
N
int
required
Number of users to sample.

​
Return value

​
W
ndarray
Shape [N, len(alpha)]. Each row is a probability vector over reward models (sums to 1).

​
create_sparse_tensor()

Randomly sub-samples a fraction of prompts from each user’s feature matrix, simulating limited interaction data. 
from utils import create_sparse_tensor

train_features = create_sparse_tensor(
    dense_tensor=all_feature_diff,   # Tensor [N, M, F]
    sample_percentage=0.005,
)
# train_features: list of N tensors, each [~0.005*M, F]

​
Parameters

​
dense_tensor
torch.Tensor
required
Full feature difference array, shape [N, M, F] where N is users, M is prompts, F is feature dimension. Typically the output of simulate_population().
​
sample_percentage
float
required
Fraction of prompts to retain per user, in (0, 1]. The number of sampled rows per user is floor(sample_percentage * M). Sampling is without replacement.

​
Return value

​
sparse_rows
list[Tensor]
List of N tensors, each shape [floor(sample_percentage * M), F], moved to device. Row order within each user is random.

​
Common sampling fractions in LoRe experiments

Splitsample_percentageInterpretation
Train (seen users)0.005~0.5% of prompts — highly limited
Few-shot (unseen users)0.001~0.1% of prompts — minimal signal
Test (generalization)1.0All prompts

​
create_dataset_prism()

Converts a PRISM embedding dictionary into the list-of-tensors format expected by LoRe training functions.
The current PRISM train_basis.py uses group_embeddings_by_user() (defined in that file) rather than create_dataset_prism(). create_dataset_prism() is available in utils.py for use with an alternative embedding format: a nested dict {user_id: {dialog_id: {"chosen": [...], "rejected": [...]}}}.
from utils import create_dataset_prism

train_features = create_dataset_prism(embeddings)
# train_features: list of N_users tensors, each [num_pairs_for_user, F]

​
Parameters

​
embeddings
dict
required
Nested dictionary with structure {user_id: {dialog_id: {"chosen": [...], "rejected": [...]}}}. The innermost lists hold embedding tensors indexed by response position.

​
Return value

​
sparse_rows
list[Tensor]
List of N_users tensors, each shape [total_pairs_for_user, F], on device. Each row is chosen[i] - rejected[i] for one preference pair across all dialogs of that user.

​
create_dataset_prism_shots()

Same as create_dataset_prism() but randomly samples a fixed number of dialogs per user instead of using all available dialogs. 
from utils import create_dataset_prism_shots

few_shot_features = create_dataset_prism_shots(embeddings, shots=5)
# few_shot_features: list of N_users tensors, each [pairs_in_sampled_dialogs, F]

​
Parameters

​
embeddings
dict
required
Same nested dictionary format as create_dataset_prism().
​
shots
int
required
Number of dialogs to randomly sample per user (without replacement from that user’s dialog list). Each selected dialog contributes all of its preference pairs.

​
Return value

​
sparse_rows
list[Tensor]
List of N_users tensors on device. Each tensor contains the feature differences from the sampled dialogs only. The number of rows varies by user depending on how many preference pairs fall within the sampled dialogs.
shots here refers to the number of dialogs sampled, not individual preference pairs. If a dialog contains multiple (chosen, rejected) pairs, all of them are included. Use sample_shots() instead if you need to control the exact number of preference pairs per user.
Like create_dataset_prism(), this function expects the nested {user_id: {dialog_id: ...}} dict format. The current PRISM pipeline stores embeddings as a flat list of dataset entries; use group_embeddings_by_user() from PRISM/train_basis.py to work with that format.

Build docs developers (and LLMs) love

Get started for freeTalk to us