노트북 04 — 다항식 환(polynomial ring)과 NTT#
목표: ML-KEM이 행렬 대신 다항식을 사용하는 이유와 NTT를 가능하게 만드는 요소를 이해합니다.
지금까지 각 LWE 샘플은 n개의 스칼라 (a, b)였습니다. ML-KEM은 대신 R_q = Z_q[x]/(x^n + 1)의 다항식을 사용합니다. 단일 다항식이 n개의 계수를 저장하고, 한 번의 다항식 곱셈은 구조화된 행렬-벡터 곱과 동등합니다. Number-Theoretic Transform (NTT)을 가능하게 하는 것이 바로 이 환(ring) 구조입니다.
import numpy as np
import time
import matplotlib.pyplot as plt
from pqc_edu.polyring import Poly
from pqc_edu.ntt import ntt, intt, poly_mul_ntt, Q, N
손으로 만드는 예제 — Z_17[x]/(x^4 + 1)에서 두 다항식을 곱합니다. 반순환(anti-cyclic) 규칙 x^4 = -1 때문에 차수가 4 이상인 항은 부호가 뒤집혀 되돌아옵니다.
# (1 + 2x) * (3 + x^3) in Z_17[x]/(x^4 + 1)
a = Poly(np.array([1, 2, 0, 0]), 17)
b = Poly(np.array([3, 0, 0, 1]), 17)
c = a * b
print('a =', a.coeffs)
print('b =', b.coeffs)
print('a*b =', c.coeffs)
# Verify by hand: (1+2x)*(3+x^3) = 3 + 6x + x^3 + 2x^4
# = 3 + 6x + x^3 + 2*(-1) (since x^4 = -1)
# = 1 + 6x + 0x^2 + x^3 in Z_17
print('expected: [1, 6, 0, 1]')
a = [1 2 0 0]
b = [3 0 0 1]
a*b = [1 6 0 1]
expected: [1, 6, 0, 1]
NTT 는 Z_q 위에서의 Fast Fourier Transform입니다. 다항식 곱셈을 \(O(n^2)\) 에서 pointwise 곱셈 \(O(n)\) + 두 번의 변환 \(O(n \log n)\) 으로 바꿉니다. 요구 조건: q가 prime이고 n | (q - 1). ML-KEM의 경우 q = 3329 = 13*256 + 1 이므로 256-th root of unity가 존재합니다. 사용되는 primitive root는 17입니다.
# Forward/inverse NTT roundtrip
rng = np.random.default_rng(42)
p = Poly(rng.integers(0, Q, N), Q)
p_recovered = intt(ntt(p))
print('NTT roundtrip matches:', p == p_recovered)
NTT roundtrip matches: True
# Timing: schoolbook vs NTT
rng = np.random.default_rng(0)
a = Poly(rng.integers(0, Q, N), Q)
b = Poly(rng.integers(0, Q, N), Q)
iters = 50
t0 = time.time()
for _ in range(iters):
_ = a * b # schoolbook
t_school = (time.time() - t0) / iters
t0 = time.time()
for _ in range(iters):
_ = poly_mul_ntt(a, b)
t_ntt = (time.time() - t0) / iters
print(f'schoolbook: {t_school*1000:.2f} ms/op')
print(f'NTT-based : {t_ntt*1000:.2f} ms/op')
print(f'speedup : {t_school/t_ntt:.1f}x')
schoolbook: 0.39 ms/op
NTT-based : 1.45 ms/op
speedup : 0.3x
왜 NTT가 더 느리게 나왔을까요? NTT 기반 곱셈의 점근적(asymptotic) 비용은 schoolbook의 \(O(n^2)\) 대비 \(O(n \log n)\)입니다. 그러나 우리가 만든 교육용 NTT는 가독성을 위해 명시적인 Python 루프로 작성되었고, schoolbook 곱셈은 계수마다 단일 벡터화된 numpy broadcast를 사용합니다. 실제 프로덕션 Kyber/ML-KEM에서는 NTT가 C나 손으로 최적화된 어셈블리로, 미리 계산된 zeta를 써서 작성됩니다 — 그런 환경에서는 n=256에서 schoolbook보다 약 10–20배 더 빠릅니다. 핵심은 점근적, 알고리즘적인 이야기입니다: n을 두 배로 늘려도 NTT는 거의 느려지지 않지만(butterfly 레이어 한 층 추가), schoolbook 비용은 4배가 됩니다. 우리의 측정은 Python의 상수 인수 비용을 반영한 것이지 근본 알고리즘을 반영한 것이 아닙니다.
핵심 요약: ML-KEM에서 피연산자는 NTT 도메인에 저장됩니다. 덧셈은 계수별로 수행되며(두 도메인에서 동일합니다), 곱셈은 pointwise_mul이 됩니다. 시간 도메인 다항식이 필요해지는 마지막 단계에서만 \(O(n \log n)\)의 INTT 비용을 지불합니다. 다음: 이 기계를 암호 시스템으로 바꾸는 법 → 05_ring_lwe.ipynb.