feat(models): GAT encoder with multi-head attention and residual connections #20

New Issue

forbes · 2026-02-04T20:22:41Z

forbes commented

2026-02-04 20:22:41 +00:00

Summary

Implement solver/models/gat.py with a GATEncoder(nn.Module) using GATv2Conv.

Depends on #18.

Architecture

Input: node features [n, 22], edge_index [2, E], edge_attr [E, 22]
num_layers GATv2Conv layers with num_heads attention heads
Residual connections between layers (togglable via residual config flag)
Edge features via GATv2Conv's edge_dim parameter
Output: node embeddings [n, hidden_dim] and graph-level embedding [hidden_dim]

Config (from `configs/model/gat.yaml`)

architecture: gat
encoder:
  num_layers: 4
  hidden_dim: 256
  num_heads: 8
  dropout: 0.1
  residual: true

Registration

Register with @register_encoder("gat").

Files

solver/models/gat.py
tests/models/test_gat.py

Acceptance criteria

Forward pass produces correct output shapes
Attention weights accessible for interpretability
Residual connections togglable via config
Works with variable-size batched graphs
Registered and resolved by build_model
Passes ruff, mypy, pytest

## Summary Implement `solver/models/gat.py` with a `GATEncoder(nn.Module)` using GATv2Conv. Depends on #18. ## Architecture - Input: node features `[n, 22]`, edge_index `[2, E]`, edge_attr `[E, 22]` - `num_layers` GATv2Conv layers with `num_heads` attention heads - Residual connections between layers (togglable via `residual` config flag) - Edge features via GATv2Conv's `edge_dim` parameter - Output: node embeddings `[n, hidden_dim]` and graph-level embedding `[hidden_dim]` ## Config (from `configs/model/gat.yaml`) ```yaml architecture: gat encoder: num_layers: 4 hidden_dim: 256 num_heads: 8 dropout: 0.1 residual: true ``` ## Registration Register with `@register_encoder("gat")`. ## Files - `solver/models/gat.py` - `tests/models/test_gat.py` ## Acceptance criteria - [ ] Forward pass produces correct output shapes - [ ] Attention weights accessible for interpretability - [ ] Residual connections togglable via config - [ ] Works with variable-size batched graphs - [ ] Registered and resolved by `build_model` - [ ] Passes ruff, mypy, pytest

forbes referenced this issue

2026-02-04 20:22:53 +00:00

feat(models): multi-task prediction heads for constraint analysis #21

forbes referenced this issue

2026-02-04 20:23:08 +00:00

feat(models): assembled multi-task model and loss function #22

forbes closed this issue

2026-02-25 19:04:51 +00:00

Sign in to join this conversation.