# How to Implement Custom Attention Mechanisms in the Hierarchical Attention Network

> Learn to implement custom attention mechanisms in your Hierarchical Attention Network by subclassing and overriding the forward method. Explore dot-product, multi-head, and normalized variants easily.

- Repository: [jerrygao/rpdnn](https://github.com/jerrygaolondon/rpdnn)
- Tags: how-to-guide
- Published: 2026-03-04

---

**Subclass the `HierarchicalAttentionNet` class in [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py) and override the scoring logic in the `forward` method to implement custom attention mechanisms such as dot-product, multi-head, or normalized attention variants.**

The `jerrygaolondon/rpdnn` repository provides a production-ready hierarchical attention network for rumor detection tasks. While the baseline implementation in [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py) uses additive (Bahdanau) attention to compute context vectors, its modular design allows developers to swap in custom attention mechanisms by subclassing the base class and modifying the scoring function.

## Core Architecture of the Baseline Attention Layer

The `HierarchicalAttentionNet` class implements a standard additive attention mechanism over RNN hidden states. According to the source code in [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py), the constructor initializes learnable parameters for the attention scoring function:

```python

# From src/attention.py, lines 55-71

def __init__(self, feature_dim, step_dim, bias=True, **kwargs):
    super(HierarchicalAttentionNet, self).__init__(**kwargs)
    self.feature_dim = feature_dim
    self.step_dim = step_dim
    self.bias = bias
    self.weight = nn.Parameter(torch.Tensor(feature_dim, 1))
    if bias:
        self.b = nn.Parameter(torch.Tensor(step_dim))
    # ... initialization code

```

The `forward` method (lines 72-88) executes the attention computation in six distinct steps:

1. **Linear projection**: Computes raw scores via `eij = x·self.weight`
2. **Bias addition**: Adds the optional per-timestep bias
3. **Non-linearity**: Applies `tanh` activation
4. **Masking**: Fills padded positions with `-inf` to ignore them in softmax
5. **Softmax**: Converts scores to attention distribution `a`
6. **Weighted sum**: Computes `context_weighted_sum = Σ_t a_t * x_t`

The method returns a tuple of `(context_weighted_sum, weighted_input, attention_weights)`, making it compatible with downstream layers in the `RumorTweetsClassifer`.

## Strategies for Customizing the Attention Mechanism

The simplicity of the baseline `HierarchicalAttentionNet` makes it an ideal foundation for implementing custom attention mechanisms. You can modify the scoring function, add normalization, or implement multi-head variants without changing the interface.

### Replace the Scoring Function with Dot-Product Attention

Instead of the additive `tanh(x·W + b)` scoring, implement **dot-product attention** by computing the similarity between each hidden state and a learnable context vector:

```python
import torch
from torch import nn
import torch.nn.functional as F

class DotProductHierarchicalAttention(nn.Module):
    """
    Custom hierarchical attention that computes a dot-product between each hidden state
    and a learnable context vector. Optional mask handling follows the original
    implementation.
    """
    def __init__(self, feature_dim: int, step_dim: int, bias: bool = True, **kwargs):
        super().__init__(**kwargs)
        self.feature_dim = feature_dim
        self.step_dim = step_dim
        self.bias = bias
        # Learnable context vector instead of weight matrix

        self.context = nn.Parameter(torch.randn(feature_dim))
        if bias:
            self.b = nn.Parameter(torch.zeros(step_dim))

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        """
        x: (batch, steps, hidden_dim)
        mask: (batch, steps) – bool tensor where True = real token, False = padding
        Returns (context_sum, weighted_input, attn_weights)
        """
        # Dot-product scoring: (batch, steps, hidden) · (hidden,) -> (batch, steps)

        scores = torch.matmul(x, self.context)
        
        if self.bias:
            scores = scores + self.b
            
        scores = torch.tanh(scores)  # Optional: change to relu/gelu

        
        if mask is not None:
            scores = scores.masked_fill(~mask, float("-inf"))
            
        attn_weights = F.softmax(scores, dim=1)
        weighted_input = x * attn_weights.unsqueeze(-1)
        context_sum = weighted_input.sum(dim=1)
        
        return context_sum, weighted_input, attn_weights

```

### Implement Multi-Head Attention

For capturing different representation subspaces, split the hidden dimension into multiple heads. The repository's `StructuredSelfAttention` class in [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py) demonstrates single-head self-attention patterns that you can extend:

```python
class MultiHeadHierarchicalAttention(nn.Module):
    """
    Splits the hidden dimension into `num_heads` parts and applies an independent
    HierarchicalAttentionNet per head. The final context vector is the concatenation
    of all heads.
    """
    def __init__(self, feature_dim: int, step_dim: int,
                 num_heads: int = 4, bias: bool = True, **kwargs):
        super().__init__(**kwargs)
        assert feature_dim % num_heads == 0, "feature_dim must be divisible by num_heads"
        self.head_dim = feature_dim // num_heads
        self.num_heads = num_heads
        
        # Import the base class from the repository

        from attention import HierarchicalAttentionNet
        
        self.heads = nn.ModuleList([
            HierarchicalAttentionNet(self.head_dim, step_dim, bias, **kwargs)
            for _ in range(num_heads)
        ])
        
    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        # Reshape to split hidden dim: (batch, steps, heads, head_dim)

        x_split = x.view(x.size(0), x.size(1), self.num_heads, self.head_dim)
        
        # Process each head

        contexts, weighted_inputs, attn_weights = [], [], []
        for i, head in enumerate(self.heads):
            ctx, w, a = head(x_split[:, :, i, :], mask)
            contexts.append(ctx)
            weighted_inputs.append(w)
            attn_weights.append(a)
            
        # Concatenate head outputs: (batch, feature_dim)

        final_context = torch.cat(contexts, dim=1)
        final_weighted = torch.cat(weighted_inputs, dim=2)
        final_attn = torch.stack(attn_weights, dim=1)  # (batch, heads, steps)

        
        return final_context, final_weighted, final_attn

```

### Add Layer Normalization

The repository includes `MyLayerNorm` in [`src/my_layer_norm.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/my_layer_norm.py). Apply normalization before the attention computation to stabilize training:

```python
from my_layer_norm import MyLayerNorm

class NormalizedHierarchicalAttention(HierarchicalAttentionNet):
    def __init__(self, feature_dim, step_dim, bias=True, **kwargs):
        super().__init__(feature_dim, step_dim, bias, **kwargs)
        self.layer_norm = MyLayerNorm(feature_dim)
        
    def forward(self, x, mask=None):
        # Normalize before attention scoring

        x = self.layer_norm(x)
        return super().forward(x, mask)

```

## Integrating Custom Attention into the Rumor Classifier

The `RumorTweetsClassifer` in [`src/allennlp_rumor_classifier.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/allennlp_rumor_classifier.py) instantiates attention layers inside the `set_attention_mechanism` method (lines 326-352). To use your custom implementation:

1. **Import your subclass** at the top of [`src/allennlp_rumor_classifier.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/allennlp_rumor_classifier.py)
2. **Replace the constructor call** in `set_attention_mechanism`:

```python

# Original instantiation (lines 326-352)

self.cxt_content_attention = HierarchicalAttentionNet(
    self.cxt_content_encoder.get_output_dim(),
    self.max_cxt_size)

# Replace with custom implementation

from attention import DotProductHierarchicalAttention
self.cxt_content_attention = DotProductHierarchicalAttention(
    self.cxt_content_encoder.get_output_dim(),
    self.max_cxt_size)

```

The classifier expects the attention object to return a tuple of `(context_vector, weighted_input, attention_weights)`, so ensure your custom implementation maintains this interface.

## Key Files for Reference

| File | Purpose | Key Components |
|------|---------|----------------|
| [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py) | Core attention implementations | `HierarchicalAttentionNet`, `StructuredSelfAttention` |
| [`src/allennlp_rumor_classifier.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/allennlp_rumor_classifier.py) | Model integration point | `set_attention_mechanism`, `RumorTweetsClassifer` |
| [`src/my_layer_norm.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/my_layer_norm.py) | Normalization utilities | `MyLayerNorm` |
| [`src/embeddings/embedding_layer.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/embeddings/embedding_layer.py) | Input preprocessing | ELMo embedding layer |

These files illustrate how attention is built, instantiated, and consumed in the repository, making them the natural extension points for any custom attention mechanism.

## Summary

- The `HierarchicalAttentionNet` class in [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py) provides a clean, extensible implementation of additive attention for rumor detection models.
- **Subclassing** is the recommended approach for implementing custom attention mechanisms—override the `forward` method to modify scoring functions while preserving the return signature `(context_sum, weighted_input, attn_weights)`.
- Common customizations include **dot-product scoring**, **multi-head attention**, and **layer normalization**, all achievable with minimal code changes.
- Integration requires only swapping the class instantiation in [`src/allennlp_rumor_classifier.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/allennlp_rumor_classifier.py) within the `set_attention_mechanism` method.

## Frequently Asked Questions

### What is the default attention mechanism used in the hierarchical attention network?

The default implementation uses **additive (Bahdanau) attention**, which computes scores via a linear projection followed by a `tanh` non-linearity. This is implemented in the `HierarchicalAttentionNet` class in [`src/attention.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/attention.py) lines 55-88.

### Can I use multi-head attention with the existing codebase?

Yes. You can implement multi-head attention by subclassing `HierarchicalAttentionNet` and splitting the hidden dimension across multiple heads, as demonstrated in the `MultiHeadHierarchicalAttention` example. Each head processes a slice of the feature dimension independently, and their outputs are concatenated to form the final context vector.

### How do I ensure my custom attention layer works with the rumor classifier?

Your custom layer must maintain the same **return signature** as the base class: `(context_sum, weighted_input, attention_weights)`. The `context_sum` should have shape `(batch, feature_dim)`, `weighted_input` should have shape `(batch, steps, feature_dim)`, and `attention_weights` should have shape `(batch, steps)`. As long as these shapes match, you can drop your custom class into [`src/allennlp_rumor_classifier.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/allennlp_rumor_classifier.py) without modifying the classifier logic.

### Where should I add layer normalization in the attention computation?

Add normalization **before** the attention scoring to stabilize hidden state distributions. You can use the repository's `MyLayerNorm` from [`src/my_layer_norm.py`](https://github.com/jerrygaolondon/rpdnn/blob/main/src/my_layer_norm.py) by calling it at the beginning of your custom `forward` method before computing attention scores, or subclass `HierarchicalAttentionNet` and wrap the input `x` with the normalization layer.