---
title: 'Uni-DocDiff: Unified Document Restoration'
url: https://www.emergentmind.com/topics/uni-docdiff
type: topic
---

# Uni-DocDiff: Unified Document Restoration

Uni-DocDiff is a unified and scalable model for document restoration based on a conditional diffusion framework. It is designed to address a wide spectrum of document degradations—deblurring, deshadowing, illumination correction, binarization, handwriting removal, and geometric dewarping—within a single architecture. Unlike previous approaches which rely on separate expert models for each restoration task or handcrafted multi-task unification strategies, Uni-DocDiff employs learnable task prompts, a set of classical feature-based priors, and an adaptive fusion mechanism. Extensive empirical results demonstrate state-of-the-art or superior performance compared to both specialized and prior unified models [2508.04055].

## 1. Conditional Diffusion Backbone

Uni-DocDiff’s core is a conditional diffusion model, specifically a Pixel Prediction Branch (PPB) modeled after Denoising Diffusion Probabilistic Models. During training, the clean document image \(x_0\) is corrupted through a fixed, progressively noised forward process:
\[
q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{\alpha_t} x_{t-1}, (1-\alpha_t)I)
\]
with the marginal
\[
x_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)
\]
and \(\bar\alpha_t = \prod_{i=1}^t \alpha_i\).

The reverse process reconstructs the original image, conditioned on the degraded input \(x_d\) and a learnable task prompt \(\tau\):
\[
p_\theta(x_{t-1}\mid x_t, x_d, \tau) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, x_d, \tau), \Sigma_\theta(\tau))
\]

Uni-DocDiff predicts \(\hat x_0\) directly for sharper results:
\[
\hat x_0 = \mathcal{F}(x_t; x_d, \tau, P)
\]
enabling the reconstruction mean
\[
\mu_\theta = \frac{\sqrt{\bar\alpha_{t-1}}(x_t - \sqrt{1-\bar\alpha_t}\epsilon_\theta(x_t)) + \sqrt{1-\bar\alpha_{t-1}}\hat x_0}{\sqrt{\bar\alpha_t}}
\]
The loss for training deviates from the standard MSE (“simple” loss), instead using \(\ell_1\) loss and frequency-weighted auxiliary terms according to task frequency (see Section 5).

## 2. Learnable Task Prompt Encoding

Task specification in Uni-DocDiff is achieved via a learnable, low-dimensional vector \(\tau \in \mathbb{R}^d\) that represents the restoration task (for example, deblurring or binarization). At each level \(l\) in the backbone, this task prompt is processed by a multilayer perceptron (MLP):
\[
w_{\mathrm{task}}^l = \mathrm{MLP}_\mathrm{task}^l(\tau) \in \mathbb{R}^{c_l}
\]
These weights modulate the gating of prior features via the Prior Fusion Module (PFM), contextualizing U-Net features for the task at hand. This learnable prompt mechanism imparts scalability and extensibility, as new tasks can be supported by embedding new prompts and retraining lightweight modules, without retraining the full model.

## 3. Prior Pool Construction

The Prior Pool (\(P \in \mathbb{R}^{H \times W \times C}\)) supplies explicit local high-frequency and global low-frequency cues derived from classical image processing operators:
- **High-frequency features:** Sobel filters (\(\mathrm{Sobel}_x(x_d), \mathrm{Sobel}_y(x_d)\)) and Canny edge detection (\(\mathrm{Canny}(x_d)\)), capturing structural, text, and edge details.
- **Low-frequency features:** Median filtering (\(\mathrm{Med}(x_d)\)), Gaussian blur (\(\mathrm{Gauss}(x_d)\)), and DCT-based low-frequency reconstruction (retaining the top \(K\%\) DCT coefficients), emphasizing smooth content and background.

All six maps, resized to match the input resolution, are concatenated along the channel dimension to form \(P\). The use of inexpensive, interpretable filters, as opposed to learned feature extractors, provides explicit access to frequency-tailored information for restoration tasks that are sensitive to specific content frequencies.

## 4. Prior Fusion Module (PFM)

PFM mediates the integration of classical priors and learned features at each U-Net level. The module consumes the U-Net activation map \(f^l\), the Prior Pool \(P\), and the task prompt \(\tau\), and fuses them as follows:
1. **Stage-specific adaptation:** Adapt priors through a learned convolution:
   \[
   P^l = \mathrm{Conv}_{\mathrm{adapt}}^l(P)
   \]
2. **Content-driven weighting:** Compute weights from global average-pooled features:
   \[
   w_{\mathrm{content}}^l = \mathrm{MLP}_{\mathrm{content}}^l(\mathrm{GAP}(f^l))
   \]
3. **Task-driven weighting:** Project the task prompt using an MLP:
   \[
   w_{\mathrm{task}}^l = \mathrm{MLP}_{\mathrm{task}}^l(\tau)
   \]
4. **Channel-wise weighting and fusion:**
   \[
   P_{\mathrm{content}}^l = P^l \otimes w_{\mathrm{content}}^l, \quad P_{\mathrm{task}}^l = P^l \otimes w_{\mathrm{task}}^l
   \]
   (\(\otimes\) denotes broadcasted channel-wise multiplication.)
5. **Integration with intermediate features:**
   \[
   f_{\mathrm{recon}}^l = \mathrm{ResBlock}^l(\mathrm{Concat}(f^l, P_{\mathrm{content}}^l + P_{\mathrm{task}}^l))
   \]
This mechanism enables adaptive selection of relevant priors per task and content state, mitigating multi-task interference by disentangling frequency and task-specific features.

## 5. Training Paradigm and Multi-Task Loss

Uni-DocDiff employs a two-stage training pipeline:
- **Stage 1—Pixel prediction:** Joint training on five restoration tasks (deblurring, deshadowing, illumination correction, binarization, handwriting removal) via mini-batch sampling across tasks. Losses are task-frequency-tailored:
  \[
  \mathcal{L}_{\mathrm{low}} = \mathbb{E}\|\hat x_0 - x_0\|_1 + \beta_1 \mathbb{E}\|\phi_L(\hat x_0) - \phi_L(x_0)\|_1
  \]
  for low-frequency tasks, and
  \[
  \mathcal{L}_{\mathrm{high}} = \mathbb{E}\|\hat x_0 - x_0\|_1 + \beta_2 \mathbb{E}\|\phi_H(\hat x_0) - \phi_H(x_0)\|_1
  \]
  for high-frequency tasks, where \(\phi_L\) and \(\phi_H\) are low-pass (Gaussian) and high-pass (Sobel) filters respectively. The overall loss is a weighted sum, normalized by batch size per task type.

- **Stage 2—Geometric prediction:** With the backbone frozen, a Coordinate Prediction Branch is trained separately on dewarping datasets using a simple \(\ell_1\) loss on backward deformation maps:
  \[
  \mathcal{L}_{\mathrm{CPB}} = \mathbb{E}\|bm - bm_{\mathrm{gt}}\|_1
  \]
Optimizer: AdamW with a learning rate of \(10^{-4}\), employing standard augmentations.

## 6. Quantitative Evaluation and Ablation

Empirical validation spans six standard document restoration benchmarks. Across deblurring (TDD), deshadowing (Jung’s, Kligler’s, OSR), illumination rectification (DocUNet, RealDAE), dewarping (DIR300, DocUNet*), binarization (DIBCO’18), and handwriting removal (EnsExam), Uni-DocDiff attains performance matching or surpassing specialist baselines and the prior unified DocRes architecture. Selected results are summarized below:

| Task             | Baseline (Metric)         | Uni-DocDiff (Metric)     |
|------------------|--------------------------|--------------------------|
| Deblurring (TDD) | DE-GAN: 0.9226/22.24     | **0.9824/28.77**         |
| Deshadowing      | BGSNet: 0.9094/17.34     | **0.9156/23.93**         |
| Illum. rect.     | GCDRNet: 0.9423/24.42    | **0.9485/24.97**         |
| Dewarping (DIR)  | UVDoc: 0.6380/6.40/0.218 | **0.6573/5.30/0.203**    |
| Binarization     | DocRes: 89.82/94.33/19.35| **90.32/93.84/19.76**    |
| Handw. removal   | EnsExam: 0.9671/35.68    | **0.9685/36.23**         |

Ablative studies show that removing either the Prior Pool or PFM degrades multi-task performance, highlighting their role in mitigating task interference. Incorporating frequency-aware loss terms yields a measurable gain (0.5–1 dB PSNR) across low- and high-frequency tasks.

A notable property is the model’s extensibility: when adapting to new tasks (e.g., novel handwriting-removal splits), retraining only lightweight modules like PFM suffices. This contrasts with alternative prompt-based unification systems that require more substantial retraining.

## 7. Significance and Architectural Implications

Uni-DocDiff advances document restoration by offering a unified, highly scalable architecture capable of expert-level performance across disparate restoration tasks. The integration of interpretable, computationally light priors with frequency- and task-adaptive fusion mechanisms exemplifies a modular, extensible system. The empirical superiority over disjoint and prior unified baselines, coupled with graceful task-adaptability, marks Uni-DocDiff as a substantive development in the progression toward generalized document analysis and restoration pipelines [2508.04055].

Source: https://www.emergentmind.com/topics/uni-docdiff