IVF Image Analysis with DINO and Kolmogorov-Arnold Networks (KAN)

In vitro fertilization (IVF) research demands sophisticated approaches for analyzing complex biomedical images. While conventional convolutional neural networks (CNNs) have shown effectiveness in many areas of image recognition, they often struggle to capture the nuanced global context and fine-grained details crucial for embryo assessment. To address these limitations, we've developed an innovative neural network model that combines the power of DINO (Distillation with No Labels) for feature extraction with Kolmogorov-Arnold Networks (KAN) for classification.

Our approach leverages the strengths of transformer-based architectures through DINO, which utilizes Vision Transformers (ViT) to divide images into small patches and learn relationships between them using attention mechanisms. This allows for the capture of both local details and global context more effectively than traditional CNNs. The pre-trained DINO models extract high-quality, global contextual features that are particularly beneficial in medical image analysis, where intricate patterns are crucial for accurate assessment.

The extracted features are then processed by Kolmogorov-Arnold Networks, a specialized neural network architecture designed to handle highly non-linear relationships between variables. Based on the Kolmogorov-Arnold representation theorem, KAN provides a mathematically grounded approach to analyzing complex feature relationships, offering robust classification capabilities that outperform traditional neural networks in handling the intricacies of embryo images.

In our testing, this ensemble approach achieved remarkable results in identifying whether embryo images are "in focus," a critical factor in embryo quality assessment for IVF procedures. Our model demonstrated an accuracy of 96.24%, precision of 94.34%, recall of 79.37%, and an impressive ROC AUC of 99.12%. These metrics underscore the effectiveness of combining transformer-based feature extraction with specialized neural network architectures for handling highly non-linear relationships in medical image analysis.

Why Transformers and Not CNNs?

CNNs have dominated the field of image recognition, excelling at detecting local patterns within images. However, when it comes to tasks that require an understanding of the entire image context, like embryo classification in IVF, CNNs may miss out on important global structures. IVF images present a challenge due to their fine details and intricate features that require more holistic analysis. This is where transformer-based architectures shine.

DINO (Distillation with No Labels) - The Power of Transformers

DINO is a self-supervised learning approach based on Vision Transformers (ViT). Unlike CNNs, ViTs divide an image into small patches and learn relationships between them using attention mechanisms. This allows ViTs to capture both local details and global context more effectively than CNNs.

DINO’s pre-trained models extract high-quality, global contextual features that are especially beneficial in medical image analysis, where intricate patterns are important. For IVF image analysis, where identifying the focus and quality of embryo images is crucial, DINO allows us to leverage self-supervised learning to extract features without the need for labeled datasets.

KAN (Kolmogorov-Arnold Networks) - Handling Complex Feature Relationships

Once DINO extracts meaningful image features, the next step is classifying these features effectively. This is where Kolmogorov-Arnold Networks (KAN) come into play.

KANs stand out as a powerful approach that leverages the fundamental mathematics of function representation to enhance the accuracy and interpretability of image classification tasks. At the heart of KAN lies the Kolmogorov-Arnold representation theorem, a mathematical principle that demonstrates how any continuous function of multiple variables can be reconstructed using only continuous functions of one variable and addition operations. This theorem, expressed mathematically as F(x1, ..., xn) = ∑[j=1 to 2n+1] gj(∑[i=1 to n] φij(xi)), forms the theoretical foundation for a neural network architecture that excels at capturing complex, non-linear relationships in data.

In our implementation of KAN for IVF image analysis, we've translated this mathematical principle into a practical neural network structure:

```python

class KAN(nn.Module):
    def __init__(self, width, grid, k):
        super(KAN, self).__init__()
        self.fc1 = nn.Linear(width[0], width[1])
        self.fc2 = nn.Linear(width[1], width[2])

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

kan_model = KAN(width=[X_train.shape[1], 10, 1], grid=5, k=3)
```        

This implementation offers several key advantages in the context of IVF image analysis. The architecture's ability to handle complex non-linear relationships makes it particularly well-suited for capturing the subtle features and patterns present in embryo images. Moreover, the mathematical foundation of KAN provides a level of interpretability that is crucial in medical applications where understanding the model's decision-making process is essential. The network's efficiency in terms of parameter count leads to faster training times and a reduced risk of overfitting, which is particularly valuable when working with limited medical datasets.

Looking ahead, there are several promising directions for further development of KAN in medical imaging applications. We are exploring possibilities for optimizing the architecture to reduce parameter count while maintaining or improving performance, as well as experimenting with different activation functions to enhance the network's capabilities. The potential applications of KAN extend beyond IVF image analysis to other areas of medical imaging and biomedical data processing, where its unique combination of mathematical rigor and practical effectiveness could prove valuable.

The successful application of Kolmogorov-Arnold Networks in IVF image analysis demonstrates the power of combining advanced mathematical principles with modern deep learning techniques. KAN represents a significant step forward in the field of medical image analysis. As we continue to refine and expand upon this approach, the potential for improving medical diagnostics and decision-making through advanced AI techniques becomes increasingly promising.

In our case, the KAN model is well-suited to learn and predict based on the intricate patterns that DINO extracts from the images. Compared to CNNs, which struggle with the complexity of these features, KAN provides a structured and mathematically grounded approach to analyze the relationships between image features.

Combining DINO and KAN for IVF Image Classification

By integrating the powerful feature extraction capabilities of DINO with the flexibility and robustness of KAN, we have created a neural network architecture that significantly outperforms traditional CNNs in IVF image analysis tasks.

Our model successfully identifies whether embryo images are "in focus," a critical factor in embryo quality assessment for IVF procedures. During testing, the model achieved impressive metrics:

- Accuracy: 96.24%

- Precision: 94.34%

- Recall: 79.37%

- F1 Score: 86.21%

- ROC AUC: 99.12%

- MCC: 84.47%

These results demonstrate the effectiveness of combining a Vision Transformer with a specialized neural network architecture for handling highly non-linear relationships. The model outperformed CNN-based approaches, showcasing its superiority in capturing the necessary details for IVF image analysis. The implementation of this combined approach presented several challenges that required careful consideration. Hyperparameter optimization was crucial, particularly in tuning the dimensions of hidden layers. We also focused on computational efficiency, optimizing both training and inference processes through effective batch processing and careful implementation of the network architecture.

The Future of AI in IVF Image Analysis

The combination of transformers like DINO and novel neural network architectures like KAN opens new possibilities for medical image classification. By moving beyond traditional CNNs, we can achieve more accurate, efficient, and meaningful predictions in complex domains like IVF research.

This project is a testament to the transformative potential of self-supervised learning and transformer models in biomedical applications. As AI continues to evolve, we expect further improvements in IVF image analysis, unlocking new insights.

---

References

- DINO (Self-Supervised Learning via Knowledge Distillation)

You can learn more about DINO and access the pre-trained models via the [official DINO repository](https://github.com/facebookresearch/dino).

- Kolmogorov-Arnold Networks (KAN)

The KAN model, used in this project, is described in the paper by Liu et al.:

Liu, Ziming, et al. "KAN: Kolmogorov-Arnold Networks." arXiv preprint arXiv:2404.19756 (2024).

Find the KAN implementation on [GitHub](https://github.com/rotem154154/kan_classification).

- Blastocyst Dataset

For independent testing, we used the [Blastocyst Dataset](https://github.com/software-competence-center-hagenberg/Blastocyst-Dataset).

---

Acknowledgments

This project would not have been possible without the contributions from the authors of the DINO model, KAN, and the Blastocyst Dataset. We thank them for making their work available to the community.