KAN là gì? Mạng Kolmogorov-Arnold khác MLP thế nào?

Kolmogorov-Arnold Network (KAN) là một kiến trúc mạng nơ-ron trong đó mỗi cạnh không chỉ chứa một trọng số vô hướng mà mang một hàm một biến có thể học được. Đây là khác biệt cốt lõi so với Multi-Layer Perceptron (MLP), nơi các cạnh chủ yếu thực hiện phép nhân với trọng số còn hàm kích hoạt thường được áp dụng tại các node.

KAN được lấy cảm hứng từ định lý biểu diễn Kolmogorov-Arnold và đặc biệt đáng chú ý trong các bài toán xấp xỉ hàm, mô hình hóa khoa học và symbolic regression. Tuy nhiên, KAN không nên được hiểu là kiến trúc đã thay thế MLP trong mọi bài toán Deep Learning; khả năng biểu diễn, chi phí tính toán và mức độ phù hợp của hai kiến trúc khác nhau theo từng nhiệm vụ.

Tóm tắt nội dung

• MLP học chủ yếu trọng số trên các cạnh, còn KAN đặt hàm một biến học được trên từng cạnh.
• Một KAN layer cộng kết quả của nhiều hàm cạnh để tạo activation của node tiếp theo.
• KAN gốc sử dụng B-spline để tham số hóa các hàm phi tuyến trên cạnh.
• KAN có lợi thế đáng chú ý về trực quan hóa, pruning và symbolic regression, nhưng không mặc định nhanh hơn hoặc tốt hơn MLP.

MLP học các trọng số vô hướng trên cạnh, trong khi KAN thay các kết nối bằng những hàm một biến có thể học được.

Nền tảng toán học của Kolmogorov-Arnold Network

KAN được lấy cảm hứng từ định lý biểu diễn Kolmogorov-Arnold. Một cách phát biểu thường được sử dụng để giải thích ý tưởng cho biết một hàm liên tục nhiều biến trên miền thích hợp có thể được biểu diễn thông qua phép cộng và các hàm liên tục một biến.

Với một hàm có n biến đầu vào, dạng biểu diễn có thể viết:

$$f(x_1,\dots ,x_n)=\sum _{q=1}^{2n+1}{\mathrm{\Phi }}_q\left(\sum _{p=1}^n{\phi }_{q,p}\right(x_p\left)\right)$$

Điểm đáng chú ý trong biểu thức này là một hàm nhiều chiều được xây dựng từ các hàm một biến kết hợp với phép cộng. Đây là nguồn cảm hứng trực tiếp cho KAN: thay vì chỉ học các hệ số vô hướng giữa các node, ta học những phép biến đổi phi tuyến riêng biệt trên từng cạnh.

Tuy nhiên, định lý Kolmogorov-Arnold không đồng nghĩa hoàn toàn với kiến trúc KAN hiện đại. Biểu diễn cổ điển có cấu trúc khá đặc thù, trong khi KAN tổng quát hóa ý tưởng thành nhiều layer có thể xếp chồng để tăng khả năng biểu diễn.

Điểm cần nhớ: định lý Kolmogorov-Arnold là nền tảng ý tưởng. Nó không tạo ra bảo đảm rằng một KAN bất kỳ sẽ tự động hội tụ nhanh hơn hoặc đạt kết quả tốt hơn MLP trên mọi dataset.

KAN Layer khác MLP Layer như thế nào?

Để thấy rõ sự khác biệt, trước tiên xét một tầng MLP quen thuộc. Ta thực hiện phép biến đổi tuyến tính rồi áp dụng một hàm kích hoạt:

$$x^{\prime }=\sigma (Wx+b)$$

Ma trận W chứa các trọng số học được. Hàm σ thường là ReLU, GELU, tanh, sigmoid hoặc một activation cố định khác.

Trong KAN, mỗi cạnh từ đầu vào i đến đầu ra j mang một hàm riêng ký hiệu là φ_j,i. Giá trị tại node đầu ra j được tạo bằng cách cộng đầu ra của tất cả các hàm cạnh:

$$x_j^{\prime }=\sum _{i=1}^n{\phi }_{j,i}\left(x_i\right)$$

Đây là công thức quan trọng nhất để hiểu KAN. Node không cần áp dụng một activation chung sau phép cộng; phần phi tuyến đã nằm trong chính các hàm φ trên cạnh.

Ví dụ, với hai đầu vào và một đầu ra:

$$y={\phi }_1\left(x_1\right)+{\phi }_2\left(x_2\right)$$

Nếu dùng MLP, ảnh hưởng của x₁ ở một kết nối tuyến tính chủ yếu được điều khiển bởi trọng số w₁. Với KAN, φ₁ có thể học một đường cong phi tuyến hoàn chỉnh theo x₁.

Thành phần	MLP	KAN
Cạnh	Trọng số vô hướng	Hàm một biến học được
Node	Tổng tuyến tính rồi activation	Cộng các tín hiệu từ cạnh
Phi tuyến	Thường cố định ở node	Học được trên từng cạnh
Diễn giải	Một weight riêng lẻ thường khó đọc trực tiếp	Có thể quan sát đường cong φ(x)

B-spline giúp KAN học một hàm trên mỗi cạnh

Nói rằng mỗi cạnh chứa một hàm có thể học được dẫn tới một câu hỏi quan trọng: làm thế nào biểu diễn và tối ưu cả một hàm?

KAN gốc sử dụng B-spline. Một spline có thể được biểu diễn bằng tổ hợp tuyến tính của nhiều hàm cơ sở B_i(x):

$$\mathrm{spline}\left(x\right)=\sum _i^{}{c}_i{B}_i\left(x\right)$$

Các B_i(x) là các hàm cơ sở B-spline được xác định bởi grid và bậc spline, còn c_i là các hệ số có thể học bằng Backpropagation.

Trong KAN gốc, hàm trên cạnh còn được kết hợp với một hàm cơ sở b(x). Cấu trúc activation được viết:

$$\phi \left(x\right)=w\left[b\right(x)+\mathrm{spline}(x\left)\right]$$

Một hàm cơ sở thường được sử dụng trong KAN gốc là SiLU:

$$b\left(x\right)=\mathrm{SiLU}\left(x\right)=\frac{x}{1+e^{-x}}$$

Nhờ spline, một cạnh không còn bị giới hạn bởi phép biến đổi tuyến tính. Trong quá trình huấn luyện, mô hình có thể điều chỉnh các hệ số c_i để thay đổi hình dạng của hàm φ(x).

Trọng số tuyến tính chỉ điều chỉnh độ lớn và dấu của kết nối, còn một hàm spline có thể học quan hệ phi tuyến phức tạp hơn theo miền đầu vào.

Grid và bậc spline

Spline được xây dựng trên một grid chia miền giá trị đầu vào thành nhiều khoảng. Grid càng mịn, spline càng có nhiều bậc tự do để biểu diễn các biến thiên chi tiết.

Nếu ký hiệu số khoảng grid là G và bậc spline là k, số basis function trên mỗi cạnh tăng theo G và k. Vì vậy, tăng grid cũng đồng nghĩa tăng số tham số và chi phí tính toán.

Trong nhiều ví dụ của KAN, spline bậc ba được sử dụng, tương ứng k = 3. Đây thường là lựa chọn tạo đường cong đủ trơn mà vẫn giữ chi phí hợp lý.

Không nên tăng grid ngay từ đầu. Một chiến lược hợp lý là bắt đầu với grid nhỏ để học cấu trúc tổng quát, sau đó refine grid nếu mô hình cần biểu diễn chi tiết hơn.

Mạng KAN nhiều tầng được xây dựng như thế nào?

KAN không bị giới hạn ở biểu diễn hai tầng của định lý Kolmogorov-Arnold. Ta có thể xếp chồng nhiều KAN layer tương tự cách xếp nhiều tầng trong MLP.

Gọi Φ_l là toàn bộ tập hàm trên các cạnh của layer l. Khi đó:

$$x_{l+1}={\mathrm{\Phi }}_l\left(x_l\right)$$

Một mạng gồm L layer có thể biểu diễn bằng phép hợp thành:

$$\mathrm{KAN}\left(x\right)=({\mathrm{\Phi }}_{L-1}\circ {\mathrm{\Phi }}_{L-2}\circ \cdots \circ {\mathrm{\Phi }}_0)\left(x\right)$$

Ví dụ một cấu hình:

[2, 5, 1]

có thể được hiểu là hai biến đầu vào, năm node ở tầng trung gian và một đầu ra. Điểm quan trọng là kết nối giữa các tầng là các hàm φ_j,i, không chỉ là scalar weight.

KAN vẫn là một mô hình khả vi. Các hệ số spline được tối ưu thông qua Backpropagation và các optimizer. Kiến trúc KAN thay đổi cách tham số hóa phép biến đổi, chứ không loại bỏ cơ chế tính gradient.

Thực hành KAN bằng Python

Một cách trực quan để hiểu KAN là dùng nó cho bài toán xấp xỉ hàm hai biến:

$$f(x,y)=\exp \left(\sin \right(\pi x)+y^2)$$

Ví dụ sau sử dụng thư viện pykan:

import torch
from kan import KAN, create_dataset

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)


# f(x, y) = exp(sin(pi*x) + y^2)
def f(x):
    return torch.exp(
        torch.sin(torch.pi * x[:, [0]])
        + x[:, [1]] ** 2
    )


dataset = create_dataset(
    f,
    n_var=2,
    device=device
)

# 2 input → 1 hidden node → 1 output
model = KAN(
    width=[2, 1, 1],
    grid=3,
    k=3,
    seed=1,
    device=device
)

results = model.fit(
    dataset,
    opt="LBFGS",
    steps=20
)

model.plot()

Cấu hình width=[2,1,1] tạo hai đầu vào, một node trung gian và một đầu ra. grid=3 đặt grid ban đầu tương đối nhỏ, còn k=3 sử dụng spline bậc ba.

Sau huấn luyện, model.plot() cho phép trực quan hóa mạng. Khả năng quan sát trực tiếp các activation trên cạnh là một trong những đặc điểm thú vị nhất của KAN khi mục tiêu không chỉ là dự đoán mà còn muốn phân tích mô hình.

Refine grid sau khi đã học cấu trúc ban đầu

Nếu grid ban đầu quá thô, có thể tăng độ phân giải spline sau khi mô hình đã học được cấu trúc tổng quát:

# Tăng độ phân giải grid
model = model.refine(10)

results = model.fit(
    dataset,
    opt="LBFGS",
    steps=20
)

model.plot()

Refinement cho phép tăng khả năng biểu diễn của spline mà vẫn tận dụng cấu trúc đã học từ mô hình grid thô. Đây là một khác biệt thú vị so với cách tăng width hoặc depth của mạng MLP.

Vì sao KAN được quan tâm về khả năng diễn giải?

Trong một MLP lớn, một trọng số riêng lẻ thường khó mang ý nghĩa trực tiếp. Kết quả cuối cùng xuất hiện từ tương tác của nhiều ma trận trọng số và các activation liên tiếp.

Với KAN, ta có thể trực tiếp quan sát hình dạng φ(x) trên từng cạnh. Nếu một cạnh sau huấn luyện có hình dạng gần với một hàm quen thuộc, chẳng hạn:

$$\phi \left(x\right)\approx x^2$$

$$\phi \left(x\right)\approx \sin \left(x\right)$$

thì ta có cơ sở để thử biểu diễn cạnh đó bằng một hàm toán học đơn giản hơn. Đây là một trong những lý do KAN được nghiên cứu trong symbolic regression và Scientific Machine Learning.

Train
Học φ(x)

→

Sparsify
Giảm cạnh yếu

→

Prune
Loại phần dư thừa

→

Symbolic
Khớp hàm toán học

Dù vậy, khả năng diễn giải không tự động xuất hiện trong mọi KAN. Nếu mô hình có quá nhiều layer, quá nhiều cạnh và grid quá chi tiết, hệ thống vẫn có thể trở nên khó đọc. Pruning và regularization thường quan trọng khi interpretability là mục tiêu chính.

KAN có tốt hơn MLP không?

Không có câu trả lời “có” cho mọi trường hợp. KAN gốc cho thấy những kết quả đáng chú ý trong xấp xỉ hàm, giải phương trình vi phân và các bài toán khoa học. Tuy nhiên, điều đó không có nghĩa có thể thay MLP bằng KAN trong mọi Transformer, CNN hay hệ thống Deep Learning và mặc định nhận kết quả tốt hơn.

Một vấn đề quan trọng là số tham số và chi phí tính toán không giống nhau. Với KAN sử dụng spline, số tham số còn phụ thuộc vào số khoảng grid G và bậc spline k.

Với một KAN có width xấp xỉ N và L layer, số tham số spline có thể được mô tả ở cấp độ độ phức tạp như:

$$\mathrm{Parameters}\in \mathrm{O}\left(N^{2}L\right(G+k\left)\right)$$

Trong khi đó một MLP cùng width và depth có số weight ở cấp độ:

$$\mathrm{Parameters}\in \mathrm{O}\left(N^{2}L\right)$$

KAN có thể bù lại điều này bằng việc sử dụng mạng hẹp hơn trong một số bài toán, nhưng đây là điều cần kiểm chứng thực nghiệm chứ không phải bảo đảm chung.

Tiêu chí	MLP	KAN
Cạnh	Scalar weight	Hàm một biến
Phép tính chính	Matrix multiplication	Đánh giá nhiều hàm cạnh
Tối ưu GPU	Rất trưởng thành	Phụ thuộc implementation
Diễn giải	Thường cần công cụ bổ sung	Có thể trực quan hóa từng hàm cạnh
Điểm mạnh	Throughput, hệ sinh thái, mô hình quy mô lớn	Function fitting, interpretability, scientific modeling

Ít tham số không đồng nghĩa chạy nhanh hơn

MLP chủ yếu dựa vào phép nhân ma trận, một trong những phép toán được GPU hiện đại tối ưu tốt nhất. Với KAN spline, mỗi cạnh cần đánh giá nhiều basis function và kết hợp chúng.

Vì vậy, ngay cả khi một KAN đạt kết quả tốt với width nhỏ, inference latency hoặc training time vẫn phải được đo trực tiếp. Parameter efficiency và computational efficiency là hai khái niệm khác nhau.

KAN phù hợp với những bài toán nào?

KAN đặc biệt đáng thử khi đầu vào có số chiều vừa phải và bài toán chứa những quan hệ hàm cần được mô hình hóa hoặc diễn giải trực tiếp.

• Function approximation: học một hàm phi tuyến từ dữ liệu quan sát.
• Scientific Machine Learning: mô hình hóa quan hệ giữa các đại lượng khoa học.
• Symbolic regression: hỗ trợ tìm công thức toán học có thể giải thích.
• PDE: xấp xỉ các hàm xuất hiện trong phương trình vi phân.
• Tabular data: có thể thử nghiệm khi quan tâm đến quan hệ phi tuyến giữa từng feature và đầu ra.

Trong Computer Vision hoặc NLP, các representation thường có số chiều lớn và yêu cầu throughput cao. Việc thay toàn bộ MLP bằng KAN vì thế không phải luôn là lựa chọn hợp lý.

Một hướng thực nghiệm có kiểm soát hơn là dùng KAN tại các vị trí cụ thể như classification head, regression head, projection layer hoặc một module nhỏ trong kiến trúc lớn rồi benchmark với MLP tương đương.

Khi so sánh KAN với MLP: nên báo cáo đồng thời validation metric, số tham số, thời gian huấn luyện, inference latency, bộ nhớ và độ ổn định qua nhiều random seed. Chỉ so sánh training loss là chưa đủ.

Những hiểu nhầm thường gặp về KAN

KAN không chỉ là MLP với activation học được

Trong MLP tiêu chuẩn, activation nằm ở node sau phép biến đổi tuyến tính. Trong KAN, phép biến đổi phi tuyến được đặt trực tiếp trên từng cạnh trước khi các tín hiệu được cộng tại node.

KAN vẫn sử dụng Backpropagation

Các hệ số spline vẫn là tham số khả vi. Loss được lan truyền ngược qua các hàm cạnh để tính gradient, sau đó optimizer cập nhật các hệ số tương ứng.

KAN không mặc định tính toán rẻ hơn

Một cạnh KAN có năng lực biểu diễn mạnh hơn một scalar weight nhưng đồng thời phức tạp hơn. Việc một mạng KAN nhỏ hơn về width không bảo đảm nó chạy nhanh hơn một MLP được tối ưu tốt trên GPU.

Interpretability không tự động xuất hiện

KAN nhỏ với các hàm cạnh đơn giản có thể dễ quan sát. Nhưng nếu width, depth và grid đều lớn, ta vẫn có thể thu được hàng nghìn hàm spline khó phân tích. Pruning và regularization trở nên quan trọng nếu khả năng giải thích là mục tiêu chính.

KAN không làm MLP trở nên lỗi thời

MLP có lợi thế rất lớn về tính đơn giản, khả năng vector hóa, hỗ trợ phần cứng và hệ sinh thái. KAN cung cấp một cách tham số hóa khác với những điểm mạnh riêng, đặc biệt cho function fitting và Scientific Machine Learning.

Kết luận

Kolmogorov-Arnold Network thay đổi một thành phần nền tảng của mạng nơ-ron: thay vì chỉ học mức độ mạnh yếu của mỗi kết nối bằng một trọng số vô hướng, KAN có thể học hình dạng của phép biến đổi trên chính kết nối đó. Mỗi node sau đó cộng kết quả của các hàm cạnh để tạo biểu diễn mới.

KAN vì thế là một hướng thiết kế đáng chú ý cho xấp xỉ hàm, symbolic regression và Scientific Machine Learning. Tuy nhiên, lựa chọn giữa KAN và MLP vẫn phải dựa trên thực nghiệm: độ chính xác, khả năng tổng quát hóa, tốc độ, bộ nhớ và mức độ diễn giải cần thiết cho chính bài toán đang giải quyết.