Source browser

examples/transformer_reference.hpp

Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.

1/** @file transformer_reference.hpp
2 * @brief Double-precision, value-only oracle for the transformer example.
3 * This host implementation uses ordinary arrays and matrix products in
4 * double precision. Its finite differences validate the training derivatives.
5 * Optimization uses the analytic gradients produced by the native shader.
6 */
7#ifndef SHDR_EXAMPLES_TRANSFORMER_REFERENCE_HPP // Include this example helper only once.
8#define SHDR_EXAMPLES_TRANSFORMER_REFERENCE_HPP // Remember that this example helper is available.
9#include <algorithm> // Use standard copying, limits and zero filling.
10#include <array> // Use fixed-size host coefficient and prediction arrays.
11#include <cmath> // Use ordinary scalar math in host designers and oracles.
12
13namespace shdr_examples { // Keep these example helpers in their own namespace.
14namespace detail { // Keep these example helpers in their own namespace.
15
16inline void // Declare the reusable host calculation or check.
17transformer_projection(const double source[16], const double* weights, double result[16]) { // Apply a four-by-four projection to every token.
18    for (unsigned token = 0; token < 4; ++token) // Visit each of the four tokens.
19        for (unsigned column = 0; column < 4; ++column) { // Visit each of the four feature coordinates.
20            double value = 0.0; // Start this sample or projected-feature accumulator at zero.
21            for (unsigned input = 0; input < 4; ++input) // Visit every value needed by this calculation.
22                value += source[4 * token + input] * weights[4 * input + column]; // Accumulate one entry of the learned matrix product.
23            result[4 * token + column] = value; // Store this projected token feature.
24        } // Finish this channel, sample or feature pass.
25} // Finish the transformer_projection operation.
26
27inline std::array<double, 4> // Declare the reusable host calculation or check.
28transformer_reference(const float* input, const double* parameters, unsigned label) { // Compute the transformer independently in double precision.
29    double tokens[16], queries[16], keys[16], values[16], context[16], residual[16], features[16]; // Keep each intermediate reference tensor as ordinary double arrays.
30    for (unsigned i = 0; i < 16; ++i) // Visit every value needed by this calculation.
31        tokens[i] = input[i] + parameters[i]; // Add the learned positional value to this input feature.
32    transformer_projection(tokens, parameters + 16, queries); // Apply a four-by-four projection to every token.
33    transformer_projection(tokens, parameters + 32, keys); // Apply a four-by-four projection to every token.
34    transformer_projection(tokens, parameters + 48, values); // Apply a four-by-four projection to every token.
35    for (unsigned token = 0; token < 4; ++token) { // Visit each of the four tokens.
36        double scores[4] = {}, maximum = -1e300, total = 0.0; // Start the attention scores and softmax denominator at zero.
37        for (unsigned source = 0; source < 4; ++source) { // Combine the contributing source features or samples.
38            for (unsigned column = 0; column < 4; ++column) // Visit each of the four feature coordinates.
39                scores[source] += queries[4 * token + column] * keys[4 * source + column] / 2.0; // Accumulate the scaled query-key attention score.
40            maximum = std::max(maximum, scores[source]); // Keep the largest score for a stable softmax shift.
41        } // Finish this channel, sample or feature pass.
42        for (unsigned source = 0; source < 4; ++source) { // Combine the contributing source features or samples.
43            scores[source] = std::exp(scores[source] - maximum); // Exponentiate the shifted attention score.
44            total += scores[source]; // Accumulate the attention normalization denominator.
45        } // Finish this channel, sample or feature pass.
46        for (unsigned column = 0; column < 4; ++column) { // Visit each of the four feature coordinates.
47            context[4 * token + column] = 0.0; // Start the attended feature at zero.
48            for (unsigned source = 0; source < 4; ++source) // Combine the contributing source features or samples.
49                context[4 * token + column] += scores[source] * values[4 * source + column] / total; // Weight this value feature by its attention probability.
50        } // Finish this channel, sample or feature pass.
51    } // Finish this channel, sample or feature pass.
52    transformer_projection(context, parameters + 64, residual); // Apply a four-by-four projection to every token.
53    for (unsigned i = 0; i < 16; ++i) // Visit every value needed by this calculation.
54        residual[i] += tokens[i]; // Add the original token features as the attention residual.
55    for (unsigned token = 0; token < 4; ++token) { // Visit each of the four tokens.
56        double square = 1e-5, normalized[4], hidden[4]; // Add a small RMS stabilizer before accumulating squared features.
57        for (unsigned column = 0; column < 4; ++column) // Visit each of the four feature coordinates.
58            square += residual[4 * token + column] * residual[4 * token + column] / 4.0; // Accumulate the mean squared token magnitude.
59        for (unsigned column = 0; column < 4; ++column) // Visit each of the four feature coordinates.
60            normalized[column] = // Prepare the scaled RMS-normalized token feature.
61                residual[4 * token + column] * parameters[80 + column] / std::sqrt(square); // Divide by the RMS magnitude and apply the learned scale.
62        for (unsigned column = 0; column < 4; ++column) { // Visit each of the four feature coordinates.
63            double value = parameters[100 + column]; // Start the first feed-forward projection with its bias.
64            for (unsigned source = 0; source < 4; ++source) // Combine the contributing source features or samples.
65                value += normalized[source] * parameters[84 + 4 * source + column]; // Accumulate the first feed-forward projection.
66            hidden[column] = std::max(0.0, value); // Apply ReLU to the hidden feed-forward value.
67        } // Finish this channel, sample or feature pass.
68        for (unsigned column = 0; column < 4; ++column) { // Visit each of the four feature coordinates.
69            double value = normalized[column] + parameters[120 + column]; // Start the second projection with its normalized residual and bias.
70            for (unsigned source = 0; source < 4; ++source) // Combine the contributing source features or samples.
71                value += hidden[source] * parameters[104 + 4 * source + column]; // Accumulate the second feed-forward projection.
72            features[4 * token + column] = value; // Store the feed-forward residual feature.
73        } // Finish this channel, sample or feature pass.
74    } // Finish this channel, sample or feature pass.
75    double logits[2] = {parameters[132], parameters[133]}; // Initialize the two class logits with their learned biases.
76    for (unsigned column = 0; column < 4; ++column) { // Visit each of the four feature coordinates.
77        double mean = 0.0; // Start the pooled feature accumulator at zero.
78        for (unsigned token = 0; token < 4; ++token) // Visit each of the four tokens.
79            mean += features[4 * token + column] / 4.0; // Average this feature over the four tokens.
80        for (unsigned output = 0; output < 2; ++output) // Visit every value needed by this calculation.
81            logits[output] += mean * parameters[124 + 2 * column + output]; // Add the pooled feature contribution to this class logit.
82    } // Finish this channel, sample or feature pass.
83    const double maximum = std::max(logits[0], logits[1]); // Use the maximum logit to stabilize the exponentials.
84    const double first = std::exp(logits[0] - maximum), second = std::exp(logits[1] - maximum), // Compute the two stabilized class exponentials and their sum.
85                 total = first + second; // Sum the stable class exponentials.
86    return std::array<double, 4>{{first / total, // Return both probabilities, loss and predicted class.
87                                  second / total, // Return the second class probability.
88                                  std::log(total) + maximum - logits[label], // Return stable cross entropy for the known target class.
89                                  logits[1] > logits[0] ? 1.0 : 0.0}}; // Return the class with the larger logit.
90} // Finish the transformer_reference operation.
91} // namespace detail
92} // namespace shdr_examples
93#endif // Finish the conditional compilation block.