Attention Mechanism

Gated units solved one question. In a recurrent network, information from a distant step was kept from being erased on the way, and the dependency length that can be carried was measured. The form of the solution did not change, though: the information still moved forward step by step, was written into a single hidden state at every step, and was read from that single state at the end of the sequence. The gate regulated what would enter the state; it did not multiply the state itself.

The question of this lesson is one step back. What does looking directly at a distant step of the sequence gain — which number would move if the output could see all the steps at once and choose for itself how much weight to give to which. The second question follows the first and is the real measure of this lesson: when the chosen weights are read, can one understand why the network gave that answer.

• AR46. The sequence setup is rebuilt in this lesson and is fictional: a counter-reading sequence of twelve steps per subscriber, seed 20260218, 720 sequences, split 360/180/180. In the sequence of architectures/04 the marker was always in the first period; whether the weights fall on the correct step can be measured only when the marker position varies between sequences, so the setup is built for this measure.
• AR47. The readings are written as the deviation that remains after the subscriber’s own level and seasonal component are removed. Step 0 is the same calibration reading for every subscriber and its value is 0.0.
• AR48. Each sequence has a marker step; it lies between the second and the eleventh step. The label depends only on whether the reading of that step exceeds 0.35. The other eleven steps never enter the label.
• AR49. The input is two channels per step: the reading and the marker indicator. The network sees which step is marked; it does not see the label.
• AR50. The baseline is the most frequent class: giving every sequence the majority label. The table baseline of 0.7579 in the Introduction to Machine Learning through Supervised Learning courses is the number of a different task and is not carried here.
• AR51. The attention layer is written inside the lesson and is a model. The query is single and learned; the key and the value are produced by linear projection; the weights are normalized with softmax.
• AR52. Training is full-batch gradient descent: 200 epochs, learning rate 0.6. The gradient descent procedure was built in the linear-models topic of the Supervised Learning course and is not repeated.
• AR53. Every number is printed with two runs, and the runs differ only in the initialization seed. This is the third number of this course.
• AR54. The causality measure is step removal: a step is dropped from the attention distribution, the remaining weights are renormalized, and how much the output probability moved is printed.

The Sequence Setup and the Baseline

The task is deliberately narrow. There is a single step that determines the label, and that step is marked; the only thing the network has to do is find the marked step and compare its reading with the threshold. In such a task the question “did it look at the correct step” has an exact answer, because the correct step is written in the setup that generates the data.

# common.py — MODEL. FICTIONAL sequence set: a 12-step counter-reading sequence per subscriber.
import math

SEED, M32, STEPS = 20260218, 0xFFFFFFFF, 12
REGIONS = [("north", 0.28, 21), ("south", 0.22, 17), ("east", 0.18, 26),
           ("west", 0.14, 14), ("center", 0.18, 23)]


def generator(t):
    x = ((t ^ (t >> 16)) * 2246822507) & M32
    x = ((x ^ (x >> 13)) * 3266489909) & M32
    s = [(x ^ (x >> 16)) & M32]

    def next_value():
        s[0] = (s[0] * 1664525 + 1013904223) & M32
        return s[0] / 4294967296
    return next_value


def discrete(u, w):
    t = 0.0
    for i, x in enumerate(w):
        t += x
        if u < t:
            return i
    return len(w) - 1


def build_sequences(n, seed):
    D = []
    for i in range(n):
        r = generator(seed + i)
        b = REGIONS[discrete(r(), [x[1] for x in REGIONS])]
        reading = [0.0] + [round((r() + r() + r() - 1.5) * 1.9, 3) for _ in range(STEPS - 1)]
        m = 2 + int(r() * (STEPS - 2))
        D.append({"reading": reading, "marker": m, "y": 1 if reading[m] > 0.35 else 0, "region": b[0]})
    return D


DATA = build_sequences(720, SEED)
TRAIN, HELD, TEST = DATA[:360], DATA[360:540], DATA[540:]
rate = sum(x["y"] for x in TEST) / len(TEST)
BASELINE = max(rate, 1 - rate)
print(f"sequences {len(DATA)}, length {STEPS}, split {len(TRAIN)}/{len(HELD)}/{len(TEST)}")
print(f"positive rate train {sum(x['y'] for x in TRAIN) / len(TRAIN):.4f}, "
      f"held-out {rate:.4f}")
print(f"baseline (most frequent class) {BASELINE:.4f}")
print(f"share of sequences with a step above the threshold "
      f"{sum(any(x['reading'][t] > 0.35 for t in range(2, STEPS)) for x in DATA) / len(DATA):.4f}")
print(f"does the largest reading fall on the marker step "
      f"{sum(max(range(2, STEPS), key=lambda t: x['reading'][t]) == x['marker'] for x in TEST) / len(TEST):.4f}")

sequences 720, length 12, split 360/180/180
positive rate train 0.3806, held-out 0.3667
baseline (most frequent class) 0.6333
share of sequences with a step above the threshold 0.9903
does the largest reading fall on the marker step 0.0611

The last two lines show that the task really requires looking inside the sequence. In 0.9903 of the sequences there is at least one step above the threshold, so the question “is there a high reading anywhere” answers yes for almost every sequence and separates nothing. Looking at the largest reading does not work either: the highest step of the sequence and the marker step coincide in only 0.0611 of the held-out set. The step the network has to choose is neither the largest nor at a fixed place.

The baseline is 0.6333. Since the positive rate in the held-out set is 0.3667, giving every sequence the negative label yields this number. It is the plainest procedure that does not learn, and every number from here on is read next to it.

Query, Key and Value

Attention is defined with three projections. Each step produces a key and a value from its own input; the side that will produce the output carries a query. The inner product of the query and the key is a score that says how relevant that step is, the scores are passed through softmax and turned into weights that sum to one, and the output is the average of the values with these weights. The inner product is divided by the square root of the key dimension; dividing by $\sqrt{d}$ keeps the scores from growing as the dimension grows and saturating the softmax.

In classification the query does not have to come from a step; a single learned query is enough. The network asks “which step is important for this decision” once and takes one summary of the whole sequence. There is no hidden state, no recurrence, no order dependency between steps.

DK = DV = 3
SCALE = 1.0 / math.sqrt(DK)
INP = [([[x["reading"][t], 1.0 if t == x["marker"] else 0.0] for t in range(STEPS)],
        x["y"], x["marker"]) for x in DATA]
TR_IN, TE_IN = INP[:360], INP[540:]


def forward(P, X, skip=None):
    K = [[P["Wk"][i][0] * x[0] + P["Wk"][i][1] * x[1] for i in range(DK)] for x in X]
    V = [[P["Wv"][i][0] * x[0] + P["Wv"][i][1] * x[1] for i in range(DV)] for x in X]
    s = [sum(P["q"][i] * k[i] for i in range(DK)) * SCALE for k in K]
    d = [t for t in range(STEPS) if t != skip]
    mx = max(s[t] for t in d)
    e = [math.exp(s[t] - mx) if t in d else 0.0 for t in range(STEPS)]
    z = sum(e)
    a = [u / z for u in e]
    c = [sum(a[t] * V[t][i] for t in d) for i in range(DV)]
    return K, V, a, c, sum(P["w"][i] * c[i] for i in range(DV)) + P["b"]


def train_model(seed, epochs=200, lr=0.6):
    r = generator(seed)
    g = lambda: (r() + r() + r() - 1.5) * 0.9
    P = {"Wk": [[g(), g()] for _ in range(DK)], "Wv": [[g(), g()] for _ in range(DV)],
         "q": [g() for _ in range(DK)], "w": [g() for _ in range(DV)], "b": 0.0}
    for _ in range(epochs):
        gWk = [[0.0, 0.0] for _ in range(DK)]
        gWv = [[0.0, 0.0] for _ in range(DV)]
        gq, gw, gb = [0.0] * DK, [0.0] * DV, 0.0
        for X, y, _m in TR_IN:
            K, V, a, c, o = forward(P, X)
            d = 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, o)))) - y
            for i in range(DV):
                gw[i] += d * c[i]
            gb += d
            da = [sum(d * P["w"][i] * V[t][i] for i in range(DV)) for t in range(STEPS)]
            mean = sum(a[t] * da[t] for t in range(STEPS))
            for t in range(STEPS):
                ds = a[t] * (da[t] - mean) * SCALE
                for i in range(DK):
                    gq[i] += ds * K[t][i]
                    gWk[i][0] += ds * P["q"][i] * X[t][0]
                    gWk[i][1] += ds * P["q"][i] * X[t][1]
                for i in range(DV):
                    gWv[i][0] += d * a[t] * P["w"][i] * X[t][0]
                    gWv[i][1] += d * a[t] * P["w"][i] * X[t][1]
        n = len(TR_IN)
        for i in range(DK):
            P["q"][i] -= lr * gq[i] / n
            P["Wk"][i][0] -= lr * gWk[i][0] / n
            P["Wk"][i][1] -= lr * gWk[i][1] / n
        for i in range(DV):
            P["w"][i] -= lr * gw[i] / n
            P["Wv"][i][0] -= lr * gWv[i][0] / n
            P["Wv"][i][1] -= lr * gWv[i][1] / n
        P["b"] -= lr * gb / n
    return P


def accuracy(P, G):
    return sum((forward(P, X)[4] > 0) == (y == 1) for X, y, _ in G) / len(G)


NETS = [train_model(SEED + 1), train_model(SEED + 2)]
print(f"parameter count {DK * 2 + DV * 2 + DK + DV + 1}, epochs 200, "
      f"training examples {len(TR_IN)}")
print(f"baseline (most frequent class)   {BASELINE:.4f}")
for i, P in enumerate(NETS):
    print(f"attention net run {i + 1}             train {accuracy(P, TR_IN):.4f}   "
          f"held-out {accuracy(P, TE_IN):.4f}")
print(f"difference between the runs      {abs(accuracy(NETS[0], TE_IN) - accuracy(NETS[1], TE_IN)):.4f}")

parameter count 19, epochs 200, training examples 360
baseline (most frequent class)   0.6333
attention net run 1             train 0.9028   held-out 0.8889
attention net run 2             train 0.9889   held-out 0.9944
difference between the runs      0.1056

Three numbers are in the table. The baseline is 0.6333; the attention network gives 0.8889 and 0.9944 on the held-out set. Nineteen parameters learn to find the correct step in a sequence of length twelve and rise at least 0.2556 above the baseline. A recurrent network could do the same job only by carrying state for twelve steps.

The third number says more than the second. The two runs share capacity, data and epoch budget; only the initialization seeds differ. The difference between them is 0.1056 — about one ninth of this network’s held-out figure — and it comes not from what the network is but from where it started in that run. The rule of the course can be read directly here: the number of a network may come from its training rather than its capacity.

Where the Weights Fall

The weight vector that attention produces is the only readable intermediate product inside the network: it gives one number for each step, and they sum to one. This readability is the reason the weights are so often presented as “what the model looked at”. The setup suits testing this claim, because the correct step is known.

def attention_measure(P, G):
    mass, hit = 0.0, 0
    for X, _y, m in G:
        a = forward(P, X)[2]
        mass += a[m]
        hit += max(range(STEPS), key=lambda t: a[t]) == m
    return mass / len(G), hit / len(G)


print(f"weight of one step in a uniform distribution {1 / STEPS:.4f}")
print(f"{'':16}{'weight on marker step':>21}{'largest weight':>18}{'held-out':>10}")
for i, P in enumerate(NETS):
    k, s = attention_measure(P, TE_IN)
    print(f"run {i + 1:<12}{k:>21.4f}{s:>18.4f}{accuracy(P, TE_IN):>10.4f}")

X0, _y0, m0 = TE_IN[0]
a0 = forward(NETS[0], X0)[2]
a1 = forward(NETS[1], X0)[2]
print(f"single example, marker step {m0}")
print("run 1 weights " + " ".join(f"{u:.2f}" for u in a0))
print("run 2 weights " + " ".join(f"{u:.2f}" for u in a1))

weight of one step in a uniform distribution 0.0833
                weight on marker step    largest weight  held-out
run 1                          0.2721            0.3444    0.8889
run 2                          0.9691            1.0000    0.9944
single example, marker step 11
run 1 weights 0.00 0.00 0.10 0.00 0.01 0.01 0.00 0.29 0.31 0.00 0.10 0.18
run 2 weights 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.97

The second run shows the expected behavior: 0.9691 of the weight falls on the marker step, and the largest weight is on the correct step in the whole held-out set. Since the share of one step in a uniform distribution is 0.0833, this is more than eleven times the baseline.

The first run solves the same task with 0.8889 accuracy, but it sends only 0.2721 of the weight to the marker step and puts the largest weight on the correct step in 0.3444 of the examples. The single-example row shows how this is possible: the marker step is the eleventh step, run 1 gives it 0.18, and distributes the rest to the seventh and eighth steps. The network has found a compromise that uses the marker channel partially and spreads the remaining weight over steps with high readings, and this compromise produces the correct answer in most cases. While the accuracy is 0.8889, the hit rate of the weights on the correct step is 0.3444; there is no necessary link between the two numbers.

A High Weight Is Not Causation

A large weight does not mean that the step determines the output. Whether the two coincide can be tested only with an intervention: the step is dropped from the attention distribution, the remaining weights are renormalized, and how much the output probability moved is examined. This is a measure independent of the size of the weight.

def probability(P, X, skip=None):
    o = forward(P, X, skip)[4]
    return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, o))))


def causality(P, G, threshold=0.05):
    agree, high, high_inert, mean_effect = 0, 0, 0, 0.0
    for X, _y, _m in G:
        a = forward(P, X)[2]
        p0 = probability(P, X)
        effect = [abs(probability(P, X, t) - p0) for t in range(STEPS)]
        ta = max(range(STEPS), key=lambda t: a[t])
        agree += ta == max(range(STEPS), key=lambda t: effect[t])
        mean_effect += effect[ta]
        if a[ta] > 0.5:
            high += 1
            high_inert += effect[ta] < threshold
    n = len(G)
    return agree / n, mean_effect / n, high, high_inert


print(f"{'':10}{'weight-effect agreement':>25}{'effect of heaviest step':>25}"
      f"{'weight>0.5':>13}{'but effect<0.05':>17}")
for i, P in enumerate(NETS):
    u, e, y, ye = causality(P, TE_IN)
    print(f"run {i + 1:<6}{u:>25.4f}{e:>25.4f}{y:>13d}{ye:>17d}")

            weight-effect agreement  effect of heaviest step   weight>0.5  but effect<0.05
run 1                        0.5889                   0.3318           81               27
run 2                        1.0000                   0.4081          180               43

In the first run, the heaviest step and the most influential step are the same in only 0.5889 of the examples. That is, someone who reads the weight ranking and says “the model looked here” misses, in four examples out of ten, the step that actually determines the output. In the second run the agreement is 1.0000. Same architecture, same data, same budget; the explanatory value of the weights changes from run to run. The weights alone do not reveal which run one is in — both produce concentrated, orderly looking distributions.

The last column shows that the second run has a limit within itself as well. In the whole held-out set the largest weight exceeds 0.5, but in 43 of these 180 examples, removing that step moves the prediction probability by less than 0.05. The reason is clear: if the remaining steps also carry low readings, dropping the marker step leaves the average on the same side and the decision does not change. Here a high weight does not mean “this step went into the decision”; it means “this step was read”.

The rule that follows: attention weights are an intermediate product useful in debugging, not a justification document on their own. A weight reported without a removal test or the run-to-run variability beside it is an unmeasured number.

Summary

• The sequence setup was built in this lesson and is fictional: 720 sequences, length 12, split 360/180/180, seed 20260218; the label depends only on the reading of the marker step.
• The baseline is the most frequent class and is 0.6333; the attention network with nineteen parameters gives 0.8889 and 0.9944 on the held-out set, and the difference between the two runs is 0.1056.
• Attention scales the inner product of the query and the key by $\sqrt{d}$, passes it through softmax, and produces the weighted average of the values; there is no hidden state and no recurrence.
• Where the weights fall is not the same thing as accuracy: at 0.8889 accuracy, the first run sends 0.2721 of the weight to the marker step and places the largest weight with a hit rate of 0.3444.
• In the step-removal test, the agreement between weight and effect is 0.5889 in the first run, and in the second run, in 43 of the 180 examples whose weight exceeds 0.5, removing that step moves the prediction by less than 0.05.

Next Step

In this lesson a single query produced a single weight distribution and the sequence was reduced to a single summary. The task suited this: the label came from a single step and step order never entered into it. Two questions remain open. The first: if there are several patterns of interest in the same sequence — one looking for the marker step, the other for the general level — can a single distribution carry both. The second, and the sharper one: since attention looks at all the steps at the same time, how will it know the order of the steps. The next lesson multiplies the query, measures the difference the number of heads adds, and prints against the same baseline how much of the order information remains when the positional encoding is removed.