lib/accy/src/properties/gradient.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const hypothesis = @import("hypothesis");
  3 const accy = @import("accy");
  4 
  5 const tensor = accy.tensor;
  6 
  7 const tokens = 4;
  8 const channels = 3;
  9 const element_count = tokens * channels;
 10 const mask_penalty: f32 = -30.0;
 11 const probe_count = 4;
 12 const probe_step: f32 = 0.01;
 13 
 14 pub fn settings() hypothesis.Settings {
 15     return hypothesis.Settings.quick()
 16         .withSeed(0xacc7_c0de_09ad)
 17         .withDatabase("zig-out/hypothesis-failures/accy");
 18 }
 19 
 20 fn causalAttentionLoss(_: *tensor.Builder, args: []const tensor.Value) !tensor.Value {
 21     const queries = args[0];
 22     const keys = args[1];
 23     const values = args[2];
 24     const mask = args[3];
 25 
 26     const scores = try queries.contract(try keys.rename(.token, .key), .channel);
 27     const masked = try scores.add(mask);
 28     const centered = try masked.sub(try masked.max(.key));
 29     const weights = try centered.exp();
 30     const probs = try weights.div(try weights.sum(.key));
 31     const mixed = try probs.contract(try values.rename(.token, .key), .key);
 32     const squared = try mixed.mul(mixed);
 33     return try squared.sum(.{ .token, .channel });
 34 }
 35 
 36 fn defineAttentionLoss(allocator: std.mem.Allocator) !tensor.Graph {
 37     return tensor.define(allocator, "property_causal_attention", &.{
 38         tensor.spec(.f32, .{ .token = tokens, .channel = channels }),
 39         tensor.spec(.f32, .{ .token = tokens, .channel = channels }),
 40         tensor.spec(.f32, .{ .token = tokens, .channel = channels }),
 41         tensor.spec(.f32, .{ .token = tokens, .key = tokens }),
 42     }, causalAttentionLoss);
 43 }
 44 
 45 fn launchLoss(
 46     executor: *tensor.execute.Cpu,
 47     allocator: std.mem.Allocator,
 48     inputs: *const [3][element_count]f32,
 49     mask: []const f32,
 50 ) !f32 {
 51     var loss: f32 = 0.0;
 52     var outputs = [_][]u8{std.mem.asBytes(&loss)};
 53     try executor.launch(allocator, &.{
 54         std.mem.sliceAsBytes(inputs[0][0..]),
 55         std.mem.sliceAsBytes(inputs[1][0..]),
 56         std.mem.sliceAsBytes(inputs[2][0..]),
 57         std.mem.sliceAsBytes(mask),
 58     }, outputs[0..]);
 59     return loss;
 60 }
 61 
 62 pub const AttentionGradientProperty = struct {
 63     pub fn property(data: *hypothesis.ConjectureData, _: std.mem.Allocator) !void {
 64         const allocator = std.testing.allocator;
 65 
 66         var inputs: [3][element_count]f32 = undefined;
 67         for (&inputs) |*tensor_input| {
 68             for (tensor_input) |*slot| {
 69                 const raw = try data.drawInteger(0, 2000, 1000);
 70                 slot.* = (@as(f32, @floatFromInt(raw)) - 1000.0) / 1000.0;
 71             }
 72         }
 73         var mask: [tokens * tokens]f32 = undefined;
 74         for (0..tokens) |row| {
 75             for (0..tokens) |col| {
 76                 mask[row * tokens + col] = if (col <= row) 0.0 else mask_penalty;
 77             }
 78         }
 79 
 80         var source = try defineAttentionLoss(allocator);
 81         defer source.deinit();
 82         var combined = try tensor.valueAndGrad(allocator, &source, .{ .wrt = &.{ 0, 1, 2 } });
 83         defer combined.deinit();
 84 
 85         var loss_value: f32 = 0.0;
 86         var gradients: [3][element_count]f32 = @splat(@splat(0.0));
 87         var grad_outputs = [_][]u8{
 88             std.mem.asBytes(&loss_value),
 89             std.mem.sliceAsBytes(gradients[0][0..]),
 90             std.mem.sliceAsBytes(gradients[1][0..]),
 91             std.mem.sliceAsBytes(gradients[2][0..]),
 92         };
 93         try tensor.execute.runCpu(allocator, &combined, &.{
 94             std.mem.sliceAsBytes(inputs[0][0..]),
 95             std.mem.sliceAsBytes(inputs[1][0..]),
 96             std.mem.sliceAsBytes(inputs[2][0..]),
 97             std.mem.sliceAsBytes(mask[0..]),
 98         }, grad_outputs[0..]);
 99 
100         var loss_executor = try tensor.execute.Cpu.init(allocator, &source);
101         defer loss_executor.deinit();
102 
103         for (0..probe_count) |_| {
104             const tensor_index: usize = @intCast(try data.drawInteger(0, 2, 0));
105             const element_index: usize = @intCast(try data.drawInteger(0, element_count - 1, 0));
106             var probes: [2]f32 = undefined;
107             for (&probes, [_]f32{ probe_step, -probe_step }) |*probe, offset| {
108                 var perturbed = inputs;
109                 perturbed[tensor_index][element_index] += offset;
110                 probe.* = try launchLoss(&loss_executor, allocator, &perturbed, mask[0..]);
111             }
112             const finite_difference = (probes[0] - probes[1]) / (2.0 * probe_step);
113             const analytic = gradients[tensor_index][element_index];
114             if (@abs(finite_difference - analytic) > 0.002 + 0.02 * @abs(finite_difference)) {
115                 return error.PropertyFailed;
116             }
117         }
118     }
119 };
120 
121 test "property: causal attention gradients match finite differences" {
122     try hypothesis.checkNamed(AttentionGradientProperty, "accy-causal-attention-gradient", settings());
123 }