텐서 연산을 반복문처럼 읽는 표기법, einx

ICLR 2026 구두 발표 ‘It’s All Just Vectorization’ 짧게 읽기

It’s All Just Vectorization: einx, a Universal Notation for Tensor OperationsICLR 2026 Oral

짧은 읽기
효율, 시스템
텐서 프로그래밍
gather, matmul, reshape처럼 제각각인 텐서 함수들을 기본 연산 몇 개와 벡터화 문자열로 다시 쓰는 einx 표기법을 짧게 읽는다. 괄호 하나로 연산에 넘기는 축과 되풀이하는 축을 가른다.
공개

2026년 10월 6일

한 줄 요약

텐서 함수가 저마다 다른 규칙을 쓰는 까닭은 벡터화를 적는 방식이 제각각이기 때문이고, 벡터화를 반복문에 빗댄 문자열 하나로 적으면 수십 개의 함수가 기본 연산 몇 개로 줄어든다.

무엇이 달라졌나

넘파이(NumPy)나 파이토치에서 색인 하나만 해도 take, gather, index_select, take_along_dim이 따로 있고, 인자 규칙도 다 다르다. 저자들이 보기에 이 함수들은 모두 같은 일을 한다. 좌표 하나로 텐서에서 값 하나를 꺼내는 기본 연산이 있고, 함수마다 다른 점은 그 연산을 어느 축을 따라 되풀이하느냐, 곧 벡터화(vectorization)의 모양뿐이다.

이 관점으로 정리하면 표가 깔끔해진다. 논문 표 1은 넘파이식 호출 24가지를 einx의 기본 연산 네 개로 옮긴다. 색인 8가지는 get_at, 외적과 크로네커 곱 같은 곱 4가지는 multiply, matmul과 tensordot 같은 내적 4가지는 dot, 전치와 reshape, concatenate 같은 모양 바꾸기 8가지는 id(항등 함수)가 된다. 호출끼리 다른 부분은 괄호 안의 벡터화 문자열뿐이다.

기존 도구와 비교하면 차이가 선명하다. einsum은 곱과 합, 축 순서 바꾸기만 다루고, einops는 여기에 모양 바꾸기와 max, mean 같은 축약을 더한 정도다(표 2). einx는 소프트맥스든 사용자가 짠 함수든 어떤 연산에나 같은 문자열 규칙을 쓴다(4.2절).

반복문으로 쓴 행렬 곱z[a, c] = dot(x[a, :], y[:, c])einx로 쓴 같은 연산einx.dot("a [b], [b] c -> a c")a, c 반복 변수 → 벡터화 축[b] 콜론 → 연산에 넘기는 축

(가) 반복문에서 einx 식으로

넘파이식 호출 (표 1)einx 연산색인 8가지take, gather, index_select 등get_at곱 4가지outer, kron, khatri_rao 등multiply내적 4가지matmul, dot, tensordot, innerdot모양 8가지transpose, reshape, stack 등id

(나) 호출 24가지, 연산 네 개

(가) 행렬 곱을 반복문으로 쓰면 반복 변수 a, c와 콜론(:)이 보인다. einx는 반복 변수를 괄호 없는 축으로, 콜론을 대괄호로 감싼 새 축 [b]로 적는다. 대괄호 축은 기본 연산 dot에 통째로 넘기고, 나머지 축에서는 그 연산을 되풀이한다(3.2절, 4.1절). (나) 논문 표 1에 실린 넘파이식 호출 24가지는 einx의 기본 연산 네 개로 모이고, 서로 다른 점은 벡터화 문자열뿐이다.

어떻게 보였나

출발점은 반복문이다. 행렬 곱을 z[a, c] = dot(x[a, :], y[:, c])처럼 쓰면, 반복 변수 a, c는 연산을 되풀이하는 축이고 콜론은 기본 연산 dot에 통째로 넘기는 축이다. einx는 이 식을 거의 그대로 옮긴다. 반복 변수는 괄호 없이 쓰고, 콜론 자리에는 새 축 이름을 대괄호로 감싸 넣는다. 그러면 einx.dot("a [b], [b] c -> a c", x, y)가 된다(4.1절). 대괄호가 붙지 않은 축은 모두 벡터화한다는 규칙 하나가 모든 연산에 똑같이 적용된다.

이 위에 einops에서 가져온 축 합치기 (a b), 이어 붙이기를 뜻하는 (a + b), 축 이름 뒤에 붙여 여러 축으로 펼치는 말줄임표 s...가 얹힌다. 대괄호가 특히 쓸모 있는 곳은 오타를 잡을 때다. einsum에서 "ij,jk->ik"를 "ij,ik->ik"로 잘못 치면 오류 없이 행렬 곱 대신 j를 따라 합하는 전혀 다른 연산이 돌아간다. einx.dot에서는 같은 실수가 대괄호 규칙이나 내적 모양 검사에 걸려 바로 오류가 난다(4.3절).

구현은 einx 식을 해당 프레임워크의 평범한 함수 호출로 번역해 파이썬 코드로 만들고, 같은 모양의 호출에는 캐시를 다시 쓴다. 넘파이에서 처음 번역하는 데 6.8ms에서 23.5ms가 들고, 그 뒤 캐시를 찾는 비용은 호출마다 0.1ms 미만이다(부록 표 3). jax.jit으로 묶으면 이 비용도 사라진다고 저자들은 적었다(4.4절).

눈여겨볼 점

첫째, 정량 평가가 없다. einx가 모양 오류를 줄이거나 코드를 빨리 읽게 한다는 주장은 예제로만 뒷받침된다. 공개 리뷰 네 건의 점수는 4, 8, 6, 6이었고, 가장 낮은 점수를 준 리뷰어는 사용자 연구도 채택 데이터도 없다는 점을 짚었다. 다른 리뷰어는 기본 예제는 반복문으로 쉽게 풀리지만 다중 헤드 어텐션 정도만 돼도 표기에 익숙해야 읽힌다고 지적했다.

둘째, 한 리뷰어가 짚었듯 “모든 것이 벡터화”라는 관찰 자체는 새롭지 않다. jax.vmap이 이미 연산을 한 축씩 들어 올리는 변환이고, 이름 붙은 텐서(named tensor) 연구도 기본 연산을 비슷하게 나눴다(2.3절). 블로그의 해석으로는, 이 논문의 몫은 그 관점을 einsum 계열의 문자열 문법 하나로 묶어 기존 프레임워크 위에서 바로 쓰게 한 데 있다.

셋째, 채택의 벽이 높다. 부록 표 4에 따르면 주요 학회 논문에 딸린 공개 저장소 11,898개 가운데 einsum을 쓰는 곳이 35.27%, einops가 21.19%이고, 이름 붙은 텐서 도구들은 1%에도 못 미친다. 새 표기법은 좋기만 해서는 퍼지지 않는다는 뜻이기도 하다.

열린 질문

  1. 벡터화 문자열은 사람이 읽는 주석이면서 기계가 풀 수 있는 모양 명세이기도 하다. 이 명세를 실제 코드의 모양 오류를 미리 잡거나, 넘파이식 코드를 einx 식으로 자동으로 옮기는 데 쓸 수 있는지, 그리고 그렇게 바꾼 코드가 정말 오류를 줄이는지 잴 수 있을까.

이해 확인

어텐션 행렬 A에 einx.softmax("b q [k] h", A)를 적용했다. 옳은 설명은?

  1. 소프트맥스는 k 축 벡터 하나를 받아 계산하고, b, q, h의 모든 조합에서 이를 되풀이한다
  2. k 축을 따라 값을 더해 없애고, 결과는 b q h 모양이 된다
  3. b, q, h 세 축을 한데 묶어 한 번에 정규화한다
  4. 대괄호는 k 축을 계산에서 빼라는 표시이므로 k를 뺀 나머지 축을 따라 정규화한다

einx에서 대괄호로 감싼 축은 기본 연산에 통째로 넘기는 축이고, 괄호 없는 축은 연산을 되풀이하는 벡터화 축이다(4.1절). 그래서 소프트맥스는 길이 k인 벡터를 받아 같은 길이의 벡터를 돌려주고, 출력 모양은 입력과 같은 b q k h다. 합으로 축을 없애는 것은 einx.sum이 하는 일이다. 논문의 다중 헤드 어텐션 예제도 이 식으로 키 축을 따라 소프트맥스를 건다(5.2절).

논문 정보