Preference Packing: Efficient Preference Optimization for Large Language Models
| Authors | Jaekyung Cho |
| Year | 2026 |
| Field | NLP |
| arXiv | 2602.24082 |
| Download | |
| Categories | cs.CL, cs.AI |
Abstract
Resource-efficient training optimization techniques are becoming increasingly important as the size of large language models (LLMs) continues to grow. In particular, batch packing is commonly used in pre-training and supervised fine-tuning to achieve resource-efficient training. We propose preference packing, a method to enhance resource efficiency in training techniques that use data with different responses for the same input prompt, such as reward models or Direct Preference Optimization (DPO). Preference packing improves resource efficiency by reducing the attention operations for duplicate input prompts and decreasing KV cache memory usage. We conducted experiments on text-only datasets and image-included datasets and achieved at least 37% reduction in training time. Notably, this method can be applied alongside existing optimization techniques such as batch sorting, resulting in a 3.22x speedup.
Engineering Breakdown
Plain English
This paper introduces preference packing, a technique to make preference optimization training (like DPO and reward model training) significantly more efficient by reducing redundant computation. The key insight is that when training on preference data, the same input prompt appears multiple times paired with different responses—preference packing bundles these duplicate prompts into a single sequence, eliminating redundant attention computation and reducing KV cache memory usage. The authors achieved at least 37% reduction in training time across text-only and image-included datasets, making it a practical optimization for scaling preference-based fine-tuning of large language models.
Core Technical Contribution
The core innovation is identifying and exploiting a specific inefficiency in preference optimization workflows: duplicate input prompts appear multiple times during training when using comparative ranking losses (DPO, reward modeling, etc.), and standard batch packing doesn't handle this case. Preference packing extends the batch packing paradigm by intelligently grouping multiple responses to the same prompt into a single sequence, allowing the model to compute the input prompt's representation once and reuse it for all response branches. This is fundamentally different from standard batch packing used in pre-training and SFT, which doesn't need to handle this prompt-duplication pattern since each training example has a unique prompt. The technique reduces both forward-pass FLOPs and GPU memory consumption without changing model weights or final model behavior.
How It Works
Preference packing works by restructuring how preference training data is batched and processed through the model. During normal DPO or reward model training, you have triplets of (prompt, preferred_response, rejected_response) or (prompt, response_A, response_B). Instead of running the prompt through the attention mechanism multiple times—once for each response variant—preference packing concatenates all responses to the same prompt into a single sequence: [prompt] [response_A] [response_B]. This allows the transformer to compute attention over the prompt only once, and the different response branches reuse the prompt's computed key-value cache entries. The loss computation still operates on the individual response embeddings (pooling the final token or taking specific positions), but the underlying attention operations and KV cache storage are shared. The packing algorithm needs to track prompt boundaries and ensure gradient computation correctly attributes loss to the right response variants while maintaining numerical stability and proper attention masking across the concatenated responses.
Production Impact
For production teams training preference-optimized models (RLHF, DPO, reward models), preference packing translates directly into lower GPU memory requirements and faster training wall-clock time—37% speedup means you can either train more models in the same time budget, use smaller GPUs, or reduce cloud costs substantially. This is particularly valuable when fine-tuning billion+ parameter models where memory pressure is the primary bottleneck. The implementation integrates into existing DPO and reward modeling pipelines without requiring architectural changes; it's a data batching and attention optimization applied during training only, with zero impact on inference. Trade-offs include slightly increased implementation complexity in data loading/batching code, potential sensitivity to sequence length limits if responses are long, and the requirement that your training framework supports flexible attention masking (most modern ones do—PyTorch, JAX, etc.). For smaller models or when GPU memory isn't a constraint, the benefit is minimal, but for large-scale alignment tuning (the primary use case for DPO), this is a straightforward win.
Limitations and When Not to Use This
Preference packing assumes that multiple responses to the same prompt will fit together in a single sequence without exceeding context length limits; if responses are very long or you have many variants per prompt, this breaks down and you fall back to standard batching. The paper doesn't thoroughly analyze numerical stability when computing attention across long concatenated sequences or how gradient flow changes when sharing KV cache entries—there may be subtle optimization dynamics that differ from standard training. The approach is most beneficial for preference optimization workflows; it has limited or no benefit for standard next-token prediction training (pre-training, SFT) where each example already has a unique prompt. The paper also doesn't discuss how preference packing interacts with other efficiency techniques like LoRA, quantization, or gradient accumulation, which is important for practitioners trying to combine multiple optimizations. Finally, the 37% speedup is reported on specific datasets and hardware; generalization to different model sizes, preference datasets, or training stacks (e.g., multi-node distributed training) remains unclear.
Research Context
This work builds directly on the success of batch packing in pre-training (used in models like Chinchilla, PaLM) and extends it to the preference optimization domain that has become central to modern LLM alignment (DPO, IPO, SimPO, etc.). The paper addresses a real efficiency gap: while batch packing is standard for SFT, the preference optimization community hasn't systematically optimized for the prompt-duplication pattern inherent to comparative losses. This fits into the broader push for training efficiency as models scale—alongside work on LoRA, attention approximations, and mixed-precision training. The research opens opportunities for further optimization: adaptive prompt sharing strategies, hardware-aware packing algorithms, and application to other comparative learning paradigms (contrastive learning, ranking losses in information retrieval). Given the widespread adoption of DPO across research labs and companies, preference packing could become a standard technique in the alignment engineer's toolkit.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::
