MatrixFormer가 결측 행렬을 한 번의 순전파로 채운다

무엇인가

결측 행렬 완성 문제는 표 형식 데이터 대치, 패널 데이터의 인과 추론, 추천 시스템 평점 예측, 언어모델 벤치마크 점수 복원까지 폭넓게 깔려 있다. 그런데 기존 표 형식 파운데이션 모델인 TabImpute는 이를 항목별 예측으로 취급한다. R×C 행렬의 각 셀을 행·열 컨텍스트를 담은 하나의 지도학습 예제로 펼쳐 RC개 예제 표를 만들기 때문에, 같은 관측값이 모든 타깃마다 반복되고 행렬의 2차원 구조는 사라진다. 논문은 이 확장이 큰 행렬에서 메모리와 실행 시간 비용을 키운다고 지적하고, 확장된 특징 표를 공유된 2차원 표현으로 대체한다.

어떻게 동작하나

MatrixFormer-Impute의 핵심 설계는 셀 토큰 격자에 트랜스포머 셀프 어텐션을 직접 적용하는 것이다. 관측 셀은 값 임베딩 MLP와 관측 상태 임베딩을 받고, 결측 셀은 0 대신 학습된 mask 벡터와 결측 상태 임베딩을 받는다. 여기에 행 CLS 토큰 K_r개 열과 열 CLS 토큰 K_c개 행을 덧붙여 격자를 확장한다. 각 블록은 두 개의 축 방향 어텐션을 순차로 수행한다. feature attention은 각 행 안에서 열 방향으로 적용되어 같은 데이터포인트의 다른 변수 정보를 쓰게 하고, datapoint attention은 각 열 안에서 행 방향으로 적용되어 같은 변수가 다른 관측에서 보이는 경험적 구조를 쓰게 한다. 두 어텐션과 SwiGLU 게이팅 FFN은 RMSNorm과 학습된 스칼라 게인을 곱한 pre-norm 잔차로 감싸고, RoPE는 어텐션 축에 맞춰 열 위치 또는 행 위치로 적용한다. 구성은 L=8 블록, 폭 d=768, 어텐션 헤드 16개, FFN 폭 1536, K_r=K_c=12로 총 78M 파라미터다.

무엇과 다른가

디코더는 각 셀의 문맥 표현을 공유 MLP에 통과시켜 5000개 고정 빈 위의 bar distribution 파라미터를 출력한다. 양 끝 빈은 half-normal 꼬리로 대체해 지원 범위가 실수 전체가 되며, 이 분포는 자체 파라미터가 없어 학습 파라미터는 MLP 출력을 통해서만 들어간다. 점 예측값은 이 셀별 예측 분포의 평균이고, 관측 항목은 그대로 둔다. 디코더가 셀마다 독립적으로 작동하고 파라미터를 행·열에 걸쳐 공유하기 때문에 모든 결측 항목을 단일 순전파 안에서 병렬로 예측할 수 있다. 어텐션 비용도 평탄화된 격자에 dense 어텐션을 걸 때의 O(R²C²) 대신 O(RC(R+C))로 줄어든다.

어떻게 쓰나

사전학습은 합성 완전 행렬만으로 이뤄지며 세 단계다. 작은 잠재요인 행렬과 MCAR 마스킹을 쓰는 베이스 단계, 데이터와 목적함수가 다양한 메인 단계, 최대 2000행×1000열까지 확장하는 짧은 크기 확장 단계가 이어진다. 메인 단계는 저랭크 신호, 완만한 꼬리와 급한 꼬리를 가진 두 스펙트럼 패밀리, 잠재요인 행렬의 네 패밀리에 고정 40% MCAR 마스킹 또는 MCAR·MAR·MNAR을 섞은 커널(명목 결측률 5~75%)을 짝지어 8개의 동등 가중 구성요소를 만든다. 에피소드의 75%에서는 일부 행이나 열을 추가로 솎아내 관측 지원이 적은 경우를 모사한다. 목적함수는 결측 항목에 대한 음의 로그우도, 같은 마스킹 입력을 받는 두 동결 교사의 평균으로 증류한 Huber 점 손실(가중치 2), 빈 중간점을 임계값으로 하는 ranked probability score 형태의 CDF 정렬 손실(가중치 0.05)의 합이다. AdamW로 128k 업데이트, 배치당 행렬 2개, FP32 파라미터와 BF16 연산을 쓴다.

전제와 한계

배포 형태인 MatrixFormer는 범용 전문가 MatrixFormer-Impute와 추천 특화 전문가 MatrixFormer-RecSys의 볼록 결합이다. 입력 행렬마다 관측 셀에서 두 전문가의 원시 예측과 실제값의 정규화 제곱 오차를 최소화하는 가중치 하나를 해석적 최소제곱으로 구하고 [0,1]로 자른다. 이 가중치가 그 행렬의 모든 결측 셀 예측을 결합하며, 두 네트워크는 사전학습 파라미터를 그대로 유지한다. 추론에는 두 전문가의 순전파가 모두 필요하다.

실험 수치는 원문에 구체적으로 제시된다. H200에서 1024×10 합성 행렬을 완성하는 중앙값 시간은 0.0389초로 TabImpute의 4.167초 대비 107배 빠르며, 16행에서는 2.1배에 그친다. 인과 패널 벤치마크에서는 6개 패널 중 5개에서 최저 정규화 점수를 기록했고, California smoking 처리블록 RMSE는 0.159로 SDiD의 0.212보다, CPS는 0.177로 0.215보다 낮았다. Kansas 효과 오차만 HyperImpute의 0.119가 MatrixFormer의 0.125보다 좋았다. 표 대치에서는 MissBench 462개 케이스에서 NRMSE 1.580으로 TabImpute legacy의 1.585를 근소하게 앞섰고, UCI 1800개 케이스에서는 0.998로 ICE의 1.172를 앞섰다. 다만 MissBench의 MCAR에서는 TabImpute가 0.819 대 0.866으로, MAR에서는 K-NN이 1.075 대 1.086으로 더 좋았다. 추천에서는 MovieLens 100K RMSE/MAE 0.8874/0.7279, Netflix small 0.8617/0.7330으로 모든 베이스라인보다 낮았지만 Netflix의 SVD 0.8621 대비 차이는 작아 유의성 주장이 아니라고 밝힌다. LLM 벤치마크 점수 복원에서는 MTEB, Merged, BenchPress에서 k=5,10,15 모두 가우시안 추정기보다 높은 R²를 냈고(Merged 0.6464 대 0.3625, BenchPress 0.4687 대 -0.0253, k=15), MMLU에서는 가우시안이 0.9266 대 0.9155로 앞섰다. 180개 설정 중 119개에서 개선, 61개에서 악화됐다.

개발자 관점에서 이 논문이 주는 실무적 의미는 결측이 있는 사용자-아이템 행렬, 패널 데이터, 모델×벤치마크 실험 결과 표 같은 것을 파인튜닝 없이 동일 가중치로 채울 수 있다는 주장이다. 특히 항목마다 컨텍스트를 반복하는 방식 대비 행 수가 커질수록 런타임 이점이 커진다는 점이 실무적으로 중요하다. 도입 전에 확인할 것은 세 가지다. 첫째, UCI 실험에서 쓴 추론 윈도우가 최대 2000행×1000열로 사전학습에서 본 가장 큰 행렬 크기이며 그보다 큰 행렬은 분할 추론과 겹치는 꼬리 윈도우 평균이 필요하다. 둘째, 추천 데이터에서는 예측값을 관측된 학습 평점 범위로 클리핑한다. 셋째, 앙상블 가중치가 관측 셀에서 적합되므로 입력에 포함된 관측값을 증거로 쓰는 in-sample 절차라는 점이다.

저자들이 직접 밝힌 한계와 전제는 다음과 같다. 모델 선택은 합성 데이터로만 했지만 개발 과정에서 다운스트림 벤치마크를 들여다봤고, UCI 컨텍스트 선택은 블라인드 검증이 아닌 탐색적 후속 실험이라고 명시한다. 윈도우 한도 역시 전체 행렬 결과를 먼저 본 뒤 채택한 탐색적 프로토콜 선택이다. 인과 패널 결과와 fresh·natural-missingness 대치 실험은 크기 확장 이전 체크포인트로 계산됐다. 앙상블 가중치는 관측 셀에서 적합되어 in-sample 증거를 사용하며, 별도로 숨긴 셀에서 블렌드를 평가했다. 향후 과제로는 예측 분포를 치료 효과 구간처럼 구간 추정에 활용하는 것, 더 큰 행렬로 컨텍스트를 확장하는 것, 홀드아웃 항목에서 앙상블 가중치를 선택하는 것을 제시한다.