理解 Transformer 如何利用上下文

import java.util.*;
public class Main {
    static double[] attend(double[] q,double[][] k,double[][] v,int visible){
        if(q.length==0||k.length!=v.length||visible<1||visible>k.length)throw new IllegalArgumentException("shape");
        int width=v[0].length;if(width==0)throw new IllegalArgumentException("empty value");
        double[] s=new double[visible];double max=Double.NEGATIVE_INFINITY;
        for(double x:q)if(!Double.isFinite(x))throw new IllegalArgumentException("query");
        for(int i=0;i<visible;i++){
            if(k[i].length!=q.length||v[i].length!=width)throw new IllegalArgumentException("dimension");
            for(int j=0;j<q.length;j++){
                if(!Double.isFinite(k[i][j]))throw new IllegalArgumentException("key");
                s[i]+=q[j]*k[i][j];
            }
            s[i]/=Math.sqrt(q.length);
            if(!Double.isFinite(s[i]))throw new IllegalArgumentException("score overflow");
            max=Math.max(max,s[i]);
        }
        double sum=0;
        for(int i=0;i<visible;i++){s[i]=Math.exp(s[i]-max);sum+=s[i];}// 稳定 Softmax。
        double[] out=new double[width];
        for(int i=0;i<visible;i++)for(int j=0;j<width;j++){
            if(!Double.isFinite(v[i][j]))throw new IllegalArgumentException("value");
            out[j]+=(s[i]/sum)*v[i][j]; // 先归一权重，再混合值向量。
        }
        for(double x:out)if(!Double.isFinite(x))throw new IllegalArgumentException("output overflow");
        return out;
    }
    static void near(double expected,double got){if(!Double.isFinite(got)||Math.abs(expected-got)>1e-10)throw new AssertionError(got);}
    public static void main(String[] args){
        double[][] values={{2},{6}};
        near(4,attend(new double[]{1},new double[][]{{0},{0}},values,2)[0]);
        near(5,attend(new double[]{1},new double[][]{{0},{Math.log(3)}},values,2)[0]);
        near(2,attend(new double[]{1},new double[][]{{0},{1000}},values,1)[0]);// 未来项不进入分母。
        near(4,attend(new double[]{1},new double[][]{{1000},{1000}},values,2)[0]);
        try{attend(new double[]{1},new double[][]{{0}},new double[][]{{2}},0);throw new AssertionError();}catch(IllegalArgumentException expected){}
        System.out.println("PASS 5 tests");
    }
}