Foreign fields 5: Optimized multiplication of sums

[mitschabaude]

This PR introduces a heavily optimized way to constrain foreign field expressions involving the product of two values or sums of values on the left side and a sum on the right side. Example:

(x - y) * z = a + b

A close analysis of the ffmul gate shows that for expressions like these, the ffmul gate is valid even without doing range checks or bounds checks on the sums x - y and a + b (assuming that the original values x, y, a, b are range-checked as usual).

Note that this is violating the assumptions that went into the original soundness proof of the ffmul gate, in multiple ways. At the bottom, we sketch an argument for why these assumptions can be weakened.

The API we introduce for adding these multiplication constraints uses a new "lazy sum" type, which allows optimal chaining and constraint-savings without leaking any of the involved complexity to the developer:

let xMinusY = ForeignField.Sum(x).sub(y); // create a lazy sum
let aPlusB = ForeignField.Sum(a).add(b); // another one

// assert that (x - y) * z === a + b
// `assertMul()` accepts either lazy sums or plain foreign field elements
ForeignField.assertMul(xMinusY, z, aPlusB, f);

What happens under the hood

• Sums like a + b that go into the remainder are wired to the multiplication result directly, without any additional constraints. This has to do with the fact that the ffadd gate constrains its result a + b in compact limb form [r01, r2], and ffmul only uses the remainder in its compact form, so ffadd tells us enough about its output to be safely used in ffmul.
  • Note: To exploit this, we had to give the new API the form of an assertion (assertMul(x, y, z)) rather than a computation (mul(x, y) -> z).
• For sums like x - y which become a multiplication input, ffmul uses the two lower limbs individually, which aren't constrained by ffadd. Therefore, we add additional generic gates which perform roughly the ffadd logic, just on the lowest limb. This is enough to guarantee that ffadd-equivalent equations hold on all 3 limbs. This is cheaper than ensuring limbwise constraints using a multi-range check on x - y.
• For the left input sum, we ensure that ffadd is chained into the ffmul gate

Results

Empirically, using assertMul() for the example equation (x - y)(x + y) === x^2 - y^2 lowers the constraint count from 32 to 20 compared to a naive version with mul(), add() and sub(). Note that the naive version also suffers because it computes the left and right side separately, and so applies range-checks to the same result twice. This is another motivation for why an "assertion" API is better than a "computation" API.

These savings are crucial for EC addition, which essentially only consists of three calls to assertMul() with different types of sum.

Summary of changes

• adds ForeignField.assertMul(), ForeignField.Sum()
  • including tests
• new optimized gadget for constraining the carry to 0, 1 or -1 using generic gates
  • starting a new file of gadgets, basic.ts
• new Unconstrained API. to facilitate the generic gate computation in Sum, I had to do computations on an auxiliary bigint that are interleaved with provable code. I added the Unconstrained class as a safe wrapper for auxiliary values like this, trying to realize some of @nholland94's ideas about boundaries between proof/auxiliary execution contexts

Sketch: Soundness of assertMul()

• There are two high-level ingredients to the soundness proof for ffmul:
  1. limb-wise equations mod n must actually hold over the integers. this is ensured with range-checking all participating terms
  2. in the final equation, we must ensure bounds on the combined bigints to go from "mod 2^3l n" to "over the integers"
• Re (1): The result of an ffadd chain, even without range check, is guaranteed to be a sum of small terms mod n. For a proof, see the PR description of Foreign fields 3: Addition #1220 (recommended read to understand the reasoning here)
• To be precise, this is true for the combined lower two limbs r01, and for the high limb r2
• If we add another equation that guarantees r0 = x0 + sign*y0 - overflow*f0 - carry*2^l, we can use that in the eq. for both lower limbs and conclude the same thing for r1 separately
• If we replace "small value thanks to range-check" with "sum of small values" in the ffmul proof, the conclusions of the form |LHS - RHS| < n which guarantee non-overflow are still valid. We can even allow "small negative" values (close to the native modulus) which might be caused by an adversarial carry in ffadd, nothing breaks.
• The limbwise equations that are being proved over the integers now prove ff addition and multiplication in one step, instead of only multiplication. this is fine
• Re (2), additional bounds checks: We can understand bounds on the inputs of a sum as a weaker version of bounds on the result of a sum, e.g. |x - y| <= |x| + |y|. If x and y are small enough, the bound we need on |x - y| is still guaranteed. The assertMul() gadget enforces a rigorous version of this argument which takes into account the number of summands, see here. We maintain the same assumption used in other gadgets, that inputs are bounds-checked relative to f, i.e. x[2] <= f[2], but depending on the number of summands we enforce a stronger bound on f than elsewhere (other gadgets enforce f < 2^259, assertMul() enforces at least f < 2^258 because we can no longer assume that r is "positive". It doesn't matter because in reality, we're only interested in f < 2^256 and that's supported even for pretty long sums.

[querolita]

Before giving an actual code review, I will start by writing comments on the above explanation.

the ffadd gate constrains its result a + b in compact limb form [r01, r2]

This is true for FFMul, where the remainder is indeed given in compact form in the first 2 columns of the next row. But in the FFAdd gate, the result was given in 3-limb form, in the first 3 columns of the next row so it could chain easily with FFAdd and FFMul. Thus, what exactly are you referring to with the above sentence?

the naive version also suffers because it computes the left and right side separately, and so applies range-checks to the same result twice.

You mean one would naively range check x and y twice? Meaning, when computing $x^2$, once for "left" $x$ and again for "right" $x$? If this is what you meant then yes, this should be avoided and only range check it once to avoid duplicates.

[querolita]

Leaving some comments for now regarding collateral changes introduced in the PR, which are not just the FFMul optimization. Leaving the rest of it (the most remarkable part) for later (but not too much).

[mitschabaude]

This is true for FFMul, where the remainder is indeed given in compact form in the first 2 columns of the next row. But in the FFAdd gate, the result was given in 3-limb form, in the first 3 columns of the next row so it could chain easily with FFAdd and FFMul. Thus, what exactly are you referring to with the above sentence?

@querolita the equation that ffadd adds is $x_{01} + s y_{01} + \ldots$, see here:
https://github.com/o1-labs/proof-systems/blob/ad91893408543d47e893cbcd6bcd14c955a8ac0d/kimchi/src/circuits/polynomials/foreign_field_add/circuitgates.rs#L184-L191

[mitschabaude]

You mean one would naively range check x and y twice? Meaning, when computing , once for "left" and again for "right" ?

I mean that you'd check x^2 - y^2 when calling sub(), and you'd also range-check (x - y)(x + y) when calling mul(), and then you assert they are equal. But if you already know they should be equal you could avoid one of the range checks. It's a minor point, of course we could add a flag to sub() or mul() that tells it to skip the range check. The nice thing about this API here is that it doesn't need any such flags

[mitschabaude]

This is the full constraint system for the assertMul() example I use in the test. It does (x - y)(x + y) === x2 + y2, where x2 and y2 are just witnessed (not constrained)

0    Generic         [1 0 0 1 0] [-1 0 -1 1 0]      0->2, 1->3, 2->5, 3->4, 4->(1,3), 5->0
1    Generic         [1 1 -1 0 0] [2894.. 2894.. -1 0 0] 0->(2,3), 1->5, 2->(6,4), 3->(0,1), 4->(2,6), 5->1
2    ForeignFieldAdd [3094.. 3094.. 1208.. 1]       0->(6,3), 1->(9,1), 2->(9,2), 3->(8,3), 4->(9,4), 5->(9,5), 6->(1,4)
3    Zero                                           0->(6,5), 1->(10,4), 2->(10,5)
4    ForeignFieldAdd [3094.. 3094.. 1208.. -1]      
5    Zero                                           0->(20,0), 1->(20,1), 2->(11,1)
6    Generic         [-1 0 -1 1 0] [1 1 -1 0 0]     0->1, 1->(7,0), 2->(7,3), 3->(8,0), 4->(1,2), 5->(10,3)
7    Generic         [2894.. 2894.. -1 0 0] [1 0 0 1 0] 0->4, 1->(9,6), 2->(8,4), 3->5, 4->(6,0), 5->(6,2)
8    Generic         [1 1 -1 0 0] [-1 1 -1 0 0]     0->(9,0), 1->5, 2->(10,0), 3->(9,3), 4->(7,2), 5->1
9    ForeignFieldAdd [3094.. 3094.. 1208.. -1]      0->(2,0), 1->(2,1), 2->(2,2), 3->(1,0), 4->(2,4), 5->(2,5), 6->(7,1)
10   ForeignFieldMul [1208.. 4294.. 0 3082..]       0->(8,2), 3->(3,0), 4->(3,1), 5->(3,2), 6->(12,0)
11   Zero                                           0->(20,2), 1->(5,2), 2->(16,0), 3->(17,0), 4->(18,0), 5->(14,0), 6->(13,0)
12   RangeCheck0     [0]                            0->(10,6), 1->(15,3), 2->(15,4)
13   RangeCheck0     [0]                            0->(11,6), 1->(15,5), 2->(15,6)
14   RangeCheck1                                    0->(11,5), 1->(18,1)
15   Zero                                           3->(12,1), 4->(12,2), 5->(13,1), 6->(13,2)
16   RangeCheck0     [0]                            0->(11,2), 1->(19,3), 2->(19,4)
17   RangeCheck0     [0]                            0->(11,3), 1->(19,5), 2->(19,6)
18   RangeCheck1                                    0->(11,4), 1->(20,3)
19   Zero                                           3->(16,1), 4->(16,2), 5->(17,1), 6->(17,2)
20   Generic         [1 3094.. -1 0 0] [1 0 0 0 0]  0->(5,0), 1->(5,1), 2->(11,0), 3->(14,1)

[Trivo25]

first pass, looks great so far! Will continue later

[Trivo25]

I like the use of assert in our code base, it makes the code flow much better and easier to read than constantly interrupting the flow with a bunch of verbose if (!/** **/) throw Error("..") statements

[jspada]

Spotted one possible thing (below); other than that looks amazing and makes sense.