CAP-WAE for DR Grading

• The paper introduces CAP-WAE, a framework that integrates a constrained asymmetric prior with latent geometry regularization and direction-aware ordinal supervision for improved DR grading.
• It leverages a non-Gaussian asymmetric prior and a Margin-Aware Orthogonality and Compactness loss to capture heavy-tailed, skewed latent structures and enforce distinct grade clusters.
• Experimental results on four DR benchmarks show state-of-the-art performance improvements in Quadratic Weighted Kappa, accuracy, and F1 score compared to conventional methods.

Constrained Asymmetric Prior Wasserstein Autoencoder (CAP-WAE) is a generative–discriminative framework for diabetic retinopathy (DR) grading that combines a Wasserstein Autoencoder (WAE) with a constrained asymmetric prior, latent geometry regularization, and direction-aware ordinal supervision. It was introduced for settings in which labels are both ordinal and long-tailed, so conventional isotropic Gaussian priors and symmetric loss functions misalign the latent representation with clinically asymmetric error costs and the skewed structure of minority grades (Shaik et al., 30 Sep 2025).

1. Problem setting and motivation

CAP-WAE is formulated for DR grading, where labels lie on an ordered severity scale, typically the 5-grade set $\mathcal{Y}=\{0,1,2,3,4\}$ or a finer 7-grade variant. The ordering is clinically meaningful: misclassifying grade 3 as 2 is less serious than misclassifying 3 as 0. At the same time, real DR datasets are long-tailed, with many mild or moderate cases and relatively few severe or no-DR cases. Minority grades are described as scarce, heterogeneous, and clinically critical, especially because under-grading a severe case is substantially more consequential than a small over-grade (Shaik et al., 30 Sep 2025).

The model is motivated by several limitations attributed to standard alternatives. Conventional classifiers based on one-hot labels and symmetric cross-entropy treat all misclassifications equally. Ordinal-specific methods such as CORN and OLDL model order, but usually assume symmetric label distributions around the true grade. Standard variational autoencoders use the instance-wise objective

$$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$

typically with $p(z)=\mathcal{N}(0,I)$, which the paper argues can suppress non-Gaussian minority-class structure and create overlapping clusters. Plain WAEs improve the regularization pattern by matching the aggregate posterior rather than each sample-wise posterior, but when they retain $p(z)=\mathcal{N}(0,I)$, they still impose a symmetric Gaussian latent geometry that does not represent skew or heavy tails in minority grades (Shaik et al., 30 Sep 2025).

This motivates three coupled design requirements: a non-Gaussian asymmetric prior tuned to latent statistics, explicit shaping of latent geometry by class and order, and supervision that distinguishes the direction of ordinal error.

2. Position within the WAE and OT autoencoder lineage

CAP-WAE inherits its generative backbone from the Wasserstein Auto-Encoder framework, in which generative modeling is formulated as minimizing a penalized form of the Wasserstein distance between the model distribution and the target distribution. In the canonical WAE formulation,

$$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$

where $Q_Z$ is the aggregated posterior. A defining property of WAE is that the regularizer acts on the aggregated posterior $Q_Z$, rather than on each conditional $Q(Z|X=x)$, which differentiates it from the standard VAE family (Tolstikhin et al., 2017).

Within the broader OT-autoencoder literature, CAP-WAE occupies a specific methodological niche. Symmetric Wasserstein Autoencoders introduce a learnable prior and symmetrically match the joint distributions of the observed data and the latent representation induced by encoder and decoder, jointly optimizing losses in data and latent spaces (Sun et al., 2021). Sinkhorn AutoEncoders show that latent-space Wasserstein penalties can be computed by differentiable Sinkhorn iterations and emphasize that the framework can work with arbitrary priors and metric spaces (Patrini et al., 2018). A separate WAE line in scientific modeling uses a hyperspherical prior with projection-based sampling rather than direct von Mises–Fisher sampling, illustrating that non-Euclidean or non-Gaussian priors are compatible with WAE training (Shukla et al., 2021).

CAP-WAE differs from these variants in the source and purpose of its asymmetry. Its prior is not merely learnable or geometrically non-Euclidean; it is explicitly designed to preserve heavy-tailed and skewed latent structure associated with imbalanced ordinal classes, and it is coupled to auxiliary losses that enforce grade-ordered separability and directional clinical costs. A common misconception is therefore to treat CAP-WAE as only a prior-substituted WAE. In the published formulation, it is a compound model whose prior, geometry, and supervision are jointly specialized to ordinal long-tailed grading.

3. Architecture and constrained asymmetric prior

The architecture consists of an encoder $f_\phi$, a decoder $g_\theta$, and three latent-space heads. The encoder is a VGG16 backbone, ImageNet-pretrained, followed by two fully connected layers, mapping a fundus image $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$0 to a latent vector $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$1. The decoder is a symmetric transposed-convolution network producing the reconstruction $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$2. On top of $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$3, CAP-WAE attaches a classification head $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$4, an asymmetric Gaussian head $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$5, and an ordinal regression head $$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$6 (Shaik et al., 30 Sep 2025).

Its WAE component starts from the standard reconstruction-plus-distribution-matching template,

$$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$7

and instantiates it as

$$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$8

$$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - \mathrm{KL}\!\left(q_\phi(z|x)\,\|\,p(z)\right),$$9

and

$p(z)=\mathcal{N}(0,I)$0

The central modification is the constrained asymmetric prior $p(z)=\mathcal{N}(0,I)$1. The prior factorizes over latent dimensions,

$p(z)=\mathcal{N}(0,I)$2

with each scalar coordinate following an Asymmetric Generalized Gaussian Distribution (AGGD): $p(z)=\mathcal{N}(0,I)$3 Here $p(z)=\mathcal{N}(0,I)$4 is the location, $p(z)=\mathcal{N}(0,I)$5 is the shape, and $p(z)=\mathcal{N}(0,I)$6 are left and right scales. The paper notes that $p(z)=\mathcal{N}(0,I)$7 yields a Gaussian and $p(z)=\mathcal{N}(0,I)$8 yields a Laplace distribution; smaller $p(z)=\mathcal{N}(0,I)$9 produces heavier tails. Parameters are estimated empirically from latent samples,

$p(z)=\mathcal{N}(0,I)$0

$p(z)=\mathcal{N}(0,I)$1

and are then treated as fixed. The paper explicitly characterizes this as the “constrained” aspect of the prior: it is not arbitrarily learned, but stabilized by empirical latent statistics (Shaik et al., 30 Sep 2025).

An important clarification concerns class structure. The prior is not explicitly class-conditional. Instead, it is fit on latents from all classes after those latents have already been shaped by the geometry and supervision terms. The paper states that minority classes with more variation contribute to heavier tails, while asymmetry in class distributions yields nonzero skew captured by different left and right scales. This suggests that the prior functions as a global distributional envelope rather than a direct class-indexed prior.

4. Margin-Aware Orthogonality and Compactness

CAP-WAE adds a dedicated latent-geometry regularizer, Margin-Aware Orthogonality and Compactness (MAOC), to enforce class separation and within-class concentration. Let $p(z)=\mathcal{N}(0,I)$2 be the batch latents, $p(z)=\mathcal{N}(0,I)$3 the corresponding grades, and $p(z)=\mathcal{N}(0,I)$4 the running mean or prototype of class $p(z)=\mathcal{N}(0,I)$5. With normalized prototypes $p(z)=\mathcal{N}(0,I)$6, the loss is

$p(z)=\mathcal{N}(0,I)$7

The orthogonality term penalizes prototype inner products above the margin $p(z)=\mathcal{N}(0,I)$8, while the compactness term pulls each sample toward its class prototype (Shaik et al., 30 Sep 2025).

The geometric intent is twofold. First, prototype directions are pushed apart so that different grades occupy approximately orthogonal directions with an angular margin. Second, each grade forms a compact cluster around its prototype. The paper interprets the combined effect as yielding well-separated, compact clusters arranged so that movement between grades follows well-defined directions in latent space.

MAOC interacts with the asymmetric prior rather than replacing it. The AGGD prior determines global distributional shape—skew and heavy tails—while MAOC organizes class-wise geometry inside that shape. The paper’s interpretation is that heavy tails provide room for skewed, long-tailed manifolds, whereas MAOC reduces cross-grade overlap by sharpening prototypes and collapsing within-grade dispersion. A common simplification is to treat the prior alone as the source of latent order; the published formulation rejects this, because ordered separability is enforced explicitly through MAOC and not derived automatically from AGGD fitting.

5. Direction-aware ordinal supervision

The supervision stack combines hard classification, asymmetric label-distribution learning, and an auxiliary ordinal regression objective. For each latent code $p(z)=\mathcal{N}(0,I)$9, the asymmetric Gaussian head predicts left and right dispersions,

$$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$0

which are clamped to $$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$1. These dispersions define a direction-aware soft target over grades,

$$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$2

where $$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$3 and $$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$4 (Shaik et al., 30 Sep 2025).

The corresponding direction-aware ordinal loss is a soft cross-entropy,

$$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$5

where $$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$6 are the classification logits. This differs from both standard cross-entropy and symmetric label-distribution learning because the mass around the true grade is skewed rather than symmetric. The paper frames this asymmetry as a mechanism for encoding clinical priorities, for example by allowing a larger spread on the over-grade side and a narrower spread on the under-grade side.

The model also retains hard classification through

$$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$7

and adds an ordinal regression head producing a scalar severity score $$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$8, trained with a Huber loss,

$$D_{\mathrm{WAE}}(P_X,P_G):= \inf_{Q(Z|X)\in\mathcal{Q}} \mathbb{E}_{P_X}\,\mathbb{E}_{Q(Z|X)} \bigl[c(X,G(Z))\bigr] + \lambda \cdot \mathcal{D}_Z(Q_Z, P_Z),$$9

The paper states that this auxiliary regression aligns the latent embedding with an underlying severity axis. Taken together, the three supervised terms serve different roles: hard class discrimination, direction-aware distributional supervision, and regression-style ordinal consistency. This suggests that CAP-WAE should be understood as an ordinal label-distribution learning model as much as an autoencoder.

6. Optimization, empirical results, and limitations

CAP-WAE balances reconstruction, prior alignment, latent geometry, and three supervised losses through an adaptive multi-task weighting scheme. For $Q_Z$0, the model introduces learnable log-variance parameters $Q_Z$1, yielding the total objective

$Q_Z$2

In this formulation, $Q_Z$3 acts as an inverse variance, so larger uncertainty downweights the corresponding supervised term, while the additive $Q_Z$4 prevents trivial solutions (Shaik et al., 30 Sep 2025).

The reported implementation uses latent dimension $Q_Z$5, batch size 32, 100 epochs, AdamW with learning rate $Q_Z$6 and weight decay $Q_Z$7, $Q_Z$8, $Q_Z$9, angular margin $Q_Z$0, compactness weight $Q_Z$1, AG-soft dispersions clamped to $Q_Z$2, and AGGD shape parameter $Q_Z$3. Training uses mixed precision, gradient norm clipping at 1.0, and a ReduceLROnPlateau scheduler on validation QWK with factor 0.2 and patience 7; the best QWK checkpoint is retained. The per-minibatch pipeline is: encode images to $Q_Z$4, reconstruct $Q_Z$5, evaluate the three heads, update class means for MAOC, sample prior codes $Q_Z$6, compute MMD and the remaining losses, and optimize the full objective with AdamW (Shaik et al., 30 Sep 2025).

The model is evaluated on four public DR benchmarks: Zenodo-DR-7, IDRiD, APTOS-2019, and Messidor-2, with images resized to $Q_Z$7 and normalized to ImageNet statistics. The principal metrics are Quadratic Weighted Kappa (QWK), accuracy, and macro-F1.

Dataset	Reported QWK	Other reported metrics
Zenodo-DR-7	0.94	Acc 91.8%, F1 89.9%
APTOS-2019	0.90	F1 83.64%
Messidor-2	0.89	F1 83.0%
IDRiD	0.87	F1 81.2%

The paper reports state-of-the-art performance across these datasets and states that CAP-WAE outperforms conventional imbalance methods, ordinal baselines, latent generative baselines such as VAE-KL and WAE-MMD, and additional disease-grading baselines. The ablation study attributes cumulative gains to the asymmetric prior, WAE alignment, AG-soft supervision, ORM, MAOC, and adaptive weighting. Specifically, the transition from VAE-KL to VAE-KL with asymmetric prior yields $Q_Z$8 to $Q_Z$9 QWK, while WAE-MMD with asymmetric prior yields $Q(Z|X=x)$0 QWK and F1 boosts; subsequent additions of AG-soft, ORM, and MAOC further improve QWK and macro-F1. The t-SNE analysis is described as showing a progression from overlapping clusters in VAE-KL, to partial separation in WAE-MMD, to compact grade-ordered clusters with reduced overlap in full CAP-WAE.

Several limitations are explicitly noted. The AGGD prior parameters are estimated from training data and then fixed, which may be suboptimal under domain shift. MAOC provides no formal guarantee of globally optimal geometry or perfect ordering. The method depends on reliable ordinal labels despite known inter-grader variability in DR annotation. It also incurs computational overhead from the decoder, WAE alignment, and multiple supervision heads. Future directions proposed in the paper include adaptive prior learning under domain shift, self-supervised or multi-modal extensions, longitudinal progression modeling, and transfer to other heavy-tailed ordinal medical tasks such as cancer staging and fibrosis scoring.