Focusing the Diffusion: Importance Sampling for One-Shot Subject-Driven Image Generation

DreamBooth-style finetuning tends to overfit to backgrounds, poses, and other contextual cues rather than learning fine-grained subject-specific features with extremely small datasets. This behavior degrades true subject fidelity and severely limits prompt-driven editability.We propose P-LoRA (Probability LoRA), a simple training-time modification that biases diffusion timestep sampling while using importance weighting to preserve an unbiased objective. By increasing the probability of sampling low-noise timesteps where fine-grained subject details are most salient P-LoRA encourages learning identity-relevant features instead of contextual memorization. Experiments in one-shot and few-shot personalization regimes show that a moderate low-noise bias improves subject reconstruction fidelity while maintaining editability on par with the DreamBooth baseline. Shifting the bias toward higher-noise timesteps reduces background memorization but consistently lowers subject fidelity and weakens prompt adherence, revealing a clear trade-off between identity retention and controllability across timestep biases. To disentangle true subject learning from memorized context, we introduce a clean reconstruction metric that segments the foreground before computing CLIP similarity. Together, these results establish timestep reweighting as a lightweight and effective mechanism for navigating fidelity vs editability trade-offs under extreme data scarcity.