Preference and RL Alignment

How to align SameDiff language models with preference pairs, rewards, DPO, GRPO, PPO, ORPO, KTO, reward models, and RLAlignmentPipeline.

Preference and RL alignment are the post-SFT steps used to make a model prefer better answers, follow reward signals, or satisfy task-specific behavior constraints. SameDiff exposes both direct trainer classes and a higher-level RLAlignmentPipeline.

The main runnable example is RLAlignmentTrainingExample.java. It demonstrates DPO, GRPO, PPO, KTO, ORPO, SimPO, DAPO, DrGRPO, GSPO, reward-model training, RLPipelineConfig, and RLAlignmentPipeline.

Where Alignment Fits

A typical LLM training stack is:

1. Continued pretraining or domain-adaptive pretraining.

2. Instruction SFT with response masks.

3. Preference or reward alignment.

4. Evaluation and serving.

Preference methods assume the policy already produces usable responses. Start with Instruction SFT, LoRA, and QLoRA before applying DPO, GRPO, PPO, or related algorithms.

Choose an Alignment Method

Method	Data	Reference model	Use when
DPO	Chosen/rejected response pairs	Yes	You have offline preference pairs and want a stable default
IPO / RDPO	Chosen/rejected pairs	Yes	You need DPO variants for overfit resistance or noisy labels
ORPO	Chosen/rejected pairs	No	You want preference optimization without a reference model
SimPO	Chosen/rejected pairs	No	You want a simpler length-normalized preference objective
KTO	Desirable/undesirable labels	Yes	You have unpaired feedback rather than paired comparisons
GRPO	Prompt groups and rewards	Usually	You want group-relative rewards, useful for reasoning-style tasks
DrGRPO	Prompt groups and rewards	Usually	You need GRPO with reduced length bias
DAPO	Prompt groups and rewards	Usually	You want asymmetric clipping and token-level KL control
GSPO	Prompt groups and rewards	Usually	You want group-stable policy optimization
PPO	Rollouts, rewards, values	Yes	You need classic online RLHF machinery
Reward model	Preference pairs	No	You need a learned reward function for online or batch RL

DPO Trainer

DPO is the simplest default when you have chosen/rejected pairs. The trainer compares policy and reference log probabilities and pushes the policy toward the chosen response.

DPOConfig dpoConfig = DPOConfig.standard(
    "policy_logits",
    "chosen",
    "rejected");

dpoConfig.setLogits2D(true); // useful for toy MLP examples

DPOTrainer trainer = new DPOTrainer(policyModel, referenceModel, dpoConfig);

double loss = trainer.trainStep(inputs);

For a real LLM, chosen and rejected should be tokenized response sequences with the same prompt context. Keep prompt formatting identical to SFT and generation so the policy and reference compare the same conditional distribution.

RLAlignmentPipeline

RLAlignmentPipeline wraps the trainer and the outer loop. The example uses a DPO config and a pipeline config:

DPOConfig dpoConfig = DPOConfig.standard(
    "policy_logits",
    "chosen",
    "rejected");

dpoConfig.setLogits2D(true);

RLPipelineConfig pipelineConfig = RLPipelineConfig.builder()
    .numEpochs(1)
    .learningRate(5e-7)
    .logEveryNSteps(1)
    .build();

RLAlignmentPipeline pipeline = RLAlignmentPipeline.create(
    policyModel,
    referenceModel,
    dpoConfig,
    pipelineConfig);

Preference pairs are plain objects at the pipeline boundary:

List<PreferencePair> pairs = new ArrayList<>();

pairs.add(PreferencePair.builder()
    .prompt("Explain backpropagation.")
    .chosen("Backpropagation computes gradients by applying the chain rule through the network.")
    .rejected("It is a type of neural network.")
    .build());

TrainingResult result = pipeline.trainDPO(pairs);
SameDiff trainedModel = pipeline.getTrainedModel();

The pipeline entry points are:

RLAlignmentPipeline.create(policy, rlConfig, pipelineConfig);
RLAlignmentPipeline.create(policy, reference, rlConfig, pipelineConfig);
RLAlignmentPipeline.create(policy, reference, rlConfig, pipelineConfig, sampler, rewardFn);

pipeline.trainDPO(preferencePairs);
pipeline.trainOnline(prompts, rewardFunction);
pipeline.train(datasetIterator, maxSteps);

Use the no-reference overload only for methods that do not require a reference model, such as ORPO-style objectives. Use the sampler/reward overload for online RL where responses are generated during training.

Reward Models

Reward-model training learns a scalar scoring function from preference pairs. The example covers RewardModelConfig and RewardModelTrainer with a Bradley-Terry objective.

Use a reward model when:

• You need online RL with generated rollouts.

• Preferences are expensive and you want a reusable reward approximation.

• You want to score responses before running PPO or GRPO.

For direct preference optimization on static pairs, DPO or ORPO is usually simpler.

GRPO and Online Rewards

GRPO-style training uses groups of responses for the same prompt and normalizes rewards inside the group. That makes it useful when absolute reward scale is noisy but relative quality within each prompt is meaningful.

Important settings to track:

Setting	Meaning
groupSize	Number of sampled responses per prompt
clipEpsilon	Policy update clipping threshold
klPenalty	Penalty for drifting from the reference model
reward function	Task score, model score, or external evaluator

For online methods, the generation stack and training stack meet. The sampler produces candidate responses, the reward function scores them, and the trainer updates the policy.

PEFT for Alignment

Alignment often uses adapters rather than full-model updates. The practical recipe is:

1. Load the SFT model or base model.

2. Wrap it with PeftModel using LoRA or QLoRA.

3. Use the PEFT-wrapped policy as the trainable policy.

4. Keep the reference model frozen.

5. Train DPO, GRPO, PPO, or ORPO.

6. Save the adapter separately or merge for deployment.

This keeps alignment checkpoints small and makes it easy to compare multiple alignment runs over the same SFT base.

Data Preparation

Preference data must preserve the original prompt context. Each record should include:

Field	Purpose
prompt	The full user/system context
chosen	Preferred response
rejected	Less preferred response
optional score	Reward score, label, or metadata
optional domain	Used for stratified splits and weighted mixing

Keep the same tokenizer and chat template used in SFT. If the SFT model trained on ChatML but DPO compares untemplated strings, the policy sees a different distribution than it learned.

Evaluation Checklist

• Track held-out preference accuracy: chosen response should score above rejected.

• Track KL drift from the reference model.

• Compare generation quality before and after alignment on fixed prompts.

• Watch response length, refusal rate, and format adherence.

• For online rewards, validate reward hacking with human or task-specific checks.

• For adapter alignment, verify base weights remain frozen.

Run the Example

cd ../deeplearning4j-examples/samediff-examples
mvn exec:java -Dexec.mainClass="org.nd4j.examples.samediff.quickstart.training.RLAlignmentTrainingExample"

Related Pages

• LLM Fine-Tuning Overview

• Instruction SFT, LoRA, and QLoRA

• Knowledge Distillation

• Quantized and Memory-Efficient Training

• PEFT and RL Alignment Training