🤖 AI

OptiModNet: A UNet-Transformer Hybrid with Grouped-Query and Channel Attention for Optic Disc and Cup Segmentation

The paper introduces OptiModNet, a lightweight hybrid UNet-Transformer architecture incorporating grouped-query and channel attention mechanisms along with an Aggregated Pyramid Loss, which achieves state-of-the-art performance in optic disc and cup segmentation on the REFUGE2 dataset while maintaining high computational efficiency.

Original authors: Soumili Ghosh, Debapriya Roy, Aryan Das, Bikash Santra

Published 2026-08-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Soumili Ghosh, Debapriya Roy, Aryan Das, Bikash Santra

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

Glaucoma is a leading cause of irreversible blindness worldwide, often developing silently without early symptoms. Because the damage is permanent, catching the disease in its earliest stages is critical. Doctors rely on examining the back of the eye, specifically looking at two circular structures: the optic disc, where the nerve enters the eye, and the optic cup, a smaller depression inside it. By measuring the size of the cup relative to the disc, clinicians can calculate a ratio that serves as a vital warning sign for the disease. However, drawing these boundaries by hand is difficult, time-consuming, and prone to human error, especially when the images are blurry or the contrast is low. To solve this, scientists have turned to artificial intelligence, teaching computers to recognize these structures automatically. While current computer programs can do this, the most accurate ones are often so heavy and complex that they require powerful, expensive hardware, making them difficult to use in everyday clinics or large-scale screening programs where speed and simplicity are essential.

A team of researchers has developed a new approach that balances high accuracy with lightness, creating a system they call OptiModNet. Their goal was to build a model that could see the whole picture of the eye's anatomy while also noticing the tiny, fine details of the edges, without needing a supercomputer to run it. They achieved this by combining two different types of artificial intelligence architectures. One type, known as a UNet, is excellent at spotting local details and edges, much like a person looking closely at a specific part of a map. The other type, based on a transformer, is better at understanding the big picture and how different parts of the image relate to each other over long distances. Previous attempts to mix these two often resulted in models that were still too heavy or complex. The researchers in this study refined the mixing process by introducing two specific improvements. First, they changed how the model processes information in its "thinking" layer to group questions together, which significantly reduces the amount of calculation needed. Second, they added a mechanism in the decoding layer that helps the model focus on the most important features while ignoring the less useful ones, sharpening the final outline of the optic disc and cup.

To train this new system, the researchers used a large collection of retinal images known as the REFUGE2 dataset, which contains thousands of eye scans with expert-drawn boundaries. They also tested their model on a second dataset called ORIGA to ensure it worked well on different types of images. The results were striking. On the first dataset, the new model correctly identified the optic disc with a score of 99.45% and the optic cup with 93.23%, outperforming the previous best methods by more than 2.5%. On the second dataset, it achieved the highest accuracy ever reported for that specific test. Perhaps more importantly, the model is incredibly efficient. It requires only 1.93 million parameters to function, which is roughly 96% fewer than the most complex competing models. In terms of computational power, it uses just 3.73 GFLOPs, a fraction of the energy required by other top-performing systems. This means the model could potentially run on standard medical equipment rather than needing specialized, high-end servers.

The researchers also tested how each part of their design contributed to the final success. They found that adding the grouping mechanism and the focus-enhancing feature layer individually improved the results, but combining them produced the best outcome. They also introduced a new way of teaching the model during training, which checks the model's progress at several different stages of learning rather than just at the end. This method helped the model learn more consistently and produce smoother, more accurate boundaries. When they compared the visual output of their model against other systems, the new model produced outlines that aligned much more closely with the expert-drawn ground truth, even in difficult cases where blood vessels crossed the area or the lighting was uneven. The study concludes that by carefully blending these different techniques, it is possible to create a tool that is both highly accurate and light enough for widespread use. This advancement suggests that high-quality, automated screening for glaucoma could soon become a practical reality in resource-limited settings, potentially saving sight for many more people.

Technical Summary: OptiModNet

Problem Statement

Precise segmentation of the optic disc (OD) and optic cup (OC) in retinal fundus images is a critical prerequisite for the early detection and diagnosis of glaucoma, primarily through the calculation of the cup-to-disc ratio (CDR). However, achieving high-performance segmentation across diverse datasets while maintaining low computational overhead remains a significant challenge. Existing deep learning approaches face specific trade-offs:

UNet-based models excel at capturing local features but struggle to model long-range dependencies and global anatomical context.
Vision Transformers (ViTs) effectively capture global relationships via self-attention but incur substantial computational costs (quadratic complexity) and require large annotated datasets, making them resource-intensive for high-resolution fundus imaging.
Hybrid architectures (e.g., TransUNet, UNetR) attempt to bridge this gap but often inherit the computational inefficiencies of vanilla Transformers.
Diffusion-based models show competitive results but suffer from high training and inference costs, limiting their suitability for real-time clinical workflows.

Consequently, there is a need for a lightweight, efficient architecture that balances global context modeling with local boundary precision without prohibitive computational overhead.

Methodology

The authors propose OptiModNet, a novel lightweight hybrid architecture tailored for 2D optic disc and cup segmentation. The model integrates a Transformer-based encoder with a Convolutional decoder, enhanced by specific attention mechanisms and a multi-stage loss function.

1. Architecture Overview

Encoder (Transformer-based): The input image is divided into non-overlapping patches and embedded with positional information. The encoder consists of 12 stacked transformer blocks.
- Grouped Query Attention (GQA): To reduce computational complexity compared to standard Multi-Head Self-Attention (MSA), the encoder employs GQA. This mechanism splits queries into groups, where each group shares a single key-value pair. This streamlines the attention mechanism, significantly lowering the computational load while maintaining the ability to model long-range dependencies in high-resolution inputs.
Decoder (Convolutional): The decoder incrementally upsamples feature maps using transposed convolutions and integrates multi-resolution skip connections from the encoder.
- Channel Attention (CA): A novel component for the decoder, Channel Attention is applied to recalibrate feature importance. It uses Global Average Pooling followed by fully connected layers to generate attention weights, emphasizing informative feature channels and attenuating less relevant ones. This specifically aids in the precise delineation of OD/OC boundaries.

2. Aggregated Pyramid Loss

To optimize both region-level overlap and pixel-wise classification, the authors introduce the Aggregated Pyramid Loss.

Composition: It combines Soft Dice Loss and Binary Cross-Entropy (BCE) Loss.
Multi-Stage Supervision: Unlike standard approaches that only supervise the final output, this loss function computes the combined loss at multiple intermediate decoder stages.
Weighted Aggregation: The total loss is a weighted sum of losses from intermediate and final outputs ( $L_{total} = \sum \lambda_k L_k$ ). This design promotes better gradient flow throughout the network and ensures structural consistency in the segmentation of fine boundaries.

Key Contributions

The paper outlines four primary contributions:

Grouped Query Attention in the Encoder: Integration of GQA to reduce computational complexity while effectively modeling long-range dependencies, addressing the inefficiency of standard Transformers.
Channel Attention in the Decoder: Implementation of CA to focus on salient feature channels, thereby improving the boundary delineation between the optic disc and cup.
Aggregated Pyramid Loss: A novel loss function that supervises multiple decoder stages to enhance gradient flow and fine detail segmentation.
Efficient Hybrid Design: A model that leverages the complementary strengths of UNet, Transformers, and attention mechanisms to achieve high accuracy with a significantly reduced parameter count.

Experimental Results

The model was evaluated on the REFUGE2 and ORIGA datasets.

Quantitative Performance

REFUGE2 Dataset: OptiModNet achieved State-of-the-Art (SOTA) performance, surpassing existing methods by over 2.5%.
- Optic Disc: Dice Score of 99.45% and IoU of 98.16%.
- Optic Cup: Dice Score of 93.23% and IoU of 89.41%.
- These results exceeded the previous best method (MedSegDiff-V2) by 2.75% (Dice) and 2.26% (IoU) for disc segmentation.
ORIGA Dataset: The model achieved the highest reported scores for 3-class classification (optic cup, optic disc, background), with an IoU of 98.92% and a Dice score of 99.31%.

Efficiency and Complexity

Parameters: OptiModNet utilizes only 1.93 million parameters, representing a 96% reduction compared to complex SOTA models like MedSegDiff-V2 (46M parameters).
Computational Cost: The model requires only 3.73 GFLOPs, which is approximately 0.38% of the computational cost of MedSegDiff-V2 (983 GFLOPs).

Ablation Studies

Component Analysis: Removing either Channel Attention or Grouped Query Attention resulted in a measurable drop in performance, confirming that both modules contribute meaningfully to the final accuracy.
Loss Function: Varying the weights ( $\lambda$ ) in the Aggregated Pyramid Loss demonstrated that the specific weighting scheme used (0.2, 0.3, 0.5 for the three intermediate stages and 1.0 for the final output) yielded optimal results.

Significance and Claims

The paper claims that OptiModNet successfully addresses the trade-off between segmentation accuracy and computational efficiency. By integrating Grouped Query Attention and Channel Attention within a UNetR-based framework, the model achieves superior boundary delineation and global context modeling without the heavy computational burden of standard Transformers or Diffusion models.

The authors emphasize that the model's high efficiency (low parameters and GFLOPs) makes it particularly suitable for rapid, large-scale screening and deployment in resource-limited clinical environments, where high-resolution fundus imaging often faces hardware constraints. The results suggest that OptiModNet provides a practical, high-performance solution for automated glaucoma screening, reducing inter-observer variability and facilitating early diagnosis.