{
  "schemaVersion": 1,
  "slug": "bahdanau-attention",
  "title": "Choose source context for the next decoder step",
  "className": "AdditiveDecoderStep",
  "question": "How does a recurrent decoder decide which source positions to consult?",
  "description": "Trace bidirectional annotations into additive alignment, then change the previous decoder state and inspect its weighted source context.",
  "category": "Language milestones",
  "tags": [
    "Bahdanau attention",
    "Additive attention",
    "Alignment",
    "GRU",
    "Translation"
  ],
  "curator": "TensorViz",
  "attribution": "Original TensorViz teaching example. PyTorch provides the underlying operators.",
  "license": "No separate redistribution license has been declared for these project examples.",
  "limitations": "Untrained attention subsystem and a single decoder step, not a complete translation system. It omits the paper's exact decoder/maxout output formulation, variable-length padding and training recipe. Publication began as a 2014 preprint and appeared at ICLR 2015.",
  "openingStep": "align",
  "provenance": {
    "capturedAt": "2026-09-17T06:32:41.318529+00:00",
    "productRevision": "cf0ab1c9c58c1dc640235b742fa032cade00b43c",
    "sourceSha256": "5fadad4ddf6e107b09d66c42307c4804cae1bbd5346ef8e4c927b6db0034a44c",
    "python": "3.13.13",
    "torch": "2.7.1",
    "execution": {
      "inputs": {
        "args": [
          {
            "shape": [
              1,
              5
            ],
            "dtype": "int64",
            "fill": "ones"
          },
          {
            "shape": [
              1,
              8
            ],
            "dtype": "float32",
            "fill": "randn"
          },
          {
            "shape": [
              1
            ],
            "dtype": "int64",
            "fill": "ones"
          }
        ]
      },
      "device": "cpu",
      "mode": "eval",
      "seed": 0,
      "deterministic": true
    },
    "outcome": "passed",
    "outputShape": [
      1,
      16
    ],
    "method": "TensorViz multifile projection and isolated execution worker",
    "scope": "One reduced additive-attention decoder step with a bidirectional GRU encoder. Alignment normalization, context-sum, bidirectional influence and decoder-state sensitivity checks.",
    "numericalChecks": [
      "Alignment weights are positive and sum to one over source positions",
      "The context equals an independent weighted sum of bidirectional annotations",
      "A changed final source token can affect the first annotation through the backward GRU",
      "Changing decoder state changes alignment; zero scores give the mean source annotation"
    ]
  },
  "layers": [
    {
      "id": "input",
      "label": "input",
      "op": "Input",
      "kind": "input",
      "parent": null,
      "params": {
        "shape": "1,8"
      }
    },
    {
      "id": "source_embedding",
      "label": "source_embedding",
      "op": "Embedding",
      "kind": "layer",
      "parent": null,
      "params": {
        "num_embeddings": 16,
        "embedding_dim": 4
      },
      "source": {
        "line": 31,
        "endLine": 31
      }
    },
    {
      "id": "encoder",
      "label": "encoder",
      "op": "GRU",
      "kind": "layer",
      "parent": null,
      "params": {
        "input_size": 4,
        "hidden_size": 4,
        "batch_first": true,
        "bidirectional": true
      },
      "source": {
        "line": 32,
        "endLine": 32
      }
    },
    {
      "id": "alignment",
      "label": "alignment",
      "op": "Module",
      "kind": "module",
      "parent": null,
      "params": {},
      "source": {
        "line": 33,
        "endLine": 33
      }
    },
    {
      "id": "target_embedding",
      "label": "target_embedding",
      "op": "Embedding",
      "kind": "layer",
      "parent": null,
      "params": {
        "num_embeddings": 16,
        "embedding_dim": 4
      },
      "source": {
        "line": 34,
        "endLine": 34
      }
    },
    {
      "id": "cat@AdditiveDecoderStep#0",
      "label": "cat@AdditiveDecoderStep#0",
      "op": "torch.cat",
      "kind": "layer",
      "parent": null,
      "params": {
        "dim": -1
      },
      "source": {
        "line": 42,
        "endLine": 44
      }
    },
    {
      "id": "decoder",
      "label": "decoder",
      "op": "GRUCell",
      "kind": "layer",
      "parent": null,
      "params": {
        "input_size": 12,
        "hidden_size": 8
      },
      "source": {
        "line": 35,
        "endLine": 35
      }
    },
    {
      "id": "output_head",
      "label": "output_head",
      "op": "Linear",
      "kind": "layer",
      "parent": null,
      "params": {
        "in_features": 8,
        "out_features": 16
      },
      "source": {
        "line": 36,
        "endLine": 36
      }
    },
    {
      "id": "output",
      "label": "output",
      "op": "Output",
      "kind": "output",
      "parent": null,
      "params": {}
    },
    {
      "id": "alignment.annotation_projection",
      "label": "alignment.annotation_projection",
      "op": "Linear",
      "kind": "layer",
      "parent": "alignment",
      "params": {
        "in_features": 8,
        "out_features": 8,
        "bias": false
      },
      "source": {
        "line": 13,
        "endLine": 13
      }
    },
    {
      "id": "alignment.state_projection",
      "label": "alignment.state_projection",
      "op": "Linear",
      "kind": "layer",
      "parent": "alignment",
      "params": {
        "in_features": 8,
        "out_features": 8,
        "bias": false
      },
      "source": {
        "line": 14,
        "endLine": 14
      }
    },
    {
      "id": "add@alignment#0",
      "label": "add@alignment#0",
      "op": "torch.add",
      "kind": "layer",
      "parent": "alignment",
      "params": {},
      "source": {
        "line": 20,
        "endLine": 23
      }
    },
    {
      "id": "alignment.tanh",
      "label": "alignment.tanh",
      "op": "Tanh",
      "kind": "layer",
      "parent": "alignment",
      "params": {},
      "source": {
        "line": 15,
        "endLine": 15
      }
    },
    {
      "id": "alignment.score",
      "label": "alignment.score",
      "op": "Linear",
      "kind": "layer",
      "parent": "alignment",
      "params": {
        "in_features": 8,
        "out_features": 1,
        "bias": false
      },
      "source": {
        "line": 16,
        "endLine": 16
      }
    },
    {
      "id": "alignment.softmax",
      "label": "alignment.softmax",
      "op": "Softmax",
      "kind": "layer",
      "parent": "alignment",
      "params": {
        "dim": 1
      },
      "source": {
        "line": 17,
        "endLine": 17
      }
    },
    {
      "id": "mul@alignment#0",
      "label": "mul@alignment#0",
      "op": "torch.mul",
      "kind": "layer",
      "parent": "alignment",
      "params": {},
      "source": {
        "line": 25,
        "endLine": 25
      }
    },
    {
      "id": "python@alignment#0",
      "label": "python@alignment#0",
      "op": "Python",
      "kind": "layer",
      "parent": "alignment",
      "params": {},
      "source": {
        "line": 25,
        "endLine": 25
      }
    }
  ],
  "connections": [
    {
      "id": "edge-1",
      "source": "alignment.annotation_projection",
      "target": "add@alignment#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-2",
      "source": "alignment.state_projection",
      "target": "add@alignment#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-3",
      "source": "add@alignment#0",
      "target": "alignment.tanh",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-4",
      "source": "alignment.tanh",
      "target": "alignment.score",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-5",
      "source": "alignment.score",
      "target": "alignment.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-6",
      "source": "alignment.softmax",
      "target": "mul@alignment#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-7",
      "source": "mul@alignment#0",
      "target": "python@alignment#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-8",
      "source": "input",
      "target": "source_embedding",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-9",
      "source": "source_embedding",
      "target": "encoder",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-10",
      "source": "encoder",
      "target": "alignment.annotation_projection",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-11",
      "source": "encoder",
      "target": "alignment.state_projection",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-12",
      "source": "encoder",
      "target": "mul@alignment#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-13",
      "source": "encoder",
      "target": "alignment",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-14",
      "source": "input",
      "target": "alignment.annotation_projection",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-15",
      "source": "input",
      "target": "alignment.state_projection",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-16",
      "source": "input",
      "target": "mul@alignment#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-17",
      "source": "input",
      "target": "alignment",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-18",
      "source": "input",
      "target": "target_embedding",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-19",
      "source": "target_embedding",
      "target": "cat@AdditiveDecoderStep#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-20",
      "source": "python@alignment#0",
      "target": "cat@AdditiveDecoderStep#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-21",
      "source": "alignment",
      "target": "cat@AdditiveDecoderStep#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-22",
      "source": "cat@AdditiveDecoderStep#0",
      "target": "decoder",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-23",
      "source": "input",
      "target": "decoder",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-24",
      "source": "decoder",
      "target": "output_head",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-25",
      "source": "output_head",
      "target": "output",
      "sourcePort": null,
      "targetPort": null
    }
  ],
  "steps": [
    {
      "id": "encode",
      "target": "encoder",
      "title": "Keep an annotation for every source position",
      "note": "Forward and backward GRUs produce four features each, giving an eight-feature annotation at each of five positions. The decoder can consult this entire sequence instead of relying on one fixed summary."
    },
    {
      "id": "align",
      "target": "alignment.tanh",
      "title": "Combine source annotations with the previous state",
      "note": "Learned projections of each annotation and the previous decoder state are added before tanh. A learned scalar projection turns these joint features into an alignment score."
    },
    {
      "id": "context",
      "target": "alignment.softmax",
      "title": "Normalize alignment and form a weighted context",
      "note": "Softmax normalizes the five scores across source positions. Their weighted annotation sum becomes an eight-feature context vector. The recorded comparison changes the previous decoder state while holding the source fixed."
    },
    {
      "id": "decode",
      "target": "decoder",
      "title": "Use the context to update the decoder",
      "note": "The previous output-token embedding and the context are concatenated into twelve input features. A GRUCell updates the previous eight-feature state, and the output head produces sixteen next-token logits."
    }
  ],
  "milestone": {
    "name": "Bahdanau additive attention",
    "year": 2014,
    "kind": "Component",
    "fidelity": "Five source tokens with width-four embeddings; bidirectional GRU annotations of width eight; additive alignment and one width-eight GRUCell decoder update. A simple linear vocabulary head replaces the original output architecture.",
    "papers": [
      {
        "title": "Neural Machine Translation by Jointly Learning to Align and Translate",
        "url": "https://arxiv.org/abs/1409.0473"
      }
    ]
  },
  "experiment": {
    "kind": "matrices",
    "title": "Ask the source a different question at each step.",
    "description": "Recorded additive alignment for a fixed source sequence under three decoder states. In translation, the evolving decoder state makes this context change at every generated token. This lesson captures one such step.",
    "controlLabel": "Previous decoder state",
    "cases": [
      {
        "id": "zero",
        "label": "Zero previous decoder state",
        "target": "alignment.softmax",
        "note": "The same five synthetic source IDs produce fixed bidirectional annotations. Only the previous decoder state changes here. The learned additive score determines the weighting used by this step. These untrained alignments do not map translated words or demonstrate translation quality.",
        "matrices": [
          {
            "label": "Source annotations · five positions × eight features",
            "values": [
              [
                -0.12753956004672473,
                -0.010182755237456913,
                0.24205580434611398,
                -0.006486028054295113,
                -0.17509282153837888,
                0.006958515441951252,
                0.10051159809073074,
                0.08204099708252638
              ],
              [
                -0.0804471440248747,
                0.1992614234598769,
                0.29115099896269575,
                0.3643333629739864,
                0.14285933833558961,
                0.28697952455997605,
                -0.21023592053671425,
                0.24920484244207416
              ],
              [
                0.018382988969698144,
                0.3666580044553704,
                0.31440636536360494,
                0.5996918391125939,
                0.272370216111587,
                0.48810228865427047,
                -0.11815556788997617,
                0.1885173475576282
              ],
              [
                -0.32190792116523237,
                0.2817262926221353,
                0.4325787696977407,
                0.6003770636368334,
                0.23025333767429998,
                0.6346657876475686,
                0.11543248602880651,
                0.02378627795358143
              ],
              [
                -0.1695092048058527,
                0.3599568286633815,
                0.4324697436135898,
                0.5783880559169529,
                0.12421300302613837,
                0.165863742223493,
                -0.06991413852559067,
                0.07757263418012064
              ]
            ]
          },
          {
            "label": "Alignment · one decoder step × five source positions",
            "values": [
              [
                0.1867198585575624,
                0.198990558816515,
                0.2095750168437494,
                0.2075075224321882,
                0.197207043349985
              ]
            ]
          }
        ],
        "vectors": [
          {
            "label": "Previous decoder state",
            "values": [
              0,
              0,
              0,
              0,
              0,
              0,
              0,
              0
            ]
          },
          {
            "label": "Weighted context",
            "values": [
              -0.13619649980682053,
              0.24403852374718826,
              0.3440740730840808,
              0.4356132119525959,
              0.12509132402973094,
              0.32510697771644986,
              -0.03766435842570114,
              0.12465032195579218
            ]
          }
        ],
        "metrics": [
          {
            "label": "Alignment weight sum",
            "value": 1
          },
          {
            "label": "Most weighted source position · zero-based",
            "value": 2
          }
        ]
      },
      {
        "id": "positive",
        "label": "Positive previous decoder state",
        "target": "alignment.softmax",
        "note": "The same five synthetic source IDs produce fixed bidirectional annotations. Only the previous decoder state changes here. The learned additive score determines the weighting used by this step. These untrained alignments do not map translated words or demonstrate translation quality.",
        "matrices": [
          {
            "label": "Source annotations · five positions × eight features",
            "values": [
              [
                -0.12753956004672473,
                -0.010182755237456913,
                0.24205580434611398,
                -0.006486028054295113,
                -0.17509282153837888,
                0.006958515441951252,
                0.10051159809073074,
                0.08204099708252638
              ],
              [
                -0.0804471440248747,
                0.1992614234598769,
                0.29115099896269575,
                0.3643333629739864,
                0.14285933833558961,
                0.28697952455997605,
                -0.21023592053671425,
                0.24920484244207416
              ],
              [
                0.018382988969698144,
                0.3666580044553704,
                0.31440636536360494,
                0.5996918391125939,
                0.272370216111587,
                0.48810228865427047,
                -0.11815556788997617,
                0.1885173475576282
              ],
              [
                -0.32190792116523237,
                0.2817262926221353,
                0.4325787696977407,
                0.6003770636368334,
                0.23025333767429998,
                0.6346657876475686,
                0.11543248602880651,
                0.02378627795358143
              ],
              [
                -0.1695092048058527,
                0.3599568286633815,
                0.4324697436135898,
                0.5783880559169529,
                0.12421300302613837,
                0.165863742223493,
                -0.06991413852559067,
                0.07757263418012064
              ]
            ]
          },
          {
            "label": "Alignment · one decoder step × five source positions",
            "values": [
              [
                0.19767351232879388,
                0.20032882675912617,
                0.20289220137491626,
                0.20146124506447577,
                0.19764421447268785
              ]
            ]
          }
        ],
        "vectors": [
          {
            "label": "Previous decoder state",
            "values": [
              3,
              3,
              3,
              3,
              3,
              3,
              3,
              3
            ]
          },
          {
            "label": "Weighted context",
            "values": [
              -0.13595179389782547,
              0.24019731018709012,
              0.34258755894299875,
              0.4286449203118779,
              0.1202065288576556,
              0.31854050230876424,
              -0.03678363058644463,
              0.12451274091978563
            ]
          }
        ],
        "metrics": [
          {
            "label": "Alignment weight sum",
            "value": 0.9999999999999999
          },
          {
            "label": "Most weighted source position · zero-based",
            "value": 2
          }
        ]
      },
      {
        "id": "negative",
        "label": "Negative previous decoder state",
        "target": "alignment.softmax",
        "note": "The same five synthetic source IDs produce fixed bidirectional annotations. Only the previous decoder state changes here. The learned additive score determines the weighting used by this step. These untrained alignments do not map translated words or demonstrate translation quality.",
        "matrices": [
          {
            "label": "Source annotations · five positions × eight features",
            "values": [
              [
                -0.12753956004672473,
                -0.010182755237456913,
                0.24205580434611398,
                -0.006486028054295113,
                -0.17509282153837888,
                0.006958515441951252,
                0.10051159809073074,
                0.08204099708252638
              ],
              [
                -0.0804471440248747,
                0.1992614234598769,
                0.29115099896269575,
                0.3643333629739864,
                0.14285933833558961,
                0.28697952455997605,
                -0.21023592053671425,
                0.24920484244207416
              ],
              [
                0.018382988969698144,
                0.3666580044553704,
                0.31440636536360494,
                0.5996918391125939,
                0.272370216111587,
                0.48810228865427047,
                -0.11815556788997617,
                0.1885173475576282
              ],
              [
                -0.32190792116523237,
                0.2817262926221353,
                0.4325787696977407,
                0.6003770636368334,
                0.23025333767429998,
                0.6346657876475686,
                0.11543248602880651,
                0.02378627795358143
              ],
              [
                -0.1695092048058527,
                0.3599568286633815,
                0.4324697436135898,
                0.5783880559169529,
                0.12421300302613837,
                0.165863742223493,
                -0.06991413852559067,
                0.07757263418012064
              ]
            ]
          },
          {
            "label": "Alignment · one decoder step × five source positions",
            "values": [
              [
                0.1975477567897031,
                0.20043415623763802,
                0.20281020543790013,
                0.2016360664448145,
                0.19757181508994423
              ]
            ]
          }
        ],
        "vectors": [
          {
            "label": "Previous decoder state",
            "values": [
              -3,
              -3,
              -3,
              -3,
              -3,
              -3,
              -3,
              -3
            ]
          },
          {
            "label": "Weighted context",
            "values": [
              -0.1359897399031528,
              0.24021270548728954,
              0.34260631929832114,
              0.428698022523444,
              0.12025252205991319,
              0.318628776952545,
              -0.03678348443300138,
              0.12451175690862279
            ]
          }
        ],
        "metrics": [
          {
            "label": "Alignment weight sum",
            "value": 1
          },
          {
            "label": "Most weighted source position · zero-based",
            "value": 2
          }
        ]
      }
    ]
  },
  "revision": "6f6805ea39bc1506",
  "files": {
    "graph.tensorviz.json": {
      "url": "/models/bahdanau-attention/6f6805ea39bc1506/graph.tensorviz.json",
      "sha256": "38a2e980e056039703946b622c93544174c0b4b16625f079a1bb5f5ae8054819",
      "bytes": 48299
    },
    "model.py": {
      "url": "/models/bahdanau-attention/6f6805ea39bc1506/model.py",
      "sha256": "5fadad4ddf6e107b09d66c42307c4804cae1bbd5346ef8e4c927b6db0034a44c",
      "bytes": 6607
    },
    "run.py": {
      "url": "/models/bahdanau-attention/6f6805ea39bc1506/run.py",
      "sha256": "c253be0922afac13a76ae9952960c9094789259722e19d462fddb3de606eac2b",
      "bytes": 750
    },
    "requirements.txt": {
      "url": "/models/bahdanau-attention/6f6805ea39bc1506/requirements.txt",
      "sha256": "edb86016b42cb3cc6cdc0c490aab929d3bb098ced4515394f0b55766033e0daf",
      "bytes": 91
    },
    "preview.svg": {
      "url": "/models/bahdanau-attention/6f6805ea39bc1506/preview.svg",
      "sha256": "6b9641d777dc748130c6f8f98ce7ab984b7e16a60dc3c318d42b62eabe5c3d99",
      "bytes": 2186
    }
  }
}