KAE: Kolmogorov-Arnold Auto-Encoder for Representation Learning

2024-12-31

AI Summary

Traditional autoencoders suffer from limited capacity in modeling complex nonlinear relationships and lack structural interpretability in learned representations. To address these limitations, we propose the Kolmogorov–Arnold Auto-Encoder (KAE), the first autoencoder framework integrating a learnable-edge Kolmogorov–Arnold Network (KAN). KAE replaces conventional neuron-wise activations with learnable piecewise polynomial functions on edges, explicitly capturing high-order nonlinear interactions. This design achieves both strong representational power and inherent structural interpretability, while supporting end-to-end joint optimization. Extensive experiments across multiple benchmark datasets demonstrate that KAE significantly reduces reconstruction error and consistently outperforms standard autoencoders and existing KAN-based variants on downstream tasks—including retrieval (measured by mAP), classification (accuracy), and denoising (PSNR). These results validate KAE’s dual improvements in latent-space modeling fidelity and generalization capability.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Learning to Search
• Computer Vision: Representation Learning for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

The Kolmogorov-Arnold Network (KAN) has recently gained attention as an alternative to traditional multi-layer perceptrons (MLPs), offering improved accuracy and interpretability by employing learnable activation functions on edges. In this paper, we introduce the Kolmogorov-Arnold Auto-Encoder (KAE), which integrates KAN with autoencoders (AEs) to enhance representation learning for retrieval, classification, and denoising tasks. Leveraging the flexible polynomial functions in KAN layers, KAE captures complex data patterns and non-linear relationships. Experiments on benchmark datasets demonstrate that KAE improves latent representation quality, reduces reconstruction errors, and achieves superior performance in downstream tasks such as retrieval, classification, and denoising, compared to standard autoencoders and other KAN variants. These results suggest KAE's potential as a useful tool for representation learning. Our code is available at url{https://github.com/SciYu/KAE/}.

Problem

• Machine Learning Optimization
• Nonlinear Pattern Recognition
• Feature Quality Enhancement

Innovation

• Kolmogorov-Arnold Autoencoder
• Learnable Activation Functions
• Enhanced Representation Learning