CKKS — Key Switching and Relinearization

Table of Contents

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

In the last article, we described the ciphertext-ciphertext addition and multiplication operations. We ended with two challenges, one to restore the squared scaling factor $\Delta^2$ back to $\Delta$, and the other to convert a three-tuple ciphertext (decrypting under $(1, s, s^2)$) back to a two-tuple ciphertext (decrypting under $(1, s)$).

This article covers the latter challenge, which is called relinearization. Along the way, we’ll see how relinearization is a special case of a more general technique called key switching. In implementing key switching, we’ll hit yet another obstacle: the noise growth from a naive key switching procedure is too large. To solve it, we’ll apply another clever technique called the gadget decomposition.

As with the last few articles, we are still avoiding the RNS form for the clarity of the algorithmic ideas.

The code for this article is in this pull request.

Plaintext-ciphertext operations

Before diving into key switching, we implement plaintext-ciphertext operations. These are needed to implement key switching.

In general, when converting a program to run under FHE, you encounter static constants and runtime data known to the server (i.e., not a private input to the program). I will call both of these cases “cleartext data.” Think of a large set of trained ML model weights, for example, or the coefficients of a polynomial that you want to evaluate on an encrypted input. Literally encrypting cleartext data under the user’s public key is so wasteful as to be infeasible, because you would need a separate copy for each user or else to perform that work at runtime.

However, most FHE schemes support plaintext-ciphertext operations which are much cheaper than the corresponding ciphertext-ciphertext operations. This allows one to prepare cleartext data as plaintexts offline, and then load and use them efficiently at runtime.

The trick is to first encode your cleartext datum $m$ as a plaintext, as you would when performing encryption. This produces a polynomial $p(x)$ in the ring $(\mathbb{Z}/Q\mathbb{Z})[x] \Big / (x^N+1)$ with a particular scaling factor $\Delta$.¹ Then you put the plaintext $p(x)$ into a tuple $(p, 0)$.

This tuple can be thought of as a trivial sort of ciphertext. That is, under the decryption procedure, $(p, 0)$ decrypts as $p\cdot 1 + 0 \cdot s = p$ for every secret key $s$. This is equivalent to a ciphertext with a random sample that happens to be zero (and with zero noise).

Because this is a valid (if unluckily sampled) ciphertext, the same formulas we saw last time apply, and they pleasantly simplify. In particular, ciphertext-plaintext addition becomes $(p, 0) + (c_0, c_1) = (p + c_0, 0 + c_1) = (p+c_0, c_1)$. Likewise, the tensor multiplication formula for ciphertext-ciphertext multiplication simplifies with a plaintext operand to $(p, 0) \otimes (c_0, c_1) = (pc_0, 0c_0 + pc_1, 0c_1)$, and we can simply drop the last coordinate to get $(p c_0, p c_1)$.

In code (this commit):

# ckks_types.py
PtOrCt = Plaintext | Ciphertext

# addition.py
from ckks_types import Ciphertext, Plaintext, PtOrCt

def add(v1: PtOrCt, v2: PtOrCt) -> Ciphertext:
    match v1, v2:
        case Ciphertext((c0, c1)), Ciphertext((d0, d1)):
            return Ciphertext(data=(c0 + d0, c1 + d1))
        case Ciphertext((c0, c1)), Plaintext() as p:
            return Ciphertext(data=(c0 + p, c1))
        case Plaintext() as p, Ciphertext((c0, c1)):
            return Ciphertext(data=(c0 + p, c1))
        case _:
            raise ValueError()

# multiplication.py
def mul_plain(v1: PtOrCt, v2: PtOrCt) -> Ciphertext:
    match v1, v2:
        case Ciphertext((c0, c1)), Plaintext() as p:
            pass
        case Plaintext() as p, Ciphertext((c0, c1)):
            pass
        case _:
            raise ValueError()

    return Ciphertext(data=(p * c0, p * c1))

The error analysis is also relatively clean. For addition, there is no additional noise introduced, because the plaintext operand’s “ciphertext” form had zero noise. For multiplication, the product decrypts under $s$ to $p(c_0 + c_1 s) = p m + p e$, where $e$ was the original error. Hence, the error is scaled by the plaintext $p$.

One small thing to note: as with ciphertext-ciphertext multiplication, the scaling factor on $p$ is part of the error growth. For this reason, some CKKS implementations require the rescaling procedure (in the next article) after each plaintext-ciphertext multiplication. Others use tricks with the RNS form to avoid an explicit rescale step. We will come back to these ideas in a future article.

Interpreting relinearization as key switching

Recall the problem with ciphertext-ciphertext multiplication: the product of $(c_0, c_1) \otimes (d_0, d_1)$ is a three-tuple $(t_0, t_1, t_2) = (c_0 d_0, c_0 d_1 + c_1 d_0, c_1 d_1)$ which “encrypts” under a dot product with $(1, s, s^2)$ to the right product message, but we want to convert this back to a two-tuple that encrypts the same message under $(1, s)$.

Now, because we framed the decryption circuit as a dot product, and because dot products are linear, we can split the dot product back up into two-tuples:

This is decrypting $(t_0, t_1)$ under $s$, and adding that to the “decryption” of $(0, t_2)$ under $s^2$. However, these two pieces are a bit strange as ciphertexts. They don’t really encrypt anything useful by themselves, but, whatever they decrypt to, the sum of those two things gives you back your original plaintext (plus error).

Now here is the clever idea: what if we could find a way to convert $(0, t_2)$ to a new ciphertext that encrypts the same plaintext (whatever it is) but back under the key $s$? Then, the equation above would become

where $(x, y)$ is the output of this hypothetical procedure. Then we can recombine the inner product to get

So the problem of relinearization is reduced to this more general problem of key switching.

Key switching, naively, and the noise problem

Restating the problem from the last section, we replace the $s^2$ key with an arbitrary key $s'$.

Key switching problem: Let $\mathbf{c} = (c_0, c_1)$ be an encryption of the plaintext $p$ under the secret key $s'$, with error $e$. Find a way to convert $\mathbf{c}$ to a new ciphertext $\mathbf{d} = (d_0, d_1)$ encrypting the same plaintext $p$ under a new secret key $s$.

The key switching procedure is perhaps the first deeply novel idea in this tutorial. Indeed, without giving something extra, it is clearly impossible to solve this problem. The one tool at our disposal is the ability to perform homomorphic operations.

So here’s the idea: we homomorphically evaluate the decryption circuit for $s'$ using encryption under $s$.

That decryption circuit is $(c_0, c_1) \mapsto c_0 + c_1 s'$. Evaluating this homomorphically would produce $p + e$ encrypted under $s$, which is exactly what we want. Evaluating it homomorphically also requires an encryption of $s'$ under $s$. We call this a key switching key.² With that in hand, the decryption circuit can be lifted to plaintext-ciphertext operations, as the values $c_0, c_1$ can be treated as plaintexts (even though we don’t consider them as having an encoding or scaling factor).

Let $(b, a)$ be an encryption of $s'$ under $s$, with noise $e'$.

And decrypting this under $(1, s)$ gives

At the end we see the original error $e$ from the input ciphertext $(c_0, c_1)$, but we also have this term $c_1 e'$. That is, the error from the key switching key is scaled by $c_1$. That’s a problem because $c_1$ is large in general. For a fresh symmetric encryption, it is a sampled polynomial with uniformly random coefficients mod $Q$ (of magnitude $Q/2$ in the centered representative). Without diving too deeply into the error analysis (as a general error analysis depends on the operations in each program!), one could heuristically estimate its size as $\sqrt{N} \frac{Q}{c} \sigma$, where $N$ is the polynomial degree, $\sigma$ is the error distribution parameter, and $c$ is some constant. For basically any $N$, this is already much larger than $Q/2$ and corrupts the message.

In code (this commit), first we define the key switching key type as a wrapper around a ciphertext. Later in this article, this type will be upgraded to a collection of ciphertexts, so having the explicit type will make that easier.

# ckks_types.py
@dataclass(frozen=True)
class KeySwitchingKey:
    """A key switching key."""
    data: Ciphertext

Then we generate the key switching key.

# key_gen.py
def generate_symmetric_key_switching_key(
    source_key: PrivateKey,
    target_key: PrivateKey,
    params: EncryptionParams,
    ntt_params: NTTParams,
    random_source: rng.RandomSource,
) -> KeySwitchingKey:
    """Generate a key switching key from source_key to target_key.

    This key can be used by key_switch to convert a ciphertext
    encrypted under source_key to be encrypted under target_key.
    """
    ct = encrypt_symmetric(
        source_key.to_mod_q(ntt_params), target_key, params, random_source
    )
    return KeySwitchingKey(data=ct)

Because the secret key type is an IntPolynomial, we added a small helper to convert it to a ModQPolynomial. Finally, we implement key_switch with a TODO for the noise fix.

# key_switch.py
from ckks_types import Ciphertext, KeySwitchingKey

def key_switch(ct: Ciphertext, ksk: KeySwitchingKey) -> Ciphertext:
    """Switch `ct` to encrypt the same plaintext under a new key."""
    # TODO: enhance naive key switching to reduce error growth
    c0, c1 = ct.data
    b, a = ksk.data.data
    return Ciphertext(data=(c0 + c1 * b, c1 * a))

To test this, we add a small helper function extract_noise to extract the raw noise from a ciphertext, then in the test we explicitly construct the error formula c_1 e' as above.

def test_key_switch_symmetric_noise_formula():
    seed = 837465
    message = np.array([1, 2, 3, 4], dtype=np.complex64)
    context = CKKSContext(DEFAULT_ENCODING_PARAMS)
    context.rng = SeededRandomSource(seed)

    source_key = context.generate_symmetric_private_key()
    target_key = context.generate_symmetric_private_key()
    pt = context.encode(message)
    ct = context.encrypt_symmetric(pt, source_key)
    ksk = context.generate_symmetric_key_switching_key(source_key, target_key)

    switched_ct = key_switch(ct, ksk)
    decrypted = context.decrypt_symmetric(switched_ct, target_key)
    decoded = decode(decrypted, DEFAULT_ENCODING_PARAMS)

    # Due to the lack of a gadget decomposition,
    # the noise has overwhelmed the message and it decrypts to garbage.
    # np.testing.assert_allclose(decoded, message, rtol=0, atol=0.2)

    # For now, extract the exact noise formula, subtract it from the plaintext,
    # and compare what remains. This shows consistency of our impl with the
    # algebra and noise model.
    c1 = ct.data[1]
    c1_eprime_term = c1 * extract_noise(
        ksk.data, target_key, source_key.to_mod_q(context.ntt_params)
    )

    ks_noise_removed = decrypted - c1_eprime_term
    decoded = decode(ks_noise_removed, DEFAULT_ENCODING_PARAMS)
    np.testing.assert_allclose(decoded, message, rtol=0, atol=0.2)

So the key switching idea requires another trick to ensure that the noise growth is explicitly bounded.

Key switching, with gadget decomposition

The technique to reduce noise growth during key switching (and other CKKS procedures) is called the gadget decomposition.³ I wrote about this at length in 2021-2022, both to describe a generalized version and to apply it to LWE key switching. Key switching in CKKS is different in general, so I will repeat the relevant subset of those articles here in context.

The key insight is that, when we have a product like $x \times \textup{Enc}(y)$ and we control how $y$ is encrypted, we can “move” some of the magnitude of $x$ into $y$ before it’s encrypted. That has the effect of reducing error growth corresponding to how much we move over to $y$.

Making that vague idea concrete: the “magnitude” of $x$ that we can “move” corresponds to its number base. Pick a base $B$ and write $x = \sum_{i=0}^{l-1} x_i B^i$, where $l$ is the number of powers you need to represent $Q$ and $0 \leq x_i < B$ are the unique base-$B$ digits.⁴ When we multiply the plaintext $x$ homomorphically by $\textup{Enc}(y)$, these powers of $B$ account for a lot of the error growth. If we instead encrypt $B^i y$ for each $i$, then the same product can be written as

Encrypting the powers-of-$B$-times-$y$ operand of this dot product means that when we evaluate the dot product homomorphically, the plaintext $x_i$’s each contribute at most $B e'$ noise. With $l$ such terms added together, the total noise growth is $B \cdot l \cdot e' = O(B \log(Q)) e'$. For a 64-bit $Q$, if you pick $B=16$ and $l=16$ ($16^{16} = 2^{64}$), then the noise growth is only $B \cdot \log_B(Q) e' = 256 e'$ when before it was on the order of $Q e'$.

In code, we first implement and test the basic decomposition routine, absent any encryptions or polynomials, in this commit.

# decomposition.py
import numpy as np

DecomposedInt = np.ndarray
TOTAL_BIT_LENGTH = 64


def _shifts(num_digits: int, base_log: int) -> np.ndarray:
    return np.uint64(TOTAL_BIT_LENGTH) - (
        np.arange(1, num_digits + 1, dtype=np.uint64) * np.uint64(base_log)
    )


def decompose(x: np.uint64, base_log: int) -> DecomposedInt:
    """Decompose a number in a given power-of-2 base, most significant first."""
    if TOTAL_BIT_LENGTH % base_log != 0:
        raise ValueError(f"base_log={base_log} must divide {TOTAL_BIT_LENGTH}")
    num_digits = TOTAL_BIT_LENGTH // base_log
    digit_mask = np.uint64((1 << base_log) - 1)
    return (np.uint64(x) >> _shifts(num_digits, base_log)) & digit_mask


def recompose(digits: DecomposedInt, base_log: int) -> np.uint64:
    """The inverse of decompose."""
    digits = np.asarray(digits, dtype=np.uint64)
    return np.sum(digits << _shifts(digits.shape[0], base_log), dtype=np.uint64)


# decomposition_test.py

# an example test showing a base-16 decomp
def test_decompose_base_16():
    x = np.uint64(0x0123_4567_89AB_CDEF)
    expected = np.arange(16, dtype=np.uint64)
    np.testing.assert_array_equal(decompose(x, base_log=4), expected)

# an example test showing recomposition
def test_recompose_known_digits():
    digits = np.array([1, 2, 3], dtype=np.uint64)
    # Digits are most-significant first, aligned to the top of the 64-bit word.
    expected = (1 << 60) + (2 << 56) + (3 << 52)
    assert int(recompose(digits, base_log=4)) == expected

Next we apply it to key switching in this commit. First we extend the scalar decomposition routine to apply to a polynomial coefficientwise. The output of the decompose_poly routine below is a list of polynomials, one for each digit of the decomposition, where all the coefficients in the polynomial correspond to the same place value. The order is most-significant digits first.

We will use decompose_poly to decompose the $c_1$ component of the input to the key switching routine.

def decompose_poly(poly: ModQPolynomial, base_log: int) -> list[ModQPolynomial]:
    """Decompose a polynomial's coefficients."""
    if TOTAL_BIT_LENGTH % base_log != 0:
        raise ValueError(f"base_log={base_log} must divide {TOTAL_BIT_LENGTH}")
    num_digits = TOTAL_BIT_LENGTH // base_log
    digit_mask = np.uint64((1 << base_log) - 1)

    # Coefficients are in [0, Q), and for now Q fits in a 64-bit integer.
    coeffs = poly.coefficients.astype(np.uint64)

    # Shape (num_digits, N): row i holds digit i of every coefficient.
    digits = (
        coeffs[np.newaxis, :] >> _shifts(num_digits, base_log)[:, np.newaxis]
    ) & digit_mask

    return [
        ModQPolynomial(
            row.astype(np.int64),
            modulus_degree=poly.modulus_degree,
            coefficient_modulus=poly.coefficient_modulus,
            ntt_params=poly.ntt_params,
        )
        for row in digits
    ]

Now we need to ensure the key switching key is generated using powers $B^i s$ according to the same decomposition order (from most-significant digit to least-significant).

# key_gen.py
def _scaled_source_keys(
    source_key: PrivateKey,
    ks_params: KeySwitchingKeyParams,
    ntt_params: NTTParams,
) -> list[ModQPolynomial]:
    """Compute B^i * s mod Q for each digit place i = l-1, ..., 0."""
    B = 1 << ks_params.log_base
    num_digits = TOTAL_BIT_LENGTH // ks_params.log_base
    # Scale before reducing s mod Q: s is ternary, so (B^i mod Q) * s has
    # coefficients in {-(B^i mod Q), 0, B^i mod Q}, which fit in an int64,
    # whereas scaling the reduced coefficients (up to Q-1) would overflow.
    return [
        IntPolynomial(
            pow(B, i, ntt_params.modulus) * source_key.coefficients,
            source_key.modulus_degree,
        ).to_mod_q(ntt_params)
        for i in reversed(range(num_digits))
    ]


def generate_symmetric_key_switching_key(
    source_key: PrivateKey,
    target_key: PrivateKey,
    params: EncryptionParams,
    ks_params: KeySwitchingKeyParams,
    ntt_params: NTTParams,
    random_source: rng.RandomSource,
) -> KeySwitchingKey:
    """Generate a key switching key from source_key to target_key."""
    # Encrypt B^i * s for all i
    cts = [
        encrypt_symmetric(scaled_key, target_key, params, random_source)
        for scaled_key in _scaled_source_keys(source_key, ks_params, ntt_params)
    ]
    return KeySwitchingKey(data=cts, params=ks_params)

# asymmetric is similar, omitted for brevity.

Finally, the new key switch routine forms the relevant dot product and sums the terms together.

# key_switch.py
from ckks_types import Ciphertext, KeySwitchingKey
from decomposition import decompose_poly
from polynomial import ModQPolynomial

def key_switch(ct: Ciphertext, ksk: KeySwitchingKey) -> Ciphertext:
    """Switch `ct` to encrypt the same plaintext under a new key."""
    c0, c1 = ct.data
    ks_cts = ksk.data

    c1_digits = decompose_poly(c1, ksk.params.log_base)

    # Sum the plaintext-ciphertext products d_i * (b_i, a_i),
    # with an initializer of (c0, 0).
    switched_c0 = c0
    switched_c1 = ModQPolynomial.zeros_like(c1)
    for c1_digit, ks_ct in zip(c1_digits, ks_cts):
        ks_b, ks_a = ks_ct.data
        switched_c0 += c1_digit * ks_b
        switched_c1 += c1_digit * ks_a

    return Ciphertext(data=(switched_c0, switched_c1))

The next commit generates a plot of the measured noise growth of key switching, shown below.⁵

First, we compare noise growth for a fixed $N$ as $B$ changes.

The relative comparison is against an average-case statistical estimate.⁶ It seems the smaller decomposition bases are better for a fixed $N$, but a smaller choice of $B$ incurs other costs: the data operated on is larger and there are more numbers to crunch. In particular, using $B=2$ would incur 128 products per key switch, while $B=256$ incurs just 16.

Then, fixing $B$ and growing $N$:

While the noise growth is quite large in absolute terms, what matters for us is how it relates to the scaling factor. As we will see next time, the scaling factor will divide error terms by $\Delta$, making any error significantly less than $\Delta$ become error significantly less than 1. So a good rule of thumb for choosing $B$ is to make it as large as possible subject to the constraint that the resulting key switching noise is meaningfully smaller than $\Delta$.

There are additional tricks you can use to make this error growth even smaller. The articles I wrote for gadget decomposition and LWE key switching provide details, but briefly: you can use a signed digit decomposition to halve the bound $B$, and you can drop the least-significant digits of the decomposition, which reduces error by reducing the number of terms to compute, but introduces a new error from the truncation. The value of $B$ determines how many terms you can safely drop before the truncation error outweighs the savings.

However, those tricks are not used in standard CKKS implementations. In fact, once we add in the residue number system (RNS) form of CKKS, we will revisit the key switching routine to both adapt it for RNS, and also to introduce an additional trick: performing the key switch in a new “special” modulus $Q'$ which is divisible by $Q$, and then dividing by $Q'/Q$ restores the modulus $Q$ while dividing the error by $Q'/Q$. When this version of key switching is combined with gadget decomposition, we get what is called “hybrid key switching” in the literature. This further allows us to choose a much larger base $B$, which is why I’m not lingering too much on how to choose $B$.

For now, we’ll wrap up this article by implementing relinearization in terms of key switching.

Implementing relinearization

As we already know how relinearization reduces to key switching, the only remaining work is to implement it. This is done in this commit.

First we add the necessary setup step to generate the key switching key for $s^2$.

# context.py
class CKKSContext:

    ...

    def initialize_relinearization_key(self, sk: PrivateKey):
        self.relinearization_key = self.generate_symmetric_key_switching_key(
            sk * sk, sk
        )

Note that the relinearization key does not depend on whether the ciphertexts are encrypted with symmetric or asymmetric encryption. Using the private key in the asymmetric case works the same, and produces less error overall than doing a public-key encryption of the squared secret key.

Finally, implementing relinearization:

# multiplication.py
def relinearize(ct: Deg2Ciphertext, ksk: KeySwitchingKey) -> Ciphertext:
    """Convert a degree-2 ciphertext to a degree-1 ciphertext."""
    t0, t1, t2 = ct.data
    ks_input = Ciphertext(data=(ModQPolynomial.zeros_like(t2), t2))
    k0, k1 = key_switch(ks_input, ksk).data
    return Ciphertext(data=(t0 + k0, t1 + k1))


# context.py
class CKKSContext:

    ...

    def mul(self, ct1: Ciphertext, ct2: Ciphertext) -> Ciphertext:
        if not hasattr(self, "relinearization_key"):
            raise ValueError("Must generate relinearization key first")
        unrelinearized = mul(ct1, ct2)
        return relinearize(unrelinearized, self.relinearization_key)

This almost works. One remaining issue is that our original key switching key generation routine assumed a ternary source key, and now that we need to encrypt $s^2$, the _scaled_source_keys helper overflows int64 for large enough $N$. To fix it, in this commit, we more carefully scale the coefficients mod $Q$.

Next

With relinearization implemented, we have a full multiplication procedure in CKKS, though it has a multiplicative error growth rate and we still haven’t addressed the problem of multiplication squaring the scale $\Delta$.

Next time we will see a simple method that helps with both of these problems: rescaling. Studying rescaling will lead us to a small crossroads, namely how to handle moduli larger than 64 bits. Almost all CKKS implementations take the same approach, but I would like to dwell a bit on the difference and how it impacts performance before proceeding with the standard solution.