🤖 AI

DisMix: Order-Aware Mixup for Medical Imaging via Disentangling Ordinal and Non-Ordinal Features

DisMix is an order-aware mixup framework for medical imaging that employs a dual-codebook VQ-VAE to disentangle ordinal and non-ordinal features, enabling independent mixing that preserves disease severity progression while enhancing appearance diversity to improve ordinal classification performance.

Original authors: Dileepa Pitawela, Gustavo Carneiro, Hsiang-Ting Chen

Published 2026-08-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dileepa Pitawela, Gustavo Carneiro, Hsiang-Ting Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery where the clues aren't just "guilty" or "innocent," but exist on a sliding scale of severity. In the world of medical imaging, doctors often have to grade diseases not as simple yes/no questions, but as a progression: a tiny scratch, a small bruise, a deep cut, or a shattered bone. This is called ordinal classification. It's like sorting a stack of photos from "sunny day" to "stormy night" rather than just sorting them into "blue" or "red."

To teach computers to do this, scientists use a clever trick called mixup. Imagine taking two photos, blending them together like mixing paint, and telling the computer, "This new picture is halfway between the two." Usually, this works great for simple tasks. But in medicine, it's like trying to mix a photo of a healthy knee with a photo of a broken one and hoping the result looks like a "slightly broken" knee. Often, the computer gets confused because the blending process smears the important medical clues (the severity) with random background noise (like the angle of the X-ray or the color of the skin). This paper tackles that specific mess, proposing a smarter way to blend medical images so the computer learns the right lessons without getting distracted by the noise.

The Problem: When Mixing Paints Ruins the Picture

The authors, Dileepa Pitawela, Gustavo Carneiro, and Hsiang-Ting Chen, noticed a major flaw in how computers currently learn to grade medical diseases. Standard "mixup" techniques are like a chaotic chef who throws ingredients into a blender without checking if they belong together. If you blend a mild disease image with a severe one, the computer sees a jumbled mess. It might accidentally mix the severity of the disease (the important part) with the background of the image, like the lighting or the patient's skin tone (the unimportant part).

The result? The computer learns to see "mild" and "severe" as a confusing blur. It might think a healthy knee looks like a broken one just because the lighting was weird, or it might create a "Frankenstein" image that looks like a knee joint that doesn't exist in real life. The paper argues that this "indiscriminate blending" destroys the very structure doctors rely on to grade diseases.

The Solution: DisMix (The Smart Blender)

To fix this, the team introduced DisMix. Think of DisMix as a super-smart blender that has two separate chutes. Before it blends anything, it sorts the ingredients into two piles:

The "Severity" Pile: This contains only the clues that tell the computer how bad the disease is (like the size of a lesion or the gap in a joint).
The "Style" Pile: This contains everything else that doesn't matter for grading, like the angle of the photo, the color of the skin, or random artifacts.

DisMix uses a special tool called a dual-codebook VQ-VAE (a fancy name for a machine that learns to separate these two piles). It forces the computer to learn that "how sick the patient is" and "what the photo looks like" are two different things.

Once the ingredients are sorted, DisMix blends them in a very specific way:

Blending Severity: It takes the "Severity" clues from a mild case and a severe case and mixes them to create a perfect "medium" case. This teaches the computer what a "halfway" disease looks like.
Blending Style: It takes the "Style" clues from different patients and swaps them around. This adds variety to the training data without messing up the severity score.

What They Found: A Clearer Picture

The team tested this idea on four different medical datasets, including knee X-rays (Knee Osteoarthritis), eye scans (Diabetic Retinopathy), and tissue samples. They compared DisMix against six other popular mixing methods and tested it with six different types of medical grading algorithms.

The results were promising. In 20 out of 24 different test scenarios, DisMix achieved the highest accuracy. It also reduced the average error in grading by 4–6% compared to the best existing methods.

One of the most interesting findings was how well DisMix handled "grading variability." In real life, different doctors might disagree on whether a disease is "mild" or "moderate." DisMix was robust even when the training data was messy or when there were very few images to learn from. For example, when the team limited the training data to just 10% of the usual amount, DisMix still outperformed the others, improving accuracy by about 1.5% on one dataset even with only 60 original images.

Why It Matters

The paper suggests that by separating the "what" (the disease severity) from the "how it looks" (the background noise), we can teach computers to be better at understanding the progression of diseases. The authors showed that their method creates images that look biologically plausible—like a knee joint that is actually narrowing, rather than a weird, distorted shape.

While the generated images are meant to help train the computer and not to be used directly for patient diagnosis, the study shows that this "order-aware" approach is a powerful step forward. It proves that when we stop blindly mixing everything together and start respecting the order of medical severity, our AI models become sharper, more reliable, and better at handling the messy reality of medical data.

Technical Summary: DisMix – Order-Aware Mixup for Medical Imaging via Disentangling Ordinal and Non-Ordinal Features

Problem Statement

Standard image mixup augmentation, a technique widely used to improve generalization in deep learning by generating intermediate samples, is fundamentally ill-suited for ordinal classification tasks prevalent in medical imaging (e.g., disease severity grading). In these domains, labels represent a progression of severity (e.g., Kellgren–Lawrence scale 0–4) rather than discrete, independent categories.

Conventional mixup strategies indiscriminately blend two distinct types of features:

Ordinal features: Cues that define the rank or severity (e.g., lesion size, joint space narrowing).
Non-ordinal features: Appearance-level variations unrelated to severity (e.g., imaging artifacts, background, staining differences).

This entangled mixing corrupts the underlying ordinal structure. As demonstrated in the paper, standard mixing produces samples with biologically implausible structures (e.g., disoriented knee joint gaps or multiple optic discs) and distorts the feature distribution, causing samples from different ranks to become entangled. This failure mode hinders the learning of robust ordinal classifiers, particularly in the presence of grading variability (disagreements among annotators).

Methodology: DisMix

To address the incompatibility between conventional mixup and ordinal learning, the authors propose DisMix, an order-aware mixup framework. The core innovation is the explicit disentanglement of ordinal and non-ordinal features into separate latent subspaces before mixing occurs.

1. Disentangled Representation Learning

DisMix employs a dual-codebook Vector Quantized Variational Autoencoder (VQ-VAE) to factorize input images into two distinct latent components:

Ordinal Latent ( $z_o$ ): Captures rank-defining structure.
Non-Ordinal Latent ( $z_n$ ): Encodes rank-independent variations (appearance).

The architecture is optimized using the following mechanisms:

Ordinal Branch: The encoder features are conditioned on ordinal soft labels (derived from multiple annotators) to bias the representation toward rank semantics. An attention-based quantizer selects discrete codes from an ordinal codebook ( $M_o$ ). An auxiliary linear classifier enforces that $z_o$ retains ordinal information via Soft Ordinal Regression (SOR) loss.
Non-Ordinal Branch: To suppress ordinal information, this branch utilizes an adversarial objective. A linear classifier attempts to predict the rank from $z_n$ , but its gradients are reversed via a Gradient Reversal Layer (GRL). This forces the non-ordinal quantizer to discard rank-sensitive signals.
Disentanglement Regularization: The cosine similarity between the two latent representations is minimized to ensure the subspaces remain distinct.
Adversarial Training: A multi-task discriminator ensures perceptual realism and ordinal consistency of the reconstructed images.

2. Order-Aware Mixup Policies

Once the VQ-VAE is trained and frozen, DisMix generates augmented samples by manipulating the latent codes independently. The paper defines four mixing policies (Table 1):

Ordinal Mix: Interpolates ordinal codes ( $\tilde{z}_o = \lambda z_o^i + (1-\lambda)z_o^j$ ) while keeping non-ordinal codes fixed. This generates intermediate ranks with soft labels.
Non-Ordinal Mix: Interpolates non-ordinal codes while keeping the ordinal code (and thus the rank label) fixed. This introduces appearance diversity without altering severity.
Order Swap: Swaps the ordinal and non-ordinal codes between two samples.
Generate & Mix: Generates an adjacent-rank variant of an anchor image by conditioning its ordinal code on a new rank, then applies Ordinal Mix.

These policies allow for controlled augmentation that preserves the ordinal manifold while enriching appearance diversity.

Key Contributions

DisMix Framework: Introduction of an order-aware mixup method that disentangles ordinal and non-ordinal features into distinct latent subspaces, enabling the synthesis of order-consistent augmented samples.
Disentangled Latent Policies: Proposal of specific mixing strategies (Ordinal Mix, Non-Ordinal Mix, Order Swap, Generate & Mix) that operate separately on the disentangled subspaces to prevent the corruption of ordinal signals.
Comprehensive Evaluation: Extensive benchmarking across four medical imaging datasets (Knee Osteoarthritis, Diabetic Retinopathy, Colorectal Histopathology, Prostate Histopathology) against six state-of-the-art mixup baselines and six ordinal classifiers.

Experimental Results

The authors evaluated DisMix across 24 combinations of mixup methods and ordinal classifiers. Key findings include:

Aggregate Performance: DisMix achieved the best accuracy in 20 out of 24 settings and the best Mean Absolute Error (MAE) in 15 out of 24 settings.
Statistical Significance: Compared to the strongest baseline, DisMix delivered statistically significant aggregate accuracy gains (one-sided Wilcoxon signed-rank test, $p = 0.0075$ ) and reduced MAE by 4–6%.
Data Scarcity: Under data scarcity (training on only 10% of samples per class), DisMix outperformed baselines, improving accuracy by approximately 1.5% on the IDRID dataset with minimal training data.
Robustness to Grading Variability: DisMix remained effective under high grading variability (up to 70% disagreement), whereas other methods suffered performance degradation, particularly at high mixing probabilities.
Boundary Error Reduction: DisMix consistently achieved the lowest boundary-level classification errors (e.g., reducing C2↔C3 errors by 9% on IDRID compared to the best baseline).
Qualitative Analysis: Visualizations confirmed that DisMix produces biologically plausible transitions (e.g., darker lesions, narrowed joint spaces) while preserving non-ordinal attributes, whereas other methods often generated anatomically inconsistent artifacts.

Significance and Claims

The paper argues that the fundamental incompatibility between conventional mixup and ordinal learning stems from the entanglement of rank-defining features with irrelevant appearance variations. By disentangling these factors, DisMix enables controlled mixing that respects the underlying ordinal geometry.

The authors claim that DisMix demonstrates the necessity of order-aware mixing for ordinal data. Unlike prior methods that either corrupt labels or under-utilize augmentation by ignoring appearance cues, DisMix provides a mechanism to preserve label fidelity while enhancing semantic diversity. The results underscore that disentanglement-based methods are a promising direction for robust ordinal learning in medical imaging, particularly in scenarios involving grading variability and limited data.

The paper maintains a modest stance regarding its scope, noting that the generated images are intended for data augmentation rather than direct clinical interpretation, and that the interpolated ranks serve as soft supervision rather than new clinically defined grades. Additionally, the authors acknowledge the computational overhead of training the disentangled representation, though they note this is a one-time cost per dataset amortized across downstream tasks.