Attention Mechanism Theory

1
#include <iostream>
2
#include <vector>
3
#include <cmath>
4
#include <iomanip>
5
#include <limits>
6
#include <algorithm>
7
8
using Matrix = std::vector<std::vector<double>>;
9
10
Matrix transpose(const Matrix &A) {
11
    size_t n = A.size(), m = A[0].size();
12
    Matrix T(m, std::vector<double>(n, 0.0));
13
    for (size_t i = 0; i < n; ++i)
14
        for (size_t j = 0; j < m; ++j)
15
            T[j][i] = A[i][j];
16
    return T;
17
}
18
19
Matrix matmul(const Matrix &A, const Matrix &B) {
20
    size_t n = A.size(), m = A[0].size(), p = B[0].size();
21
    Matrix C(n, std::vector<double>(p, 0.0));
22
    for (size_t i = 0; i < n; ++i)
23
        for (size_t k = 0; k < m; ++k)
24
            for (size_t j = 0; j < p; ++j)
25
                C[i][j] += A[i][k] * B[k][j];
26
    return C;
27
}
28
29
// Row-wise softmax with numerical stability
30
void row_softmax(Matrix &A) {
31
    for (auto &row : A) {
32
        double mx = -std::numeric_limits<double>::infinity();
33
        for (double x : row) mx = std::max(mx, x);
34
        double sum = 0.0;
35
        for (double &x : row) { x = std::exp(x - mx); sum += x; }
36
        for (double &x : row) x /= (sum > 0 ? sum : 1.0);
37
    }
38
}
39
40
// Apply additive mask before softmax: mask[i][j] == true means allowed (no penalty), false means disallowed (add -inf)
41
void apply_mask(Matrix &scores, const std::vector<std::vector<bool>> *mask) {
42
    if (!mask) return;
43
    double neg_inf = -1e30; // practical -inf
44
    for (size_t i = 0; i < scores.size(); ++i)
45
        for (size_t j = 0; j < scores[0].size(); ++j)
46
            if (!(*mask)[i][j]) scores[i][j] = neg_inf;
47
}
48
49
Matrix scaled_dot_product_attention(const Matrix &Q, const Matrix &K, const Matrix &V,
50
                                    const std::vector<std::vector<bool>> *mask = nullptr) {
51
    size_t d_k = Q[0].size();
52
    // Compute S = Q K^T / sqrt(d_k)
53
    Matrix KT = transpose(K);
54
    Matrix S = matmul(Q, KT);
55
    double scale = 1.0 / std::sqrt(static_cast<double>(d_k));
56
    for (auto &row : S) for (double &x : row) x *= scale;
57
58
    // Apply mask if provided, then softmax row-wise
59
    apply_mask(S, mask);
60
    row_softmax(S); // S becomes attention weights A
61
62
    // Output O = A V
63
    Matrix O = matmul(S, V);
64
    return O;
65
}
66
67
void print_matrix(const Matrix &A, const std::string &name) {
68
    std::cout << name << " (" << A.size() << "x" << A[0].size() << ")\n";
69
    std::cout << std::fixed << std::setprecision(4);
70
    for (const auto &row : A) {
71
        for (double x : row) std::cout << std::setw(8) << x << ' ';
72
        std::cout << '\n';
73
    }
74
    std::cout << '\n';
75
}
76
77
int main() {
78
    // Example: self-attention on a tiny sequence (n=3, d_k=d_v=4)
79
    Matrix X = {
80
        {0.1, 0.2, 0.3, 0.4},
81
        {0.5, 0.4, 0.3, 0.2},
82
        {0.9, 0.7, 0.1, 0.0}
83
    };
84
85
    // Here we use identity projections: Q=K=V=X for simplicity
86
    Matrix Q = X, K = X, V = X;
87
88
    // Optional: build a causal mask (lower triangular, including diagonal)
89
    std::vector<std::vector<bool>> causal_mask(3, std::vector<bool>(3, false));
90
    for (size_t i = 0; i < 3; ++i)
91
        for (size_t j = 0; j <= i; ++j)
92
            causal_mask[i][j] = true;
93
94
    Matrix O_no_mask = scaled_dot_product_attention(Q, K, V, nullptr);
95
    Matrix O_causal  = scaled_dot_product_attention(Q, K, V, &causal_mask);
96
97
    print_matrix(Q, "Q");
98
    print_matrix(K, "K");
99
    print_matrix(V, "V");
100
    print_matrix(O_no_mask, "Output without mask");
101
    print_matrix(O_causal,  "Output with causal mask");
102
103
    return 0;
104
}
105

1
#include <iostream>
2
#include <vector>
3
#include <random>
4
#include <cmath>
5
#include <iomanip>
6
#include <limits>
7
#include <algorithm>
8
9
using Matrix = std::vector<std::vector<double>>;
10
11
Matrix transpose(const Matrix &A){ size_t n=A.size(), m=A[0].size(); Matrix T(m,std::vector<double>(n)); for(size_t i=0;i<n;++i) for(size_t j=0;j<m;++j) T[j][i]=A[i][j]; return T; }
12
Matrix matmul(const Matrix &A,const Matrix &B){ size_t n=A.size(), m=A[0].size(), p=B[0].size(); Matrix C(n,std::vector<double>(p,0.0)); for(size_t i=0;i<n;++i) for(size_t k=0;k<m;++k) for(size_t j=0;j<p;++j) C[i][j]+=A[i][k]*B[k][j]; return C; }
13
void row_softmax(Matrix &A){ for(auto &row:A){ double mx=-1e30; for(double x:row) mx=std::max(mx,x); double s=0; for(double &x:row){ x=std::exp(x-mx); s+=x;} for(double &x:row) x/= (s>0? s:1.0);} }
14
15
Matrix attention(const Matrix &Q,const Matrix &K,const Matrix &V){ size_t d_k=Q[0].size(); Matrix KT=transpose(K); Matrix S=matmul(Q,KT); double scale=1.0/std::sqrt((double)d_k); for(auto &r:S) for(double &x:r) x*=scale; row_softmax(S); return matmul(S,V); }
16
17
Matrix linear(const Matrix &X,const Matrix &W){ return matmul(X,W); }
18
19
void print_matrix(const Matrix &A,const std::string &name){ std::cout<<name<<" ("<<A.size()<<"x"<<A[0].size()<<")\n"; std::cout.setf(std::ios::fixed); std::cout<<std::setprecision(4); for(const auto &r:A){ for(double x:r) std::cout<<std::setw(8)<<x<<' '; std::cout<<"\n";} std::cout<<"\n"; }
20
21
int main(){
22
    // Model sizes
23
    const size_t n=3;          // sequence length
24
    const size_t d_model=8;    // model width
25
    const size_t h=2;          // number of heads
26
    const size_t d_k=d_model/h;// per-head key/query dim
27
    const size_t d_v=d_k;      // per-head value dim
28
29
    // Input sequence X (n x d_model)
30
    Matrix X = {
31
        {0.2, -0.1, 0.0, 0.5, 0.3, -0.2, 0.1, 0.4},
32
        {0.0,  0.3, 0.2, 0.1, 0.1,  0.0, 0.2, 0.0},
33
        {0.5,  0.2, 0.1, 0.0, 0.4,  0.1, 0.0, 0.3}
34
    };
35
36
    // Initialize projections: W_Q, W_K, W_V (d_model x (h*d_k)), and W_O ((h*d_v) x d_model)
37
    std::mt19937 rng(42);
38
    std::uniform_real_distribution<double> dist(-0.5, 0.5);
39
40
    auto rand_mat = [&](size_t r,size_t c){ Matrix M(r,std::vector<double>(c)); for(size_t i=0;i<r;++i) for(size_t j=0;j<c;++j) M[i][j]=dist(rng)/std::sqrt((double)r); return M; };
41
42
    Matrix WQ = rand_mat(d_model, h*d_k);
43
    Matrix WK = rand_mat(d_model, h*d_k);
44
    Matrix WV = rand_mat(d_model, h*d_v);
45
    Matrix WO = rand_mat(h*d_v, d_model);
46
47
    // Linear projections
48
    Matrix Q = linear(X, WQ); // (n x h*d_k)
49
    Matrix K = linear(X, WK); // (n x h*d_k)
50
    Matrix V = linear(X, WV); // (n x h*d_v)
51
52
    // Split into heads and compute attention per head
53
    std::vector<Matrix> heads;
54
    for(size_t head=0; head<h; ++head){
55
        // Slice columns for this head
56
        Matrix Qh(n, std::vector<double>(d_k)), Kh(n, std::vector<double>(d_k)), Vh(n, std::vector<double>(d_v));
57
        for(size_t i=0;i<n;++i){
58
            for(size_t j=0;j<d_k;++j){ Qh[i][j]=Q[i][head*d_k + j]; Kh[i][j]=K[i][head*d_k + j]; }
59
            for(size_t j=0;j<d_v;++j){ Vh[i][j]=V[i][head*d_v + j]; }
60
        }
61
        // Self-attention for this head (Q=K=V from same X)
62
        Matrix Oh = attention(Qh, Kh, Vh); // (n x d_v)
63
        heads.push_back(Oh);
64
    }
65
66
    // Concatenate head outputs along feature dimension: (n x (h*d_v))
67
    Matrix Hcat(n, std::vector<double>(h*d_v));
68
    for(size_t i=0;i<n;++i){
69
        for(size_t head=0; head<h; ++head){
70
            for(size_t j=0;j<d_v;++j){ Hcat[i][head*d_v + j] = heads[head][i][j]; }
71
        }
72
    }
73
74
    // Final output projection: (n x d_model)
75
    Matrix Y = matmul(Hcat, WO);
76
77
    print_matrix(Y, "Multi-Head Self-Attention output");
78
    return 0;
79
}
80

1
#include <iostream>
2
#include <vector>
3
#include <cmath>
4
#include <iomanip>
5
#include <limits>
6
#include <algorithm>
7
8
using Matrix = std::vector<std::vector<double>>;
9
10
Matrix transpose(const Matrix &A){ size_t n=A.size(), m=A[0].size(); Matrix T(m,std::vector<double>(n)); for(size_t i=0;i<n;++i) for(size_t j=0;j<m;++j) T[j][i]=A[i][j]; return T; }
11
Matrix matmul(const Matrix &A,const Matrix &B){ size_t n=A.size(), m=A[0].size(), p=B[0].size(); Matrix C(n,std::vector<double>(p,0.0)); for(size_t i=0;i<n;++i) for(size_t k=0;k<m;++k) for(size_t j=0;j<p;++j) C[i][j]+=A[i][k]*B[k][j]; return C; }
12
13
void row_softmax(Matrix &A){ for(auto &row:A){ double mx=-1e30; for(double x:row) mx=std::max(mx,x); double s=0; for(double &x:row){ x=std::exp(x-mx); s+=x;} for(double &x:row) x/= (s>0? s:1.0);} }
14
15
Matrix attention_weights(const Matrix &Q,const Matrix &K, const std::vector<std::vector<bool>> &mask){
16
    size_t d_k=Q[0].size();
17
    Matrix S = matmul(Q, transpose(K));
18
    double scale = 1.0/std::sqrt((double)d_k);
19
    for(size_t i=0;i<S.size();++i){
20
        for(size_t j=0;j<S[0].size();++j){
21
            double v = S[i][j]*scale;
22
            if(!mask[i][j]) v = -1e30; // apply -inf to masked positions
23
            S[i][j]=v;
24
        }
25
    }
26
    row_softmax(S);
27
    return S; // return normalized attention weights
28
}
29
30
void print_matrix(const Matrix &A,const std::string &name){ std::cout<<name<<" ("<<A.size()<<"x"<<A[0].size()<<")\n"; std::cout.setf(std::ios::fixed); std::cout<<std::setprecision(4); for(const auto &r:A){ for(double x:r) std::cout<<std::setw(8)<<x<<' '; std::cout<<"\n";} std::cout<<"\n"; }
31
32
int main(){
33
    // Toy sequence embeddings (n=5, d=4)
34
    Matrix X = {
35
        {0.1, 0.0, 0.2, 0.1},
36
        {0.0, 0.3, 0.1, 0.0},
37
        {0.2, 0.1, 0.0, 0.4},
38
        {0.4, 0.0, 0.2, 0.1},
39
        {0.3, 0.2, 0.1, 0.0}
40
    };
41
42
    Matrix Q = X, K = X; // self-attention
43
44
    // Build strict causal mask (allow j <= i)
45
    size_t n = X.size();
46
    std::vector<std::vector<bool>> mask(n, std::vector<bool>(n,false));
47
    for(size_t i=0;i<n;++i) for(size_t j=0;j<=i;++j) mask[i][j]=true;
48
49
    Matrix A = attention_weights(Q, K, mask);
50
    print_matrix(A, "Causal attention weights");
51
52
    // Check that probabilities above the diagonal are ~0
53
    for(size_t i=0;i<n;++i){
54
        for(size_t j=i+1;j<n;++j){ if(A[i][j] > 1e-8) { std::cerr << "Mask failure at ("<<i<<","<<j<<")\n"; } }
55
    }
56
57
    return 0;
58
}
59