KAN(Kolmogorov-Arnold Networks) 논문 분석: MLP를 대체하는 학습 가능한 활성화 함수 아키텍처

KAN Layer: [2, 3] (2 inputs -> 3 outputs)

Input Layer          Learnable Activations (Edges)          Output Layer
                         (B-spline functions)

              +------ phi_1,1(x) --------+
              |                           |
  x_1 -------+------ phi_2,1(x) --------+-----> y_1 = phi_1,1(x_1) + phi_1,2(x_2)
              |                           |
              +------ phi_3,1(x) --------+
              |                           |
              |                           +-----> y_2 = phi_2,1(x_1) + phi_2,2(x_2)
              |                           |
  x_2 -------+------ phi_1,2(x) --------+
              |                           |
              +------ phi_2,2(x) --------+-----> y_3 = phi_3,1(x_1) + phi_3,2(x_2)
              |                           |
              +------ phi_3,2(x) --------+

  [Nodes: SUM only]   [Edges: Learnable]    [Nodes: SUM only]

  Total learnable functions: n_in x n_out = 2 x 3 = 6
  Each phi_{j,i} is a B-spline with (G + k) parameters
  where G = grid intervals, k = spline order

import torch
import torch.nn as nn
import numpy as np


def compute_bspline_basis(x: torch.Tensor, grid: torch.Tensor, k: int) -> torch.Tensor:
    """B-스플라인 기저 함수를 De Boor 재귀 공식으로 계산한다.

    Args:
        x: 입력 텐서 [batch_size, in_features]
        grid: 매듭(knot) 벡터 [in_features, grid_size + 2k + 1]
        k: 스플라인 차수 (보통 3, 즉 3차 스플라인)

    Returns:
        기저 함수 값 [batch_size, in_features, grid_size + k]
    """
    # x를 grid와 비교할 수 있도록 차원 확장
    # x: [batch, in_features] -> [batch, in_features, 1]
    x = x.unsqueeze(-1)

    # 0차 기저 함수: 구간 내에 있으면 1, 아니면 0
    bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).float()

    # De Boor 재귀: k번 반복하여 k차 기저 함수 계산
    for degree in range(1, k + 1):
        # 왼쪽 항: (x - t_i) / (t_{i+degree} - t_i)
        left_num = x - grid[:, :-(degree + 1)]
        left_den = grid[:, degree:-1] - grid[:, :-(degree + 1)]
        left = left_num / left_den.clamp(min=1e-8)

        # 오른쪽 항: (t_{i+degree+1} - x) / (t_{i+degree+1} - t_{i+1})
        right_num = grid[:, (degree + 1):] - x
        right_den = grid[:, (degree + 1):] - grid[:, 1:(-degree)]
        right = right_num / right_den.clamp(min=1e-8)

        # 재귀 결합
        bases = left * bases[:, :, :-1] + right * bases[:, :, 1:]

    return bases  # [batch, in_features, grid_size + k]


# 사용 예시
batch_size, in_features = 32, 4
grid_size, k = 5, 3  # 5개 그리드 구간, 3차 스플라인

# 균일 그리드 생성: [-1, 1] 범위에 k개씩 양쪽 확장
grid_points = torch.linspace(-1, 1, grid_size + 1)
extended_grid = torch.cat([
    grid_points[0:1] - torch.arange(k, 0, -1) * (grid_points[1] - grid_points[0]),
    grid_points,
    grid_points[-1:] + torch.arange(1, k + 1) * (grid_points[1] - grid_points[0])
])
# in_features 차원으로 복제
grid = extended_grid.unsqueeze(0).repeat(in_features, 1)

x = torch.randn(batch_size, in_features)
basis_values = compute_bspline_basis(x, grid, k)
print(f"기저 함수 shape: {basis_values.shape}")
# 출력: 기저 함수 shape: torch.Size([32, 4, 8])  # G + k = 5 + 3 = 8

import torch
import torch.nn as nn
import torch.nn.functional as F


class KANLayer(nn.Module):
    """Kolmogorov-Arnold Network 레이어.

    각 엣지에 학습 가능한 B-스플라인 활성화 함수를 배치한다.
    노드에서는 단순 합산만 수행한다.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        grid_size: int = 5,
        spline_order: int = 3,
        scale_base: float = 1.0,
        scale_spline: float = 1.0,
        grid_range: tuple = (-1.0, 1.0),
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.grid_size = grid_size
        self.spline_order = spline_order
        self.scale_base = scale_base
        self.scale_spline = scale_spline

        # 균일 그리드 생성 및 확장 (양쪽으로 spline_order만큼)
        h = (grid_range[1] - grid_range[0]) / grid_size
        grid = torch.arange(-spline_order, grid_size + spline_order + 1) * h + grid_range[0]
        # grid shape: [grid_size + 2 * spline_order + 1]
        # 각 (in, out) 쌍마다 별도 그리드를 가질 수 있도록 확장
        self.register_buffer(
            "grid",
            grid.unsqueeze(0).unsqueeze(0).expand(out_features, in_features, -1)
        )

        # 스플라인 계수: 각 엣지마다 (grid_size + spline_order)개
        self.spline_weight = nn.Parameter(
            torch.randn(out_features, in_features, grid_size + spline_order)
            * 0.1
        )

        # 잔차(base) 가중치
        self.base_weight = nn.Parameter(
            torch.randn(out_features, in_features) * 0.1
        )

    def b_splines(self, x: torch.Tensor) -> torch.Tensor:
        """B-스플라인 기저 함수 계산.

        Args:
            x: [batch_size, in_features]

        Returns:
            기저 함수 값 [batch_size, in_features, grid_size + spline_order]
        """
        # x를 [batch, 1, in_features, 1]로 확장하여 grid와 비교
        x = x.unsqueeze(1).unsqueeze(-1)  # [B, 1, in, 1]

        # grid: [out, in, G+2k+1] -> 방송을 위해
        grid = self.grid  # [out, in, G+2k+1]

        # 0차 기저: 구간 내에 있으면 1
        # grid[:, :, :-1]과 grid[:, :, 1:]를 사용
        # 첫 번째 out 차원의 그리드만 사용 (모두 동일하므로)
        g = grid[0]  # [in, G+2k+1]
        bases = ((x[:, 0] >= g[:, :-1]) & (x[:, 0] < g[:, 1:])).float()
        # bases: [B, in, G+2k]

        for degree in range(1, self.spline_order + 1):
            left_num = x[:, 0] - g[:, :-(degree + 1)]
            left_den = g[:, degree:-1] - g[:, :-(degree + 1)]
            left = left_num / left_den.clamp(min=1e-8)

            right_num = g[:, (degree + 1):] - x[:, 0]
            right_den = g[:, (degree + 1):] - g[:, 1:(-degree)]
            right = right_num / right_den.clamp(min=1e-8)

            bases = left * bases[:, :, :-1] + right * bases[:, :, 1:]

        return bases  # [B, in, G+k]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """순전파: 잔차 함수 + 스플라인 함수.

        Args:
            x: [batch_size, in_features]

        Returns:
            출력 [batch_size, out_features]
        """
        batch_size = x.size(0)

        # 1) 잔차 부분: SiLU 활성화 후 선형 결합
        base_output = F.silu(x)  # [B, in]
        base_output = torch.einsum("bi,oi->bo", base_output, self.base_weight)
        # base_output: [B, out]

        # 2) 스플라인 부분: B-스플라인 기저 * 계수
        spline_basis = self.b_splines(x)  # [B, in, G+k]
        spline_output = torch.einsum(
            "big,oig->bo", spline_basis, self.spline_weight
        )
        # spline_output: [B, out]

        # 3) 결합
        return self.scale_base * base_output + self.scale_spline * spline_output

class KAN(nn.Module):
    """다층 Kolmogorov-Arnold Network.

    Args:
        width: 각 레이어의 뉴런 수 리스트. 예: [2, 5, 1]
        grid_size: B-스플라인 그리드 구간 수
        spline_order: B-스플라인 차수 (기본값: 3, 3차 스플라인)
    """

    def __init__(
        self,
        width: list,
        grid_size: int = 5,
        spline_order: int = 3,
    ):
        super().__init__()
        self.width = width
        self.layers = nn.ModuleList()

        for i in range(len(width) - 1):
            self.layers.append(
                KANLayer(
                    in_features=width[i],
                    out_features=width[i + 1],
                    grid_size=grid_size,
                    spline_order=spline_order,
                )
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x)
        return x

    def parameter_count(self) -> dict:
        """레이어별 파라미터 수를 반환한다."""
        info = {}
        total = 0
        for i, layer in enumerate(self.layers):
            layer_params = sum(p.numel() for p in layer.parameters())
            info[f"layer_{i}"] = layer_params
            total += layer_params
        info["total"] = total
        return info


# 사용 예시
model = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
print(f"모델 구조: {model.width}")
print(f"파라미터 수: {model.parameter_count()}")

# 순전파 테스트
x = torch.randn(16, 2)
y = model(x)
print(f"입력: {x.shape}, 출력: {y.shape}")
# 출력 예시:
# 모델 구조: [2, 5, 1]
# 파라미터 수: {'layer_0': 90, 'layer_1': 45, 'total': 135}
# 입력: torch.Size([16, 2]), 출력: torch.Size([16, 1])

class MLP(nn.Module):
    """비교용 Multi-Layer Perceptron.

    KAN과 동일한 width 구조를 사용하되,
    고정 활성화 함수(SiLU)와 선형 가중치를 사용한다.
    """

    def __init__(self, width: list):
        super().__init__()
        layers = []
        for i in range(len(width) - 1):
            layers.append(nn.Linear(width[i], width[i + 1]))
            if i < len(width) - 2:  # 마지막 레이어 제외
                layers.append(nn.SiLU())
        self.network = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


# MLP vs KAN 파라미터 비교
mlp = MLP(width=[2, 5, 1])
kan = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)

mlp_params = sum(p.numel() for p in mlp.parameters())
kan_params = sum(p.numel() for p in kan.parameters())

print(f"MLP 파라미터: {mlp_params}")
print(f"KAN 파라미터: {kan_params}")
print(f"KAN/MLP 파라미터 비율: {kan_params / mlp_params:.1f}x")
# 출력 예시:
# MLP 파라미터: 21
# KAN 파라미터: 135
# KAN/MLP 파라미터 비율: 6.4x

import torch
import torch.nn as nn
import matplotlib.pyplot as plt


def create_dataset(f, n_train=1000, n_test=200, input_dim=2):
    """학습/테스트 데이터셋 생성.

    Args:
        f: 근사할 대상 함수
        n_train: 학습 샘플 수
        n_test: 테스트 샘플 수
        input_dim: 입력 차원
    """
    x_train = torch.rand(n_train, input_dim) * 2 - 1  # [-1, 1]
    y_train = f(x_train)
    x_test = torch.rand(n_test, input_dim) * 2 - 1
    y_test = f(x_test)
    return x_train, y_train, x_test, y_test


def train_model(model, x_train, y_train, x_test, y_test,
                epochs=500, lr=1e-2, optimizer_type="adam"):
    """KAN 또는 MLP 모델 학습.

    Args:
        model: KAN 또는 MLP 인스턴스
        optimizer_type: 'adam' 또는 'lbfgs'

    Returns:
        학습 이력 딕셔너리
    """
    if optimizer_type == "lbfgs":
        optimizer = torch.optim.LBFGS(
            model.parameters(), lr=lr, max_iter=20,
            history_size=10, line_search_fn="strong_wolfe"
        )
    else:
        optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    criterion = nn.MSELoss()
    history = {"train_loss": [], "test_loss": []}

    for epoch in range(epochs):
        model.train()

        if optimizer_type == "lbfgs":
            def closure():
                optimizer.zero_grad()
                pred = model(x_train)
                loss = criterion(pred, y_train)
                loss.backward()
                return loss
            loss = optimizer.step(closure)
        else:
            optimizer.zero_grad()
            pred = model(x_train)
            loss = criterion(pred, y_train)
            loss.backward()
            optimizer.step()

        # 평가
        model.eval()
        with torch.no_grad():
            train_loss = criterion(model(x_train), y_train).item()
            test_loss = criterion(model(x_test), y_test).item()

        history["train_loss"].append(train_loss)
        history["test_loss"].append(test_loss)

        if (epoch + 1) % 100 == 0:
            print(f"Epoch {epoch+1}/{epochs} | "
                  f"Train MSE: {train_loss:.6f} | "
                  f"Test MSE: {test_loss:.6f}")

    return history


# 대상 함수: f(x1, x2) = exp(sin(pi * x1) + x2^2)
target_fn = lambda x: torch.exp(
    torch.sin(torch.pi * x[:, 0:1]) + x[:, 1:2] ** 2
)

x_train, y_train, x_test, y_test = create_dataset(target_fn)

# KAN 학습
kan_model = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
kan_history = train_model(kan_model, x_train, y_train, x_test, y_test,
                          epochs=500, lr=1e-2, optimizer_type="adam")

# MLP 학습 (파라미터 수를 맞추기 위해 폭을 넓힘)
mlp_model = MLP(width=[2, 50, 50, 1])
mlp_history = train_model(mlp_model, x_train, y_train, x_test, y_test,
                          epochs=500, lr=1e-3, optimizer_type="adam")

# 손실 곡선 비교 시각화
fig, ax = plt.subplots(1, 1, figsize=(10, 6))
ax.semilogy(kan_history["test_loss"], label="KAN [2,5,1]", linewidth=2)
ax.semilogy(mlp_history["test_loss"], label="MLP [2,50,50,1]", linewidth=2)
ax.set_xlabel("Epoch")
ax.set_ylabel("Test MSE (log scale)")
ax.set_title("KAN vs MLP: Function Approximation")
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("kan_vs_mlp_loss.png", dpi=150)
plt.show()

# pip install pykan
from kan import KAN as PyKAN
import torch

# 대상 함수 정의
target_fn = lambda x: torch.exp(
    torch.sin(torch.pi * x[:, [0]]) + x[:, [1]] ** 2
)

# 데이터셋 생성 (pykan 내장 유틸리티 활용)
dataset = {
    "train_input": torch.rand(1000, 2) * 2 - 1,
    "test_input": torch.rand(200, 2) * 2 - 1,
}
dataset["train_label"] = target_fn(dataset["train_input"])
dataset["test_label"] = target_fn(dataset["test_input"])

# KAN 모델 생성 및 학습
model = PyKAN(width=[2, 5, 1], grid=5, k=3, seed=42)

# 학습 실행 (LBFGS 옵티마이저 사용)
results = model.fit(dataset, opt="LBFGS", steps=50, lamb=0.01)

# 학습된 활성화 함수 시각화
# 각 엣지의 학습된 B-스플라인 함수를 플롯으로 표시
model.plot()

# 기호 회귀: 학습된 함수에서 수학적 수식 추출 시도
# model.auto_symbolic()을 통해 sin, exp 등 기호 함수로 매핑
model.auto_symbolic()
formula = model.symbolic_formula()
print(f"추출된 수식: {formula}")
# 예상 출력: exp(sin(pi*x1) + x2^2)

# 그리드 범위가 입력 데이터 범위를 포함하는지 확인
x_min, x_max = x_train.min().item(), x_train.max().item()
print(f"데이터 범위: [{x_min:.2f}, {x_max:.2f}]")
# 그리드 범위를 데이터 범위보다 약간 넓게 설정
# grid_range=(-1.5, 1.5) 등으로 조정

# LBFGS 옵티마이저 사용 시 closure 패턴 필수
optimizer = torch.optim.LBFGS(
    model.parameters(),
    lr=0.1,           # LBFGS는 보통 더 큰 lr 사용
    max_iter=20,       # 내부 반복 횟수
    tolerance_grad=1e-7,
    tolerance_change=1e-9,
    line_search_fn="strong_wolfe"  # 안정적 라인서치
)

# grid_eps 조정으로 그리드 적응
# 초기에는 0.5 정도로 시작하고, 학습 후반에 0.0으로 줄이기

# 그래디언트 체크포인팅을 활용한 메모리 절약
from torch.utils.checkpoint import checkpoint

class MemoryEfficientKAN(nn.Module):
    def __init__(self, width, grid_size=5, spline_order=3):
        super().__init__()
        self.layers = nn.ModuleList([
            KANLayer(width[i], width[i+1], grid_size, spline_order)
            for i in range(len(width) - 1)
        ])

    def forward(self, x):
        for layer in self.layers:
            # 그래디언트 체크포인팅: 메모리 절약 (역전파 시 재계산)
            x = checkpoint(layer, x, use_reentrant=False)
        return x

import matplotlib.pyplot as plt
import torch
import numpy as np


def visualize_kan_activations(model, layer_idx=0, input_range=(-1, 1)):
    """KAN 레이어의 학습된 활성화 함수를 시각화한다.

    Args:
        model: KAN 모델 인스턴스
        layer_idx: 시각화할 레이어 인덱스
        input_range: 입력 범위
    """
    layer = model.layers[layer_idx]
    n_in = layer.in_features
    n_out = layer.out_features

    # 시각화용 입력 생성
    x_vis = torch.linspace(input_range[0], input_range[1], 200)

    fig, axes = plt.subplots(n_out, n_in, figsize=(4 * n_in, 3 * n_out))
    if n_out == 1:
        axes = axes.reshape(1, -1)
    if n_in == 1:
        axes = axes.reshape(-1, 1)

    for j in range(n_out):
        for i in range(n_in):
            ax = axes[j][i]

            # 단일 입력 차원에 대한 활성화 함수 계산
            x_input = torch.zeros(200, n_in)
            x_input[:, i] = x_vis
            with torch.no_grad():
                # 스플라인 부분만 추출
                basis = layer.b_splines(x_input)  # [200, in, G+k]
                spline_out = torch.einsum(
                    "big,g->bi", basis[:, i:i+1, :],
                    layer.spline_weight[j, i, :]
                )

            ax.plot(x_vis.numpy(), spline_out[:, 0].numpy(),
                    linewidth=2, color="blue")
            ax.set_title(f"phi_{j+1},{i+1}(x)", fontsize=10)
            ax.grid(True, alpha=0.3)
            ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)

    plt.suptitle(f"Layer {layer_idx} Activation Functions", fontsize=14)
    plt.tight_layout()
    plt.savefig(f"kan_layer{layer_idx}_activations.png", dpi=150)
    plt.show()


# 학습 후 시각화
# visualize_kan_activations(kan_model, layer_idx=0)