CKKS — Rescaling

Table of Contents

This is part of a tutorial series on the CKKS homomorphic encryption scheme.

In the last article, we described the key switching and relinearization algorithms. Relinearization allowed us to restore the natural “two tuple” ciphertext structure after a homomorphic multiplication, but we did not address two more problems produced by multiplication. The first is that, because CKKS uses fixed-point arithmetic on its underlying plaintexts, multiplication produces a ciphertext encoded with a squared scaling factor $\Delta^2$ relative to the operands which have scale $\Delta$. The second is that multiplication error grows too quickly: one error term is a product of the errors of the operands, and the other term looks like $\Delta \cdot e$.

This article solves both of these challenges with a technique called rescaling or modulus switching.¹

Also, this article will be the one that finally motivates us to implement the residue number system (RNS) form of CKKS, which we will save for the next article.

The code for this article is in this pull request.

Rescaling mathematically

Reiterating from previous articles, our situation is we have a ciphertext $c = (c_0, c_1)$ which encrypts a plaintext polynomial $m(x)$ with scaling factor $\Delta^2$. As we mentioned in the Notes section of the multiplication article, because this ciphertext is the output of a multiplication (and a relinearization), it doesn’t have a specific analytic form. Instead, we have the guarantee that

for the secret key polynomial $s$ and some sufficiently small-normed error polynomial $e$. We want a rescaling procedure that can take this ciphertext and produce a new ciphertext $c' = (c_0', c_1')$ for which

Here $m$ is the unscaled plaintext polynomial, to make it clear how the scaling factors evolve.

Note the scaling factor after the procedure must be $\Delta$ instead of $\Delta^2$, and the error polynomial is some new error $e'$ that will depend in some way on the original error $e$ and the procedure we take.

The simplest idea is to just divide the ciphertext by $\Delta$.

The problem is that dividing by $\Delta$ may not be a valid ring operation. For example, if the coefficient modulus $Q$ had a common prime factor with $\Delta$, then $\Delta$ would not have an inverse in the ring $\mathbb{Z}/Q\mathbb{Z}$. But even if it does, multiplying by $\Delta^{-1}$ would make the noise unacceptably large: $\Delta^{-1} e$ is not guaranteed to be small mod $Q$, and in general it is large.

A short script to demonstrate:

Q = 2**60 - 93
delta = 2**40
delta_inv = pow(delta, -1, Q)

def centered(x):
    x %= Q
    if x > Q / 2:
        return x - Q
    return x

for e in [1, -2, 3, 5, -7]:
    print(f"e = {e:3d},  delta^-1 * e mod Q = {centered(delta_inv * e):>20d}")

Which outputs

e =   1,  delta^-1 * e mod Q =   -12397005425868799
e =  -2,  delta^-1 * e mod Q =    24794010851737598
e =   3,  delta^-1 * e mod Q =   -37191016277606397
e =   5,  delta^-1 * e mod Q =   -61985027129343995
e =  -7,  delta^-1 * e mod Q =    86779037981081593

Instead of hoping for a ring operation, we charge ahead and divide by $\Delta$ as real numbers. In order to make this work, let’s analyze the phase of the ciphertext polynomial $c_0 + c_1 s$ after dividing by $\Delta$. Since we’re not in the ring anymore, let’s lift the mod $Q$ equation to a particular representative. So given

for some $t \in \mathbb{Z}[x] / (x^N + 1)$, if we divide by $\Delta$ we get

Two things stand out. The left-hand-side terms $c_0 / \Delta$ and $c_1 / \Delta$ may not be integers, so we have to round the values, which complicates it slightly. The second thing is that this looks like a valid equation in a different ring: $(\mathbb{Z}/(Q/\Delta)\mathbb{Z})[x] / (x^N + 1)$.

That is, if we could take this equation modulo $Q/\Delta$ instead of modulo $Q$, we’d have exactly what we need: the scaling factor would be restored to $\Delta$, the error would be scaled down by $\Delta$, and the term $t Q / \Delta$ would vanish modulo $Q / \Delta$.

So let’s make a new assumption: require $Q = \Delta \cdot Q'$ for some $Q'$; then our candidate rescaling procedure is as follows:

Let’s analyze the correctness and error growth of this procedure, a little bit more formally than above.

First, in order to analyze the rounding operation, we have to lift the rescale operation inputs to representatives $c_0 + t_0Q$ and $c_1 + t_1Q$ for some $t_0, t_1 \in \mathbb{Z}[x] / (x^N + 1)$. Then consider $c_0' = \textup{round}\left(\frac{c_0 + t_0Q}{\Delta}\right)$. We can write this as $c_0' = \frac{c_0 + t_0Q}{\Delta} + z_0$, where $z_0$ is a polynomial whose coefficients are in $[-1/2, 1/2]$. Similarly define $z_1$ for $c_1'$. Then,

Now the question is how big is this new error polynomial $\frac{e}{\Delta} + z_0 + z_1 s$? I want to avoid doing too much error analysis (saving that for a dedicated article). But to provide the intuition, because of how $z_0$ and $z_1$ are defined as rounding errors, and because the values of $c_0$ and $c_1$ are pseudorandom, the coefficients of $z_0$ look like independent uniform random variables in $[-1/2, 1/2]$ with variance $1/12$. Likewise, the product $z_1 s$ is a negacyclic polynomial product of a uniform random polynomial and a small-normed (ternary-coefficient) polynomial. This product $z_1 s$ is the main issue here, so I’ll ignore $z_0$ for simplicity. Taking this polynomial product flips some signs (which can be ignored due to symmetry in the randomness) and sums coefficients from $z_1$ corresponding to nonzero coefficients of $s$, and due to sign differences they will often cancel out on average, leaving coefficients of size roughly $\sqrt{h/12}$, where $h$ is the hamming weight of $s$.² After applying the canonical embedding (the step of decoding before removing the scale), this error grows to roughly $\sqrt{Nh/12}$.

And likewise, the term $e / \Delta$ is the original error $e$ scaled down by $\Delta$. The error is the same relative to the new modulus size, but if you recall from our previous article on multiplication, after the multiplication op there is an error term of the form $\Delta (m_c e_d + e_c m_d) + e_c e_d$, and after relinearization there was an additional $B \log(Q) e'$ term, for $e'$ depending only on the error in the key-switching key. After rescaling, this error is all divided by $\Delta$. So the first term becomes a linear error growth rate (while still scaled by the size of the underlying messages), and both the second term $e_c e_d$ and the third term $B \log(Q) e'$ become negligible as they are both very small relative to $\Delta$.

Rescaling in code

Implementing rescaling in code adds a small challenge: we need a scaling factor that divides $Q$ exactly. So this code outline will actually start with the test setup, in this commit.

We start by defining two encoding params. The first is for the output of the test: a scale of Delta = 2**20 and a $Q'$ that is close to $Q / \Delta$. The second is for the input to the test: a scale of $\Delta^2$ (as if freshly multiplied and relinearized) and the original $Q$ modulus.

# rescaling_test.py
def test_rescale():
    rescaled_encoding_params = EncodingParams(
        scale=2**20,
        poly_modulus_degree=8,
        # This prime is close to NTT_64_BIT_PRIME / 2**20
        coefficient_modulus=2199023255521,
    )
    squared_encoding_params = EncodingParams(
        scale=2**40,
        poly_modulus_degree=8,
        coefficient_modulus=NTT_64_BIT_PRIME,
    )

    message = np.array([1, 2, 3, 4], dtype=np.complex64)

    ...

Note $Q'$ is not exactly $Q/\Delta$. We used the exact $Q / Q'$ divisor to ensure the error term from the “lifted representative” vanishes mod $Q'$. But $Q / Q'$ is not $\Delta$. It doesn’t need to be exact; if it’s close enough then this only causes the true fixed-point scale to differ from the scale we’re tracking. This results in a small additional error at decoding time, but because decoding already naturally incurs error, if this error is small enough it is acceptable.

More importantly, after this section we will ditch this approach in favor of the residue number system (RNS) formalism, which requires a substantial rewrite of the core data types.

Proceeding with that understanding, we construct the context for encryption using the $Q$ and $\Delta^2$ context. Then we rescale, passing in the appropriate values to use for the output ciphertext’s metadata.

# rescaling_test.py
def test_rescale():
    ...

    context = CKKSContext(rescaled_encoding_params)

    sq_context = CKKSContext(squared_encoding_params)
    sk = sq_context.generate_symmetric_private_key()
    pt = encode(message, squared_encoding_params, sq_context.ntt_params)
    ct = sq_context.encrypt_symmetric(pt, sk)

    ct_rescaled = rescale(
        ct,
        rescaled_encoding_params.coefficient_modulus,
        context.ntt_params,
    )

    decrypted_product = context.decrypt_symmetric(ct_rescaled, sk)
    decoded_product = decode(decrypted_product, rescaled_encoding_params)

    np.testing.assert_allclose(decoded_product, message, rtol=0, atol=0.1)

A similar test in the same file chains together multiplication, relinearization, and rescaling, being careful to thread through the correct metadata as here.

With the proper test setup, the implementation of rescale is rather straightforward (in the same commit):

# rescaling.py
import numpy as np

from ckks_types import Ciphertext
from polynomial import ModQPolynomial
from params import NTTParams


def rescale_poly(
    poly: ModQPolynomial, new_modulus: int, new_ntt_params: NTTParams
) -> ModQPolynomial:
    # Note this only works because the NTT_PRIME fits comfortably in a 64-bit float.
    rescale_factor = np.float64(poly.coefficient_modulus) / np.float64(new_modulus)
    new_coeffs = np.round(poly.coefficients / rescale_factor)
    return ModQPolynomial(
        coefficients=new_coeffs,
        coefficient_modulus=new_modulus,
        modulus_degree=poly.modulus_degree,
        ntt_params=new_ntt_params,
    )


def rescale(ct: Ciphertext, new_modulus: int, new_ntt_params: NTTParams) -> Ciphertext:
    d0 = rescale_poly(ct.data[0], new_modulus, new_ntt_params)
    d1 = rescale_poly(ct.data[1], new_modulus, new_ntt_params)
    return Ciphertext(data=(d0, d1))

The slight deviation from the mathematical formulas is that we’re dividing not by $\Delta$, but by $Q / Q'$, which as we said above is close enough to $\Delta$. There is a lot more to say about this discrepancy, rescaling with an exact scaling factor versus an inexact one, and I plan to return to this in a future article because not all CKKS implementations take the same approach. But to understand the approaches they do take, we finally have to address the residue number system.

Modulus budget

Taking a step back: in our CKKS we now have the ability to add ciphertexts with linear noise growth, and to multiply ciphertexts (plus relinearize and rescale) such that the relative effect of noise growth is also roughly linear.

However, the cost of multiplying two ciphertexts and rescaling is that our modulus $Q$ becomes smaller by a factor of $\Delta$. Even if we were able to make $\Delta$ quite small for our application, say 10 bits, we’d still only be able to do about 6 sequential multiplications before a 64-bit modulus $Q$ is exhausted.

This incentivizes us to introduce the concept of a modulus budget. Recall, we get to pick the modulus $Q$, with only a few constraints.³ So we can pick $Q$ to be quite large, say 800 bits, and that would allow for many more sequential multiplications. Ignoring momentarily that the choice of $Q$ (and $N$) impacts cryptographic security, the problem with choosing such a huge $Q$ is performance. All the operations that previously fit in a single 64-bit int or float now require big integer math, and big integer math is slow (and particularly bad for SIMD hardware acceleration) because it has serial data dependencies across the digits of the integer being processed.^{4 5}

Thankfully, it’s known how to modify big integer math so it can be executed in parallel: you use the residue number system (RNS). Briefly: you write $Q$ as a product of distinct primes $Q = \prod_i q_i$, and then, via Sunzi’s theorem, any number mod $Q$ can be represented by its residues mod each $q_i$. A single 800-bit integer becomes an array of 13 64-bit integer residues called limbs, and a polynomial of degree 1024 becomes a two-dimensional array of shape 13x1024.

The theory behind the RNS form ensures that any ring operation applied identically to all RNS limbs corresponds to the same operation applied to the re-composed number mod $Q$. In other words, $\mathbb{Z}/Q\mathbb{Z}$ is isomorphic as a ring to the ring product $\prod_i \mathbb{Z}/q_i\mathbb{Z}$, and the same holds for the polynomial ring $(\mathbb{Z}/Q\mathbb{Z})[x]/(x^N+1)$.

In this frame, the process of rescaling corresponds to dropping a limb, which can be implemented in terms of slicing a tensor to exclude the last limb in each coefficient of a polynomial. As such, we can refine our modulus budget from tracking bits to tracking how many RNS limbs remain. The number of limbs remaining is also called the “level” of a ciphertext, and so the “modulus budget” or “limb budget” is often called a “level budget.”⁶

Next

Next time we’ll rewrite the CKKS implementation so as to incorporate RNS, and allow us to use much larger moduli. This will also likely require us to revisit the performance issues with our current implementation of NTT.

Until then!