Source-linked AI summary
A KAN-based Non-linear Factorization for Embedding Layer Compression
Oliver Giudice, Anna Russo Russo
TL;DR
Large vocabularies make embedding layers costly, while standard low-rank compression uses a linear factorization and KAN applications to NLP remain scarce. This paper integrates a KAN mapping into a factorized embedding layer and tests SVD-based initialization strategies in Skip-gram. After adaptation, KAN-based variants outperform standard low-rank factorization on average, and the paper identifies training stability as an ongoing challenge.
Problem
Embedding layers become parameter- and memory-intensive as vocabularies grow, while standard low-rank factorization is linear and KAN use in NLP remains scarce.
Source: §1 · pp. 1–2Method
Results
After three epochs of adaptation, all KAN-based variants outperform standard low-rank factorization on average, with KF-SVD+ attaining an overall average reduction of 11.7%.
Source: §6 · p. 7Takeaways & Limitations
The results support KAN-based factorized embeddings as a viable, competitive alternative to traditional low-rank factorization under compression.
Source: §7 · p. 8Training stability remains a central challenge, and the study calls for more systematic understanding of how initialization, optimization hyperparameters, and fine-tuning interact with KAN architectures.
Source: §7 · p. 8
Abstract
from arXiv · showhide
Among existing embedding compression approaches, low-rank factorization provides a simple and mathematically principled solution by approximating the embedding matrix with a reduced-rank decomposition. However, its practical effectiveness is often limited, as the resulting linear bottleneck constrains representational capacity. To overcome these limitations, we propose an enhanced factorization scheme that increases expressive power by introducing learnable non-linear projections within the compressed embedding layer. Our method is inspired by the Kolmogorov Arnold Networks (KAN) paradigm, which replaces fixed linear transformations with learnable univariate functions. Despite promising results in other domains, applications of KANs in NLP remain scarce and have so far been largely limited to final or task specific layers, rather than core representational components such as embedding layers. We evaluate the proposed approach by integrating the compressed embedding layer into the Skip-gram framework and comparing it against a standard low-rank baseline under matched parameter budgets on the Word Analogy and Semantic Similarity benchmarks. We further introduce an effective initialization strategy from a pre-trained, uncompressed embedding layer. The behavior of the model is analyzed both immediately after initialization and after a brief post-initialization refinement phase, considering training with the standard Skip-gram objective as well as variants augmented with distillation from the original embedding layer. Overall, our results show that augmenting low-rank factorization with learnable non-linear projections effectively mitigates the limitations of linear bottlenecks, yielding more faithful compressed embeddings under fixed dimensionality and suggesting a promising direction for embedding layer compression beyond purely linear methods.
1. Introduction
Large vocabularies make embedding-layer parameters and memory a scalability bottleneck, while standard low-rank compression remains strictly linear. The paper proposes KAN-based factorization and evaluates it in Skip-gram against standard low-rank factorization.
1. Introduction
- Large vocabularies make standard embedding layers’ parameter count and memory footprint a scalability bottleneck, especially in memory-constrained settings.Low-rank factorization is one of several compression approaches, but the standard factorization decomposes the embedding matrix into two smaller matrices linearly.Source: §1 · pp. 1–2
- The proposed factorization replaces the second linear projection with a KAN mapping that sums learned nonlinear transformations of individual input coordinates.KANs use learnable nonlinear univariate functions.Source: §1 · p. 2
- The paper introduces KAN-based factorized embeddings and two initialization schemes, evaluating them against standard low-rank factorization within Skip-gram.The authors state that KAN-style learnable activations had not previously been explored for embedding compression, to their knowledge.Source: §1 · p. 2
2. Related Work
Related work spans KAN applications across domains and architectural components, alongside concerns about the scope of reported advantages. In NLP, the cited applications are scarce and include task-specific classification heads.
2. Related Work
- KANs have been applied across time-series, vision, reinforcement learning, speech, and physics-informed modeling, often replacing fully connected blocks in standard architectures.Related work also explores alternative univariate function parameterizations, aggregation functions, and strategies for regularization or parameter reduction.Source: §2 · p. 2
- Some studies question comparison fairness and suggest KAN gains may be confined to smooth regular functions, small-to-medium models, or symbolic regression settings.The passage presents these as concerns raised by recent studies, not as settled boundaries.Source: §2 · pp. 2–3
- KAN applications in NLP are described as scarce, with cited examples including KAN classification heads for several language-model tasks.The passage suggests input nonsmoothness and training instability may make deployment in state-of-the-art models difficult.Source: §2 · p. 3
3. Background
KAN layers map input coordinates to outputs by summing learnable univariate functions. The original formulation uses spline-based parametrizations, while the layer definition permits other functional families.
3. Background
- A KAN layer maps an input vector to an output vector by associating a learnable univariate function with each input-output connection and summing contributions.The layer can also be viewed as a matrix of univariate functions applied coordinate-wise and summed across inputs.Source: §3 · p. 3
- A KAN with multiple layers is formed by composing the nonlinear mappings defined for its individual layers.The passage identifies each Φ(l) as the mapping defined by the layer equation.Source: §3 · p. 3
- In the original formulation, each activation function is parametrized as a sum of a base function and a spline, but the layer definition is not tied to one functional family.The spline is described as learnable, with its order and grid specified in the background passage.Source: §3 · p. 3
4. Methodology
The method keeps a rank-r factorized embedding structure but replaces its second linear mapping with a KAN layer. Two SVD-based initialization strategies provide different ways to initialize the KAN parameters.
4. Methodology
- The proposed layer retains a rank-r factorization but replaces the second linear mapping with a KAN transformation having r inputs and d outputs.The figure contrasts this layer with both a full-rank embedding layer and a low-rank factorized layer.Source: §4 · pp. 3–4, Figure 1 · p. 4
- Both initialization strategies begin with an SVD of the pre-trained embedding matrix and scale its factors before initializing the first linear transformation.The scaling uses A′ = αA and B′ = B/α, with α > 0.Source: §4 · p. 4
- KAN-SVD initializes the KAN activations to approximate linear functions with slopes from B′, aiming to reproduce the SVD-based low-rank behavior initially.The activations can be fitted by least-squares regression over points in the activation interval.Source: §4 · pp. 4–5
- KAN-BCD alternates updates to the KAN parameters and W to minimize vocabulary-wide reconstruction loss, stopping when relative improvement falls below a threshold.KAN parameters are updated with ridge-regularized least squares; W is then optimized by stochastic gradient descent, and the strategy does not use B.Source: §4 · p. 5
5. Experimental setup
The experiments compare KAN-based and standard low-rank factorized embeddings in Skip-gram, measuring information loss immediately after compression and after adaptation. Evaluation uses analogy and semantic-similarity tasks, with adaptation using either the Skip-gram objective alone or feature distillation.
5. Experimental setup
- Models are compared at the same factorization rank, before and after fine-tuning on the full-rank model’s training data, using either Skip-gram loss or feature distillation.The compressed embeddings are derived from a pre-trained full-rank Skip-gram model.Source: §5 · p. 5
5.1. Training data
- Training uses 120 million Wikipedia words, yielding approximately 37,000 tokens after filtering; each target-context pair has three frequency-sampled negative examples.The context window is 2, with up to four target-context pairs collected per occurrence.Source: §5.1 · p. 5
5.2. Hyperparameters and implementation details
- The KAN layer uses Gaussian radial basis functions on 10 uniformly spaced grid points in the activation interval [−1,1].The grid points are fixed, while the RBF weights are learnable.Source: §5.2 · p. 6
- For KAN-SVD, the RBF layer is initialized to emulate the linear transformation specified by the SVD-derived matrix B′.The implementation interpolates the identity function at grid points and applies the resulting weights to B′.Source: §5.2 · p. 6
- Post-compression adaptation uses Adam at learning rate 10^-4; with distillation, a cosine-similarity penalty is added with λ = 0.3.The distillation penalty compares each current embedding with its original full-rank embedding.Source: §5.2 · p. 6
5.3. Embedding evaluation
- Analogy is evaluated with 3CosAdd on Google Analogy, while WordSim-353 and SimLex-999 use Spearman correlation.The Google Analogy set contains 19,544 semantic and syntactic questions.Source: §5.3 · p. 7
6. Results
After brief adaptation, all KAN-based variants outperform standard low-rank factorization on average, and KF-SVD+ has the smallest reported average reduction. Immediately after compression, KF-BCD retains more performance than low-rank factorization, though SimLex-999 shows a distinct pattern.
Immediately after compression
- 16.4% average performance reduction for KF-BCD is lower than LF’s 28.8% reduction immediately after compression.On analogy tasks and WordSim-353, the average reductions are 19% for KF-BCD and 39% for LF.Source: §6 · p. 7, Table 2 · p. 7
- LF and KF-SVD slightly exceed the full-rank model on SimLex-999, a pattern that persists after adaptation without increased correlation.The authors suggest a corpus–dataset mismatch may explain the pattern.Source: §6 · p. 7
Adaptation
- 11.7% overall average reduction is achieved by KF-SVD+ after adaptation, and all KAN-based variants outperform standard low-rank factorization on average.The comparison uses the best scores within three adaptation epochs.Source: §6 · p. 7, Table 3 · p. 8
- Distillation is generally beneficial for KAN-based variants but consistently harms adaptation for LF embeddings.The authors suggest KAN functional expressivity supports recovery of task-relevant structure, while BCD initialization may yield a less favorable optimization landscape.Source: §6 · p. 8
7. Conclusion and future work
KAN-based factorized embeddings are a viable, competitive alternative to traditional low-rank factorizations, recovering substantial expressiveness lost through compression. The conclusion highlights complementary initialization strengths, ongoing training-stability challenges, and future evaluation in downstream NLP tasks and larger architectures.
7. Conclusion and future work
- KAN-based factorized embeddings are a viable, competitive alternative to traditional low-rank factorizations, with learnable nonlinearities recovering substantial expressiveness lost through compression.The paper frames non-linear factorization as a promising direction for more expressive embedding compression without sacrificing parameter efficiency.Source: §7 · p. 8
- KAN-SVD recovers information directly from full-rank embeddings, while KAN-BCD offers a more effective starting point for post-compression adaptation or possibly downstream fine-tuning.Heterogeneous behavior across post-compression adaptation strategies makes training stability a central challenge.Source: §7 · p. 8
- Future work will evaluate whether intrinsic-evaluation advantages transfer to downstream NLP tasks and test compatibility with larger architectures, especially Transformers.Source: §7 · p. 8
- KAN-based non-linear factorization is a promising direction for more expressive embedding compression without sacrificing parameter efficiency.Source: §7 · p. 8