Home / Research / Wasserstein diffusion
Research · IIIT Kottayam · Computer Science
Wasserstein Diffusion on the Annotation Simplex
Diffusion-based video summarisation averages annotator scores arithmetically and corrupts them with Gaussian noise. Annotator scores live on a probability simplex, so both steps leave it. The project proposes three corrections and keeps the rest of the pipeline as published.
Paper
Open proposal (PDF) ↗Timeline
- Proposal acceptedResearch proposal on correcting the geometry of diffusion-based video summarisation (Shang et al., AAAI-25) accepted.
- Sinkhorn projection and Dirichlet samplerPlanned.
- Wasserstein barycenter in the dataloaderPlanned.
- Combined loss and 5-split experimentsTVSum, SumMe and FPVSum. Planned.
- Ablation and write-upEach correction evaluated on its own. Planned.
- Converting to paperStatus at the end of the position.
The weekly schedule is the plan in the proposal.
How the thinking developed
Where the flaw appeared
While studying the AAAI-25 model, it became clear that each annotator’s importance scores sum to one, so they are points on a simplex, yet the model treats them as ordinary vectors.
Why the arithmetic mean fails
If annotators split into two camps, their mean falls between the camps and matches nobody’s judgement.
Why the fix is small
The Transformer noise predictor, the DR-DSN warm-start, the 200-step schedule and the Adam optimiser are all kept. Only the aggregation, the forward noise and the loss change.
Why it should not regress
When annotators agree the barycenter is close to the mean. It differs only when they disagree structurally, which is the case the datasets stress.
Pipeline: original and corrected
| Stage | Original | Proposed |
|---|---|---|
| Consensus | Arithmetic mean | Wasserstein barycenter |
| Forward process | Gaussian noise | Dirichlet noise, Sinkhorn projection onto the simplex |
| Noise predictor | Transformer, unchanged | |
| Loss | ‖ε − ε̂‖² | ‖ε − ε̂‖² + λ·W₂(x⁽ᵏ⁾, x̂₀) |
λ ≈ 0.1, chosen by cross-validation. Implemented with geomloss (PyTorch), at roughly 15–20% extra training time and unchanged inference.
Evaluation
Benchmarks: TVSum, SumMe and FPVSum with 5-split cross-validation, with an ablation of each correction on its own.
| Dataset · setting | AAAI-25 | Projected |
|---|---|---|
| TVSum · canonical | 64.8 | 65.3–65.6 |
| TVSum · transfer | 60.9 | 61.7–62.1 |
| SumMe · canonical | 58.7 | 60.2–61.2 |
| SumMe · transfer | 50.3 | 52.3–53.8 |
| FPVSum · transfer | 46.1 | 48.1–50.1 |
F-scores. The projected column holds the estimates in the proposal. They are not measured results.