KAE: Kolmogorov-Arnold Auto-Encoder for Representation Learning

Scientific quality:

It’s unclear to what extent are the improvements due to model selection, especially as the kind of activation function used was especially selected for these tasks. Moreover, how this selection was done isn’t clear - was it based on the validation accuracy? The paper doesn’t elaborate and as it currently is seems to indicate that it was based on the test accuracy as that’s the only metric which is given.

It’s unclear if the MLP and KAE have identical computational resources - how many parameters do each have? How many FLOPs? This is somewhat indicated in Figure 4, based on which it seems that the models aren’t equally powerful, at least using these simple heuristic comparisons.

As the models have different sizes but the same hyperparameter setup was used for all, it’s unclear whether improvements are chiefly due to better inductive biases or these hyperparameters suiting some more than others, eg. as bigger models overfit more or smaller ones have implicit regularisation. As the models and datasets are modestly sized a small hyperparameter search over each setup would go a long way.

Based on section 4.5, is the data standardised before feeding it into the models? If not then why? This is a standard practice that could significantly affect learning.

For section 4.4, is the runtime comparison justified given the relatively small batch sizes and different number of parameters? It’s difficult to understand how it would extrapolate to larger, more realistic settings where the GPU is fully utilised. A single forward pass’ runtime with a big batch would be more indicative.

Clarity/presentation:

The paper has imprecise statements at time, eg. “does not produce satisfactory results” (line 193).

At times the paper uses unscientific language, eg. “this remarkable theorem” (133).

Information is sometimes overly repeated or overly detailed, eg. there’s no need to say so explicitly what an autoencoder is (lines 98-103) and KANs are in practice presented several times in similar ways (lines 46-49, 164-166, 174-176, etc.).

Lots of information can be delegated to the appendix or removed as it can be considered well known/irrelevant, eg. table 1.

Novelty:

Although it’s generally worthwhile using different architectures in new domains, generally there’s a clear benefit or problem that doing so would solve. For example, [1] used transformers for vision to a) get better compute scaling and b) show that large-scale data is more important than inductive biases. While some potential deeper improvements are stated here, eg. “a more robust model for representation learning” (line 51), “providing a theoretical basis for improving AE performance” (59), “KANs have been shown to outperform MLPs in various applications (Xu et al., 2024; Bozorgasl & Chen, 2024), particularly those requiring complex, non-linear transformations of input data” (116-118), etc., these seem speculative or insufficiently rigorous as it’s unclear how the paper shows these points.

1. The experiments in the paper are insufficient to substantiate its claims (e.g., "These findings position KAE as a ‘practical’ tool for high-dimensional data analysis" and "KAE effectively captures complex data patterns"). The architecture used is too minimal to be considered empirically significant, especially given the simplicity of the datasets used. Although the authors provide experiments on four datasets, quality is more important than quantity.

2. Using learned activation functions introduces extra computational overhead compared to standard autoencoders with the same architecture but fixed activation functions. There is no section detailing this computational overhead (e.g., GPU memory usage).

• Originality:

One of the major issue of this paper is the lack of originality. The KAE model is simply replacing MLP by KAN in auto-encoder with a simple change of kernel from B-spline to polynomial.

The auto-encoder network is widely used and simple to implement while KAN is novel, it is not the contribution of this paper. It is expected that to publish in ICLR we need to either propose a novel idea that would challenge the existing belief or we should present the new system that beat the state-of-the-art. Given that the plain two-layer auto-encoder is not the state-of-the-art in either CIFAR or MINST and the two-layer architecture is trivial, we are expecting that the authors would bring new insights from either theory or in implementation point of view to show why KAE is able to challenge MLP. Unfortunately, due to the limited choice of model architecture and experiment datasets (MINST and CIFAR is hard representative of the benchmark in the ML community), this paper is not achieving this goal yet.

• Clarity: - The plot is small and the line in Figure 2,3 cannot be distinguished in grey-scale printout since the markers and line styles are the same. It is suggested to modify markers and line styles so that it can be visualized without relying on colored pdf file. - The interpretability result only covers the MINST, which is a toy dataset. It is not very convincing. On the other hand, it is not convincing that KAE is more interpretable than MLP-based AE. Both has poor interpretability but AE at least maps linearly locally to Euclidean space but KAE maps to feature space via high dimensional kernels, which is hard to interpret.

• Significance:

  • The key challenge for KAE or KAN framework is to convince the other researchers that it worth the restructuring of the basic building block. In high-level, KAN layer is a kernel non-parametric regression model, which is in contrast to the MLP layer which is a parametric model. - The benefit of non-parametric regression is that it learns from a function space which is very high dimensional compared to original parametric space. - The drawback of this framework, however, lies in the choice of kernel. Like the difficulty we faced in early 2000s, kernel machine need to select a family of kernels apriori, which is hand-crafted. KAE, KAN with B-spline, wavelet, fourier basis are all human crafted feature family. This is the same issue for authors. For your use case which is MINST and CIFAR, polynomial kernel may be good enough. But how about NLP tasks? how about ImageNet, how to choose a family of kernel? in fact, it is critical to choose a good family of kernels for this type of system to work the best. But there is no guidance on that.
  • Another issue is that why only use shadow network? is there limitation for KAN to go deep? AE can benefit strongly by going deep. This is a fundamental question for KAN, since each layer is very expressive, it seems that there is little to learn for the second KAN layer on top. This paper could demonstrate the performance gain by constructing a 6 layer + network, if possible.
  • Finally, it is the question on whether or not the KAN suffers from vanishing and exploding gradient issue more than MLP. Kernel such as polynomial kernel is not stationary. In a sense, in a long run, the error in gradient computation will grew much faster than ReLU activation. This is a critical issue if we want to go deep.

• Claims like demonstrating "superiority of the Kolmogorov-Arnold Auto-Encoder (KAE) through extensive experimental validation" are not supported by the empirical evidence because the experiments are conducted exclusively with shallow auto-encoder baselines that are quite unrealistic and not used in practice.
• Too strong claims on the power of KANs themselves like stating that MLPs "still rely on fixed activation functions, which limits their flexibility in representing complex functions". This is neither practically nor theoretically true and is also just stated without even a reference. With enough layers and/or width, MLPs should be equally powerful.
• Insufficient experimental validation: the experiments are conducted on standard benchmarks but with extremely small architectures (3 layers) and hidden dimensionality of 16 or 32. Further, image datasets should employ a convolutional architecture instead of MLPs, at least as a point of comparison what is achievable.

Overall, I think the paper could be improved by toning down claims and improving the span of experiments. Currently, while the paper is very well written and illustrated, I find it scientifically questionable due to above mentioned claims and limited experimental validation. The paper hinges on practicality, which has not been verified by the limited scale experiments. I am open to changing my score in case these two major concerns are addressed.