缩放点积注意力与掩码：亲手计算一次关注

public class Main {
  static double attention(double[] q,double[][] keys,double[] values,boolean[] allowed){
    if(q.length==0||keys.length==0||keys.length!=values.length||values.length!=allowed.length)
      throw new IllegalArgumentException();
    for(double x:q)if(!Double.isFinite(x))throw new IllegalArgumentException();
    double[] scores=new double[keys.length];
    double maximum=Double.NEGATIVE_INFINITY;
    for(int i=0;i<keys.length;i++){
      if(keys[i].length!=q.length||!Double.isFinite(values[i]))throw new IllegalArgumentException();
      for(int d=0;d<q.length;d++){
        if(!Double.isFinite(keys[i][d]))throw new IllegalArgumentException();
        scores[i]+=q[d]*keys[i][d];
      }
      scores[i]/=Math.sqrt(q.length); // 缩放依赖查询键维度
      if(!Double.isFinite(scores[i]))throw new ArithmeticException("score range");
      if(allowed[i])maximum=Math.max(maximum,scores[i]);
    }
    if(maximum==Double.NEGATIVE_INFINITY)return 0; // 显式定义全屏蔽行
    double sum=0;
    for(int i=0;i<scores.length;i++){
      scores[i]=allowed[i]?Math.exp(scores[i]-maximum):0;sum+=scores[i];
    }
    double output=0;
    for(int i=0;i<scores.length;i++)output+=(scores[i]/sum)*values[i]; // 先归一化再加权
    return output;
  }
  static void check(boolean v){if(!v)throw new AssertionError();}
  static void close(double x,double y){check(Math.abs(x-y)<1e-10);}
  public static void main(String[] args){
    double[][] k={{1,0},{0,1}};double[] v={10,20};
    close(attention(new double[]{0,0},k,v,new boolean[]{true,true}),15);
    close(attention(new double[]{1,0},k,v,new boolean[]{true,false}),10);
    close(attention(new double[]{1,0},k,v,new boolean[]{false,false}),0);
    close(attention(new double[]{1000,0},k,v,new boolean[]{true,true}),10);
    close(attention(new double[]{1,0},k,new double[]{10,999},new boolean[]{true,false}),10);
    boolean bad=false;
    try{attention(new double[0],k,v,new boolean[]{true,true});}catch(IllegalArgumentException e){bad=true;}
    check(bad);
    System.out.println("6 stable attention checks passed");
  }
}