Bayesian Edge Pruning in KANs via Regularized Horseshoe Priors

이 연구는 Kolmogorov-Arnold Networks(KANs)의 기존 pruning 방식이 사용자가 정한 임계값에 따라 구조 선택이 크게 좌우된다는 문제를 해결하는 것을 목표로 한다. KAN은 각 edge마다 학습 가능한 일변수 함수를 두어 복잡한 관계를 표현하는 신경망이다. edge별 함수 모습을 직접적으로 분석할 수 있어 해석가능하며 blackbox로 표현하는 MLP의 문제를 해결한다. 기존 제시된 모델의 크기를 줄이는 방법인 deterministic pruning은 node의 중요도가 고정된 임계값보다 작으면 node를 제거하는 방식이다. 따라서 임계값을 조금만 바꾸어도 제거되는 연결 수와 예측 성능이 크게 달라질 수 있다. 데이터나 모델 구조가 달라질 때마다 적절한 값을 다시 찾아야 한다는 점도 실제 활용과 재현성을 떨어뜨린다.
이를 개선하기 위해 본 연구에서는 Regularized Horseshoe prior와 Variational Inference를 이용한 Probabilistic pruning 방법을 node가 아닌 edge에 적용한다. 한 edge에 속한 여러 계수를 단위로 불필요한 연결은 강하게 축소하고 중요한 edge는 유지하도록 학습한다. 또한 각 edge의 중요도를 하나의 값으로 단정하지 않고 여러 추정 결과를 통해 불확실성까지 함께 분석한다. 이후 posterior에서 계산한 edge contribution score에 따라 pruning 여부를 결정하여 점추정 판단을 posterior의 확률 판단으로 바꾼다.
평가는 두 방법에서 각각 pruning된 edge의 개수를 동일하게 두어 파라미터 개수를 통일했다. 그 결과 probabilistic pruning은 deterministinc pruning보다 낮은 MSE를 보였으며 추가로 출력에 영향을 주지 않는 edge를 우선적으로 제거하는 현상까지 확인했다. 본 연구는 KAN의 해석 가능성을 구조 선택의 불확실성 정량화까지 확장하는 재현 가능한 pruning 메커니즘을 제시한다.