Using `winterfell` for neural networks

[timk11]

Hi team. I'm working on a project using winterfell in conjunction with a neural network trainer. My project repo is at https://github.com/timk11/zk-trainer and the main logic is in prover.rs. The aim is to train a small (publicly visible) neural network on a private dataset, and to use a zk-STARK to show that the training was authentic. Public inputs are hashes of the initial model and final model and a hash of the dataset. My questions are mainly about the trace table and constraint degree calculations.

Trace table - The trace table columns consist of flag, len_sample, hidden_size, learning_rate, dataset, expected, dataset_hash, wb_hash, and weights and biases. These are all one column each except for dataset (one sample per row, occupying len_sample columns, repeating for each epoch) and weights and biases (updated at each row except when flag == 1, variable number of columns). The update loop for each epoch is in lines 110-196.

Q1: Have I used the right approach here, including having the dataset in the trace table as a recurring matrix occupying several columns?

Constraint degrees - I've included two of these:

        let degrees = vec![
            TransitionConstraintDegree::new(5),
            TransitionConstraintDegree::new(30)
        ];

The first one is for the dataset hash. By my understanding rescue prime has degrees of 5, based on this answer.

The second one is for the hash of weights and biases. I've shown my working of the calculation in the image below. If I'm correct, the highest degree for weights should be 25, and then we add 5 for the rescue hash to make 30. Q2: Does this look right?

I've added some debug print lines to produce the following output on testing:

---- tests::test_train_and_prove stdout ----
num_samples: 7, len_sample: 8, hidden_size: 12, num_columns: 136
flag: 0, frame.next()[4]: 1988944142759683713, dataset_hash: 1988944142759683713, frame.next()[5]: 13443926300744107950, w_b_hash: 13443926300744107950, frame.next()[-1]: 7666499316168129282, bias_o: 7666499316168129282, result: 0, 0
flag: 0, frame.next()[4]: 15758796173594749889, dataset_hash: 15758796173594749889, frame.next()[5]: 13251977776104907959, w_b_hash: 13251977776104907959, frame.next()[-1]: 3585489129564534949, bias_o: 3585489129564534949, result: 0, 0

This shows frame values nicely matching up with hashes as anticipated by these lines:

        result[0] = (frame.next()[4] - dataset_hash);
        result[1] = (frame.next()[5] - w_b_hash);

I've also shown output bias and the corresponding frame value as an additional check. This all repeats quite happily for 64 lines, at which point it ought to stop, but then this starts happening:

flag: 0, frame.next()[4]: 9080182147223831867, dataset_hash: 9080182147223831867, frame.next()[5]: 8973157242228625980, w_b_hash: 8973157242228625980, frame.next()[-1]: 16312796124388405370, bias_o: 16312796124388405370, result: 0, 0
flag: 18210199780982572190, frame.next()[4]: 1535783160212123500, dataset_hash: 17866269203665947714, frame.next()[5]: 11252226151124759407, w_b_hash: 5171612489511717358, frame.next()[-1]: 17544902386472247398, bias_o: 17150592172401078110, result: 2116258025960760107, 6080613661613042049
flag: 3727822874259453825, frame.next()[4]: 1709029980266899968, dataset_hash: 7099683142479157042, frame.next()[5]: 10269402497121169506, w_b_hash: 5995184084187237790, frame.next()[-1]: 11594520110464754199, bias_o: 4733675032304182569, result: 13056090907202327247, 4274218412933931716

From this point on none of the hashes match the frame values and flag is all over the place instead of just 0 or 1. This continues on for several hundred lines until ending in:

thread 'tests::test_train_and_prove' panicked at /home/timk/.cargo/registry/src/index.crates.io-6f17d22bba15001f/winter-prover-0.12.0/src/constraints/evaluation_table.rs:214:9:
assertion `left == right` failed: transition constraint degrees didn't match
expected: [252, 1827]
actual:   [2041, 2047]
  left: [252, 1827]
 right: [2041, 2047]

Q3: What could be going on here? Any pointers would be very gratefully received.

[timk11]

My mistake - Rescue prime using Rp64_256 has a degree of 7, as specified here. I've changed the degrees accordingly:

        let degrees = vec![
            TransitionConstraintDegree::new(7),
            TransitionConstraintDegree::new(32)
        ];

However, I'm still getting the same type of issue as before:

flag: 1993318452562882023, frame.next()[4]: 2247888647023326769, dataset_hash: 15367546186385776563, frame.next()[5]: 5787859169783420782, w_b_hash: 10876202903457016335, frame.next()[-1]: 6894734504394840185, bias_o: 8709264048824405880, result: 5327086530052134527, 13358400335740988768
thread 'tests::test_train_and_prove' panicked at /home/timk/.cargo/registry/src/index.crates.io-6f17d22bba15001f/winter-prover-0.12.0/src/constraints/evaluation_table.rs:214:9:
assertion `left == right` failed: transition constraint degrees didn't match
expected: [378, 1953]
actual:   [2041, 2047]
  left: [378, 1953]
 right: [2041, 2047]

On thinking through this further, I've realised that there's another issue with the way I've set this up. The neural network training involves matrix multiplication of decimals, but because prime fields can only work with integers I've attempted to scale all the decimals up by a factor of 10000 as in the following example:

pub const E4: BaseElement = BaseElement::new(10_000);
const E8: BaseElement = BaseElement::new(100_000_000);

hidden_gradients[j] = hidden_errors[j] * hidden_activations[j] * (E4 - hidden_activations[j]) / E8;

As we well know (and I should have remembered!), division does not work as expected in a prime field and so the results obtained from this set of calculations are effectively meaningless.

An option might be to remove all the scaling until the very end. This might necessitate generating the proof for a single update of one row only (one student and one epoch in the example I've used) because rescaling the weights and biases would have to be done outside of the prime field. Even if the sigmoid approximation is removed this would still leave us with around a 13-degree calculation, so some of the weights and biases would be in the order of 10^56 prior to rescaling, which is way beyond what the f64 prime field can handle for this purpose. Instead of a scaling factor of 10000, it would need to be brought down to about 20 (assuming that all the calculated values fit nicely into the range -1 to +1). This is unlikely to produce a helpful model, especially without a proper activation function.

Q4: Have I missed anything here? Is there some other way around this apart from switching to another package (I know GKR is an option) or using a much simpler algorithm?

[timk11]

In view of the scaling issues outlined above I've now swapped out the neural network model for a linear regression model. My evaluate_transition function now looks like this:

    fn evaluate_transition<E: FieldElement<BaseField = BaseElement> + From<Self::BaseField>>(
        &self,
        frame: &EvaluationFrame<E>,
        _periodic_values: &[E],
        result: &mut [E], // as vector?
    ) {
        let current_state = frame.current();
        let flag = current_state[0].base_element(0).as_int();
        let len_sample = usize::from(WrappedBaseElement::wrap(current_state[1].base_element(0)));
        let learning_rate = current_state[2];
        let mut dataset_hash = current_state[3]; // Current dataset hash
        let mut w_b_hash = current_state[4]; // Hash of weights and biases
        let input = current_state.get(5..(5usize + len_sample)); // Input features
        let expected = current_state[5usize + len_sample]; // Expected output
        let mut weights = current_state.get((6usize + len_sample)..(6usize + len_sample * 2usize))
            .map(|slice| slice.iter().map(|&e| e.base_element(0)).collect::<Vec<_>>())
            .unwrap_or_default(); // Weights
        let mut bias = current_state[6usize + len_sample * 2usize]; // Bias

        let mut output = bias;

        dataset_hash = (update_hash(
            dataset_hash.base_element(0),
            &[
                input.map(|slice| slice.iter().map(|&e| e.base_element(0)).collect::<Vec<_>>())
                    .unwrap_or_default(),
                vec![expected.base_element(0)],
            ].concat(),
        ) * BaseElement::new(flag)).into();

        for i in 0..len_sample {
            output += input.unwrap()[i] * weights[i].into() / E4.into();
        }

        let error = output - expected;
        let gradient = error.base_element(0);
        for i in 0..len_sample {
            weights[i] -= (learning_rate * gradient.into() * input.unwrap()[i] * BaseElement::new(flag).into() / E8.into()).base_element(0);
        }
        bias -= learning_rate * gradient.into() * BaseElement::new(flag).into() / E4.into();
        
        w_b_hash = row_hash(&[
            weights.clone(),
            vec![bias.base_element(0)]
        ].concat()).into();

        result[0] = frame.next()[3] - dataset_hash;
        result[1] = frame.next()[4] - w_b_hash;
    }

Constraint degrees are like so:

        let degrees = vec![
            TransitionConstraintDegree::new(7),
            TransitionConstraintDegree::new(12)
        ];

However, I'm still having the same issue as before, with a similar output and resulting in this error:

thread 'tests::test_train_and_prove' panicked at /home/timk/.cargo/registry/src/index.crates.io-6f17d22bba15001f/winter-prover-0.12.0/src/constraints/evaluation_table.rs:214:9:
assertion `left == right` failed: transition constraint degrees didn't match
expected: [378, 693]
actual:   [1017, 1023]
  left: [378, 693]
 right: [1017, 1023]

Is there something fundamentally wrong with my transition constraint logic?