Robert Joseph George가 TorchLean의 텐서 API를 이용해 Lean에서 계산 함수를 작성하고 CPU·GPU에서 실행하는 방법을 소개한다. 반복 갱신, 사용자 정의 손실 함수, einsum 기반 행렬 곱과 어텐션을 예제로 들며, 함수의 수학적 성질과 GPU용 변환이 원래 계산을 보존하는지 Lean에서 증명하는 접근을 설명한다.

함수 작성과 실행

TorchLean의 텐서 라이브러리는 형상이 지정된 텐서 연산, 자동 미분(autograd), FloatLib의 수치 형식을 제공한다. 자동 미분은 역전파에 필요한 연산과 중간값을 기록한다. 이제 지원되는 Lean 함수를 컴파일해 GPU에서 실행하고 텐서 API를 통해 호출할 수 있다. 함수 자체에 장치를 고정하지 않고 실행 시 CPU 또는 GPU를 선택한다.

예를 들어 원소별 제곱은 [1, 2, 3]을 [1, 4, 9]로 바꾼다. 각 출력이 대응하는 입력에만 의존하므로 계산은 서로 독립적이다. α와 [Mul α]로 원소 형식을 일반화하고 입력에서 Float32를 선택할 수 있으며, 같은 함수를 여러 텐서 형상에 적용할 수 있다. GPU 경로는 지원되는 산술을 인식해 확인된 표현을 만들고 CUDA 코드를 생성한다. NVIDIA NVRTC가 실행 시 코드를 컴파일한다. 지원되는 지역 바인딩, 비교, 조건문도 함수에 포함할 수 있다.

두 텐서의 대응 원소를 가중 평균하는 함수에서는 두 입력 형상이 일치해야 한다. .zip은 길이를 임의로 자르거나 요청하지 않은 브로드캐스팅을 하지 않는다. 캡처된 가중치 같은 스칼라는 GPU 코드의 상수가 되므로 값이 달라지면 재컴파일이 필요할 수 있다. 런타임은 같은 소스 코드를 캐시하며, 원소별로 자주 바뀌는 값은 텐서 입력으로 전달할 수 있다.

재귀 함수 evolve로 x ↦ x²+1을 반복하면 시작값 1은 2회 후 5가 되고, 2는 26이 된다. 서로 다른 시작값은 병렬로 처리할 수 있지만, 한 값의 다음 단계는 이전 값에 의존한다. 순방향 GPU 호출에서는 프런트엔드가 재귀의 0·후속 단계 방정식에서 순서가 있는 루프를 만들고 그 루프가 방정식과 일치한다는 Lean 증명을 구성한다. 반복할 때마다 함수 본문을 복제하지 않는다. 현재 인식되는 재귀는 매개변수를 고정하고 마지막 자연수 인자의 전임자로 재귀하는 형태다. UInt64 반복 횟수를 자연수로 변환해 전달하며, 트리 재귀와 동적으로 할당되는 재귀 구조는 이 경로에서 지원하지 않는다.

갱신식 relax는 현재값에서 목표까지 남은 거리의 일정 비율을 이동한다. 시작값 [0, 2, 4], 목표 2, 비율 0.5로 세 번 갱신하면 [1.75, 2, 2.25]가 된다. 실수에서는 모든 비율·목표·초기값·자연수 단계에 대해 다음 식을 증명할 수 있다: relax rate target x steps = target + (1 - rate) ^ steps * (x - target). 귀납법과 mathlib의 ring 전술을 사용한다. 비율이 0.5이면 오차가 매번 절반으로 줄고, 1.5이면 목표를 오가며 오차가 절반으로 줄어든다. 2에서는 오차가 진동하고, 2보다 크면 오차가 커진다. 반올림이 있는 계산에는 별도의 유한 정밀도 오차 경계가 필요하다.

미분과 손실 함수

기존 기록 테이프를 이용해 반복 갱신의 기울기도 구할 수 있다. 두 단계에 걸친 미분은 각 단계의 국소 미분을 역순으로 곱한다. GPU에서 기록된 재귀는 현재 단계마다 테이프 연산을 실행하므로, 기록 없이 하나의 루프로 실행하는 순방향 경로와 실행 비용이 다르다.

energy(x) = sum(x*x)에 입력 [1, 2, 3]을 넣으면 스칼라 출력은 14, 기울기는 [2, 4, 6]이다. 출력이 텐서인 함수는 result.backward seed로 각 출력의 민감도를 지정할 수 있다. 이는 전체 야코비안 행렬을 만들지 않고 출력 민감도를 입력으로 전파하는 벡터-야코비안 곱(VJP)이다. GPU에서는 LibTorch가 지원되는 텐서 값과 국소 기울기 연산을 제공하고, 자체 자동 미분 엔진은 끈다. 조건문은 선택된 경로만 기록한다. 분기 경계에서 미분이 존재하지 않을 수 있고, GPU에서 스칼라 분기 결정을 하면 Lean과 동기화된다. 역전파는 기록을 소비하며, 순방향 값만 필요한 경우 result.close로 저장된 중간값을 해제한다.

사용자 정의 손실 함수는 예측값 [1, 2, 3], 목표값 [0, 1, 1]의 제곱 잔차를 합산한다. 손실은 6, 기울기는 [2, 2, 4]이며, 학습률 0.25로 예측값 자체를 갱신하면 [0.5, 1.5, 2], 손실은 1.5가 된다. 이 예제의 수는 모두 binary32에서 정확히 표현된다. 실제 신경망에서는 예측값의 기울기를 앞선 계층까지 전파한 뒤 매개변수를 갱신한다. 목표값은 상수로 캡처되므로 이 호출은 목표값이 아니라 제공된 예측값을 미분한다. 고정된 목표 t에 대해 (x-t)²의 미분이 2(x-t)라는 명제는 기존 autograd 전술로 증명한다. 이 특정 목적 함수에서 실수 잔차는 1-2η배가 되고 손실은 (1-2η)²배가 된다. 이는 임의의 신경망 학습이 수렴한다는 주장이 아니다.

축 연산과 어텐션

einsum은 축 이름으로 곱하고 합산할 차원을 나타낸다. 행렬 곱에서는 입력의 row feature, 가중치의 feature output을 연결하고 출력에 row output을 남긴다. 입력 [2,3]과 가중치 [3,2]의 결과는 [[22,28],[49,64]]다. 패턴에서 화살표 뒤에 없는 feature 축은 합산되고, 결과 형상은 Lean이 추론한다. 같은 축 이름을 두 피연산자에서 사용해도 출력에 남기면 합산하지 않는다. 이름은 패턴에만 속하며 텐서 형식에 영구적인 의미로 저장되지 않는다. 크기가 맞지 않는 공유 축은 거부하지만 단일 크기 축 브로드캐스팅은 지원한다. 형상 검사는 축의 의미가 의도에 맞는지까지 증명하지는 않는다.

rearrange는 행렬을 전치하거나 평탄화하고, expand는 편향에 행 축을 명시적으로 추가하며, reduce는 지정한 축을 합산해 제거한다. 전치는 좌표 대응을 바꾸며 단순히 같은 데이터의 형상만 바꾸는 것과 다르다. 평탄화도 행 우선 순서를 보존한다. 같은 원소 수를 가진 배열이라도 배치 순서는 서로 바꿔 쓸 수 없다. 실수 합과 반올림 합은 계산 순서를 다르게 다뤄야 한다.

어텐션 예제는 [batch, head, token, feature] 형상의 질의와 키를 사용해 특징 축을 합산하고 배치·헤드·질의·키 축을 유지한다. 이어 별도로 제공한 가중치와 값을 수축해 출력한다. 이 가중치는 예시용이며 표시된 점수의 소프트맥스 결과가 아니다. 완전한 어텐션은 점수를 계산한 뒤 스케일링, 마스킹, 정규화를 거친다. 예제의 두 헤드는 서로 섞이지 않으며, 마지막 rearrange는 헤드와 채널 축을 결합한다. 이 패턴 예제는 CPU 텐서 계산이다. 임의의 einsum이 패턴 검사를 통과했다고 해서 GPU 구현이나 역전파 규칙까지 자동으로 생기지는 않는다.

einsum은 증명 전술이 아니라 항(term) elaborator다. Lean이 정의를 elaboration할 때 패턴을 읽고 텐서 형상과 대조해 계산과 인증서를 만든다. 실행 때마다 축 이름을 다시 해석하지 않는다. 생성된 계산은 중간 곱 텐서를 만들지 않고 입력을 읽어 출력별로 누적할 수 있다. 예를 들어 행 256개, 특징 1024개, 출력 256개의 곱에서 중간 텐서는 binary32 원소 67,108,864개, 즉 256 MiB의 원소 데이터가 된다. 최종 [256,256] 결과는 256 KiB다. 이는 할당기 측정값이 아닌 원소 수를 기준으로 한 비교다.

einsumTensorKernel_correct는 생성 결과를 참조 계산과 연결하며 부동소수점 덧셈의 결합법칙을 요구하지 않는다. 피연산자 곱과 축소 덧셈은 지정된 순서를 유지한다. 다른 순서로 재배열하려면 추가 대수 법칙이 필요하다. 이 구현이 최적화된 GPU 행렬 라이브러리라는 뜻은 아니다. 예제는 컴파일된 Lean 텐서 경로를 사용하고, GPU 모델 행렬 곱은 LibTorch 백엔드를 사용한다.

모델과 성능 조건

MLP 예제는 입력 특징 2개, 은닉 크기 8, ReLU, 출력 1개로 구성된다. Lean은 계층 연결 지점의 형상을 검사한다. 트레이너가 목적 함수, 최적화기, 장치 설정을 제공하며 지원되는 네이티브 GPU 경로에서는 gpu를 선택할 수 있다. 변환기, 푸리에 신경 연산자, 확산 모델 등도 공통 텐서·계층·미분 API를 사용한다. 모델 매개변수와 버퍼의 스칼라 형식은 타입 상태에 기록된다. 초기화 과정에서 이미 반올림된 값을 더 넓은 형식으로 바꿔도 잃은 자릿수는 복구되지 않는다.

실행 시간은 GPU 사용 여부만으로 판단할 수 없다. 사용자 정의 코드 생성·컴파일, 커널 실행, 값 이동, 중간값 저장, 산술 비용을 따로 살펴야 한다. 런타임은 같은 CUDA 소스를 캐시하므로 첫 호출에는 NVRTC 컴파일이 포함될 수 있고, 이후 호출은 컴파일된 프로그램을 재사용할 수 있다. 캡처된 상수가 바뀌면 소스가 달라질 수 있으므로 측정은 컴파일, 준비된 실행, 또는 둘 다를 구분해야 한다. 기록된 GPU 재귀는 단계마다 연산을 실행하지만 순방향 전용 경로는 하나의 루프로 만들 수 있다.

표준 모델 연산에는 PyTorch의 기존 수치 인프라와 중복되는 자체 커널을 유지하는 대신 LibTorch를 사용한다. TorchLean이 모델과 미분 순회를 담당하고, ATen이 지원되는 텐서 값과 국소 기울기를 계산한다. 역전파의 행렬 연산은 전치 복사본을 먼저 만들지 않고 전치 플래그를 전달할 수 있다. 이는 구현 방식이지 모든 모델이 PyTorch보다 빠르다는 보장은 아니다.

보관된 A100 단일 GPU의 FP32 GPT 실험에서 같은 50회 갱신을 비교했을 때 TorchLean은 예약된 토큰 기준 초당 중앙값 5,596개, PyTorch는 19,207개였다. 최종 검증 손실은 각각 2.431321과 2.431322932였고, 이 설정에서 PyTorch가 약 3.43배 빨랐다. 훈련 손실의 최대 차이는 약 0.00001193이었다. 이 측정은 이전 FP32 구현의 기준선이며 현재 LibTorch 어댑터나 사용자 정의 컴파일러의 성능 측정값이 아니다. 별도의 보관된 체크포인트·프롬프트·11토큰 이어쓰기 실험에서는 캐시를 사용한 준비된 실행 3회가 약 168~172ms, 전체 접두부를 다시 계산한 경우 9.43~9.74초였다. 해당 작업에서 이전 계산을 재사용한 차이이며 일반적인 성능 보장은 아니다. 비교 기록과 실험 설명에 수치와 조건이 제시돼 있다.

직접 인덱싱하는 rowSum 예제는 행 우선 주소 row * width + column으로 값을 읽어 누적한다. 출력 형상만으로 모든 주소가 유효함을 보장하지 않으며, 범위를 벗어난 읽기는 오류를 낸다. 행끼리는 병렬로 실행할 수 있지만 한 행 안의 덧셈은 작성한 순서를 따른다. 이 인터페이스는 특수한 읽기나 누적 순서가 필요할 때 쓸 수 있다. 소스와 하위 표현의 대응 증명이 있다고 해서 임의 함수의 미분이나 입력 영역의 수치 경계까지 자동으로 제공되는 것은 아니다.

정밀도와 증명 범위

FloatLib는 텐서 원소의 수치 표현을 설정할 수 있다. 예제는 지수부 15비트, 가수부 112비트인 binary128을 사용해 a=1+2^-100을 표현한다. 이 추가 자릿수는 binary64에 저장하면 사라지므로 유리수에서 직접 값을 만든다. 정확한 제곱은 1+2^-99+2^-200이며 binary128 최근접 짝수 반올림은 1+2^-99를 남긴다. 예제는 CPU 결과의 정확한 유리수 값이 그 값과 같은지 검사한다. GPU 사용자 정의 경로는 넓은 값을 정수 워드로 전달해 장치에서 소프트웨어 산술을 수행하므로 정밀도를 유지하지만 하드웨어 곱셈보다 명령이 더 필요하다. 이 예제는 처리량 측정이 아니다.

설정 가능한 GPU 형식은 최근접 짝수 반올림과 점진적 언더플로를 사용하며 지수 비트 최대 30, 가수 비트 최대 4096을 지원한다. 호환 GPU의 binary16과 bfloat16은 덧셈·뺄셈·곱셈에 네이티브 연산을 쓸 수 있지만 나눗셈은 소프트웨어 경로다. 설정된 IEEE binary32·binary64는 네이티브 나눗셈도 사용할 수 있다. 십진수, 포지트 등 CPU 전용 형식은 GPU를 요청해도 형식과 정밀도를 유지하면서 CPU 대체 실행을 보고한다. 지원하지 않는 본문, 사용할 수 없는 GPU, 컴파일 실패는 오류를 낸다. 설정 형식의 GPU 기록은 예제에 나온 산술·구조 연산을 지원하지만 모든 네이티브 모델 연산을 지원하지 않는다. 별도 Graph.runBuffers 실행기는 여전히 binary32 입력과 페이로드를 요구한다. 임의의 실수는 장치에서 바로 계산할 수 있는 표현이 아니므로 실행에는 네이티브 부동소수점, FloatLib 형식, 또는 지원되는 연산의 유리수 같은 실행 가능한 표현이 필요하다.

Lean은 실수 함수의 미분과 컴파일 변환의 성질을 증명할 수 있다. autograd 전술은 미분 규칙으로 증명을 만들고 Lean 커널이 확인한다. 반올림 부동소수점 함수는 매끄러운 실수 함수와 같지 않으므로 실수 계산과 실제 수치값을 연결하려면 오차 경계가 필요하다. Program.eval_eq_reference는 검사된 평가기를 원래 계산과 연결하고, Program.lower_correct는 이를 구조화된 루프 표현까지 전달한다. 네이티브 FP32·FP64 출력 소스 정리는 산술 및 입력 버퍼 계약을 명시하지만 NVRTC나 GPU를 검증하지 않는다. 설정 형식의 CUDA 산술은 런타임 테스트에서 FloatLib와 대조하며 커널 검증 범위 밖에 있다. 물리 장치 실행까지 보장하려면 외부 산술·메모리·컴파일러 계약도 필요하다.

TorchLean은 중첩 쌍대수를 이용한 고차 방향 미분과 mathlib의 iteratedFDeriv를 연결한다. 나눗셈·로그·제곱근 등은 정의역 조건이 성립할 때 국소 규칙을 사용할 수 있다. 물리 정보 신경망 예제는 u''(x)=-2를 [-1,1]에서 만족하고 양 끝에서 0인 함수를 학습한다. 알려진 해 1-x²의 훈련 표본은 제공하지 않는다. 입력 방향을 두 번 쌓으면 2차 방향 미분을 구하며, 서로 다른 방향이면 혼합 미분을 구한다. 그 결과는 여전히 네트워크 가중치에 의존하므로 derivativeVjp가 출력 민감도를 통해 상태까지 역전파한다. 주변 예측값을 빼는 방식으로 2차 미분을 근사하지 않는다. 매끄러운 방정식에 ReLU의 꺾임점은 고전적 2차 미분이 없으므로 예제는 tanh를 사용한다. 표본에서 잔차가 작다는 사실은 훈련 근거이지 전체 구간의 보장은 아니다. 전체 구간 보장에는 균일 잔차 경계와 해 오차를 연결하는 정리가 필요하다.

정밀도 있는 학습에서는 샘플, 모델 상태, 예측값, 기울기, 손실에 선택한 스칼라 형식을 유지할 수 있고 체크포인트도 스칼라 워드를 보존한다. 다만 초기화된 Float 매개변수를 넓혀도 새 난수 비트가 생기지 않으며, 최적화기 설정도 변환 전에 Float일 수 있다. 높은 정밀도의 학습률이 필요하면 선택한 형식으로 전달하는 nn.sgdStep을 사용할 수 있다. 복소수 텐서는 실수부와 허수부를 같은 원소 형식에 저장하며, 복소 회귀 예제는 두 좌표를 미분해 실수 제곱 잔차를 최소화한다. 실수 손실의 기울기를 정칙 복소 미분과 동일시하지 않는다.

JAX의 추적 표현인 Jaxpr와 지원되는 PyTorch 그래프를 TorchLean의 텐서 중간 표현(IR)으로 가져올 수 있다. 가져오기 과정은 형상과 정밀도 메타데이터를 검사하고 그래프에 Lean 의미를 부여한다. 임의의 Python 프로그램을 번역하는 도구는 아니다. 예제는 sum(x*x) 그래프의 실수 의미를 텐서 제곱합과 연결하고, 모든 허용 길이와 순열에 대해 순열이 그래프 결과를 보존함을 증명한다. 부동소수점에서는 합산 순서가 바뀌면 결과가 달라질 수 있다. FloatLib의 interval 전술은 양의 밑 거듭제곱, 지정 밑 로그, 역쌍곡 함수 등을 포함한 실수식의 구간을 다룬다. 분류기 검증은 지원되는 인증서와 건전성 가정이 있을 때 입력 영역의 교란에도 라벨이 유지되는지 확인한다. converges 전술은 등록된 수축·경사하강 정리를 적용하되 학습률, 매끄러움 등의 가정을 유지한다. 손실 곡선이 감소했다는 사실만으로 수렴을 추론하지 않는다.

독자는 완화 갱신식이나 사용자 정의 손실을 바꿔 볼 수 있다. 실행 API를 그대로 쓸 수 있지만, 닫힌형 해와 미분 증명은 새 함수에 맞춰야 한다. 참고 자료와 전체 API는 사용자 정의 텐서 계산 안내, TorchLean 컴파일러 소스, FloatLib, 모델 예제에서 확인할 수 있다. 전체 구현 및 예제 링크에는 텐서 장, 부동소수점 장, 정밀도 API, 미분 전술 예제, 검증 API가 포함된다.

Robert Joseph George, 「Writing tensor programs in Lean」, TorchLean Verified Examples, 2026. 원문

부록: 원문 예시

text
import NN.API
open TorchLean

def model :=
 nn.Sequential![
 nn.linear 2 8,
 nn.relu,
 nn.linear 8 1
 ]

def trainer := Trainer.new model
 { objective := Trainer.Objective.mse
 optimizer := optim.adam { learningRate := 0.03 }
 device := cpu }
text
import NN.Tactic.Autograd

example (x : ℝ) : HasDerivAt (fun y : ℝ => Real.exp (y * y))
 (2 * x * Real.exp (x * x)) x := by
 autograd