서론

EZKL로 뉴럴 네트워크 (Neural Network)를 영지식 증명 (ZKP)화하려고 하면, 모델이 조금만 커져도 메모리가 부족해집니다.

그래서 사용하는 것이 Proof Splitting입니다. 큰 모델을 여러 개의 작은 회로 (Circuit)로 분할하여 증명을 생성하고, 마지막에 하나로 연결하는 메커니즘입니다.

Proof Splitting이란

거대한 뉴럴 네트워크를 하나의 회로로 다루는 것이 아니라, 여러 개의 작은 회로로 분할하여 증명을 생성하는 기능입니다.

장점은 두 가지입니다.

메모리 절감: 분할하여 피크 메모리 (Peak Memory)를 억제할 수 있으므로 한 대의 머신에서 처리 가능 -
분산 처리: 여러 머신에 부하를 분산하여 파이프라인 (Pipeline) 처리 가능

리포지토리

2가지 결합 방법

분할된 회로끼리 연결하는 (Stitch) 방법이 2가지 있습니다.

전체 파이프라인

PyTorch 모델
→ ONNX 내보내기 (network.onnx)
→ 모델 분할 (network_split_0.onnx / network_split_1.onnx)
...

구현

1. 모델 정의 및 ONNX 내보내기

class MyModel(nn.Module):
def __init__(self):
super().__init__()
...

2. ONNX를 2개로 분할

onnx.utils.extract_model을 사용하여 분할 위치를 지정합니다. Netron (https://netron.app/)에서 ONNX의 중간 노드 이름을 확인한 후 자릅니다.

import onnx
# 서브 모델 0: 입력 → conv1 + relu 까지
onnx.utils.extract_model(
...

3. 각 서브 모델의 셋업 (Setup)

async def setup(i):
model_path = f'network_split_{i}.onnx'
settings_path = f'settings_split_{i}.json'
...

매번 파일로 쓰는 이유는 두 가지입니다.

메모리 관리: 모든 서브 모델의 데이터를 변수로 유지하면 OOM (Out of Memory)이 발생함. 불필요한 데이터는 디스크로 보냄 -
분산 처리 상정: 머신 A가 서브 모델 0을 처리하여 파일을 저장하고, 머신 B가 그것을 받아 서브 모델 1을 처리하는 구성에 대응 가능

4. 증명 생성 및 연결

이 부분이 Proof Splitting의 핵심입니다.

def prove_model(i):
proof_path = f'proof_split_{i}.json'
witness_path = f'witness_split_{i}.json'
...

swap_proof_commitments가 필요한 이유

분할 증명의 가장 큰 함정은 부정 (Cheat)의 리스크입니다.

서브 모델 0: 올바른 입력으로 출력 A를 계산했다고 증명
서브 모델 1: 임의로 유리한 다른 입력 B를 사용하여 최종 출력을 계산

이렇게 되면 모델 전체로서 올바른 계산이 수행된 것이 아니게 됩니다.

swap_proof_commitments는 "전 단계의 출력과 다음 단계의 입력이 완전히 일치함"을 증명에 포함시키는 처리입니다. 이를 통해 분할된 증명이 변조 없이 바통 터치되었음을 검증자가 확인할 수 있습니다.

요약

큰 모델을 EZKL에 올릴 때는 Proof Splitting이 현실적인 선택지가 됩니다. 메모리가 부족해진다면, 우선 Netron을 사용하여 모델을 어디에서 분할할 수 있는지 확인하는 것부터 시작해 보는 것도 좋습니다.

저자의 다른 실험

LiDAR 점군 (Point Cloud) 분석, 블록체인 × AI 등 다른 프로젝트들도 정리해 두었습니다.

Discussion