Uni-DocDiff: A Unified Document Restoration Model Based on Diffusion

TL;DR

Uni-DocDiff is a scalable diffusion-based model that unifies multiple document restoration tasks, leveraging learnable prompts and a novel prior pooling mechanism to improve performance and adaptability.

Contribution

It introduces a unified diffusion model with learnable prompts, a prior pool, and a prior fusion module, enabling scalable and effective multi-task document restoration.

Findings

Achieves comparable or superior performance to task-specific models.

Demonstrates high scalability across diverse document restoration tasks.

Effectively integrates local and global features for improved restoration.

Abstract

Removing various degradations from damaged documents greatly benefits digitization, downstream document analysis, and readability. Previous methods often treat each restoration task independently with dedicated models, leading to a cumbersome and highly complex document processing system. Although recent studies attempt to unify multiple tasks, they often suffer from limited scalability due to handcrafted prompts and heavy preprocessing, and fail to fully exploit inter-task synergy within a shared architecture. To address the aforementioned challenges, we propose Uni-DocDiff, a Unified and highly scalable Document restoration model based on Diffusion. Uni-DocDiff develops a learnable task prompt design, ensuring exceptional scalability across diverse tasks. To further enhance its multi-task capabilities and address potential task interference, we devise a novel \textbf{Prior \textbf{P}ool}, a simple yet comprehensive mechanism that combines both local high-frequency features and global low-frequency features. Additionally, we design the \textbf{Prior \textbf{F}usion \textbf{M}odule (PFM)}, which enables the model to adaptively select the most relevant prior information for each specific task. Extensive experiments show that the versatile Uni-DocDiff achieves performance comparable or even superior performance compared with task-specific expert models, and simultaneously holds the task scalability for seamless adaptation to new tasks.