Functional Attention: From Pairwise Affinities to Functional Correspondences

Motivation

• Known: operator learning은 PDE 풀이, 물리 시뮬레이션 등 여러 응용에서 중요하며, FNO나 U-NO 같은 spectral 기반 neural operator와 OFormer, GNOT, Transolver 같은 transformer 기반 operator가 각각 함수 표현과 attention 기반 성능에서 좋은 결과를 보여왔다. Galerkin Transformer는 linear attention을 Petrov-Galerkin projection으로 해석하여 attention과 function space 사이의 연관성을 제시한 바 있다.
• Gap: 기존 transformer 기반 operator learning 방법들은 연속적인 함수를 discrete token으로 취급하여 global functional structure를 무시하고, scaled dot-product attention의 quadratic complexity로 인해 확장성이 떨어지며, 서로 다른 resolution이나 irregular mesh 간 일관성을 유지할 원칙적인 방법이 부재하다.
• Why: operator learning은 PDE 풀이, computational design, 물리 시뮬레이션 등 과학 계산 전반에 핵심적인 역할을 하는데, resolution에 강건하고 global structure를 반영하면서도 효율적인 attention 메커니즘을 확보하면 다양한 discretization과 도메인에 걸쳐 일반화 가능한 신경 연산자 설계가 가능해진다.
• Approach: 저자들은 attention을 query/key-value 공간에 각각 학습된 adaptive basis를 두고 이들 사이의 대응관계를 구하는 functional map 문제로 재정식화하며, softmax affinity 대신 spectral domain에서의 optimal linear solve(least-squares regression)로 attention을 계산한다.

Achievement

Figure 3. PDE solving visualization. Ground truth and error maps

1. 새로운 attention 재해석: 표준 attention 수식과 geometric functional maps 프레임워크 사이의 원칙적인 연결고리를 확립하여 attention을 functional correspondence로 재정의했다.
2. 다양한 태스크에서의 SOTA 성능: PDE solving, 3D segmentation, regression 등 폭넓은 operator learning 태스크에서 기존 SOTA와 동등하거나 우수한 성능을 달성했다.
3. Resolution-invariance 및 효율성: 서로 다른 discretization과 mesh 해상도에 걸쳐 강건성을 보이며, compact하고 generalizable한 표현을 제공한다.

How

Figure 1. Architecture Overview. Top: Input functions are encoded by MLP, processed through N FUNCATTN blocks, and decod

• 입력 함수를 MLP로 encoding한 뒤, query와 key-value 공간 각각에 대해 별도의 feed-forward 네트워크로 adaptive functional basis를 학습한다.
• Functional maps 프레임워크에서처럼, 두 basis 사이의 대응관계를 나타내는 structured linear operator(함수적 대응 행렬)를 spectral coefficient 공간에서의 least-squares 문제로 풀어 attention을 계산한다.
• 이는 Galerkin Transformer의 암묵적 basis 변화나 Transolver의 slice-and-attend 방식을 일반화한 형태로, 명시적으로 basis를 분리해 학습함으로써 더 큰 표현력을 확보한다.
• PDE solving, few-shot sinusoidal regression, 3D segmentation 등 다양한 벤치마크에서 정확도, runtime, 메모리 사용량을 기존 방법(FNO, Transolver, Galerkin Transformer 등)과 비교 평가한다.

Originality

• Attention을 token-wise pairwise affinity가 아니라 함수 공간 사이의 functional correspondence로 재해석하는 관점 자체가 독창적이며, geometry processing 분야의 functional maps 개념을 operator learning에 최초로 연결했다.
• 기존 linear attention이나 low-rank approximation 계열이 표준 softmax attention을 근사하는 데 초점을 맞춘 반면, 본 논문은 spectral domain에서의 optimal linear solve라는 전혀 다른 계산 패러다임을 제안한다.
• Galerkin Transformer의 암묵적 basis 변화, Transolver의 slice-and-attend 개념을 각각 명시적 basis 분리와 일반화된 spectral framework로 확장했다는 점에서 기존 연구 대비 이론적 진전을 보인다.

Limitation & Further Study

• Adaptive basis의 차원 수나 spectral coefficient 개수에 대한 선택이 성능에 미치는 영향과 이를 자동으로 결정하는 방법에 대한 논의가 발췌 부분에서는 충분히 드러나지 않는다.
• Functional maps 프레임워크가 원래 3D shape correspondence를 위해 설계된 만큼, 매우 복잡하거나 topology가 급격히 변하는 함수 공간에서의 이론적 보장이나 실패 사례 분석이 더 필요해 보인다.
• 후속 연구로는 다양한 도메인(예: 시계열, 그래프 기반 물리 시스템)에 대한 확장성 검증과, spectral solve의 수치적 안정성에 대한 심층 분석이 요구된다.

Evaluation

총평: Attention을 functional correspondence 관점에서 재해석하는 참신하고 이론적으로 잘 뒷받침된 프레임워크로, operator learning 분야에 새로운 설계 공간을 열어주는 의미 있는 기여를 한다.