Using Spatiotemporal Graph Neural Networks to Decode Caffeine’s Effect on Brain Connectivity

Ever wondered how your morning coffee actually changes your brain’s activity patterns? As researchers and caffeine lovers studying brain connectivity, we’ve often pondered this question. Today, we’re excited to share our recent work using cutting-edge artificial intelligence (AI) techniques to decode how caffeine — the world’s most widely consumed psychoactive substance — affects brain dynamics.

The Challenge: Decoding Dynamic Brain States

Understanding how cognitive states manifest in brain activity patterns remains one of the biggest challenges in neuroscience. While we know caffeine enhances attention, alertness, and motor function through its effects on adenosine receptors, the precise ways it reshapes communication patterns between different brain regions over time have remained elusive.

Traditional approaches to analyzing brain imaging data often treat spatial and temporal patterns separately. However, brain activity is inherently spatiotemporal — brain regions influence each other’s activity patterns across both space and time in complex ways. This is where spatiotemporal Graph Neural Networks (GNNs) come in. These sophisticated deep learning architectures can simultaneously capture both:

1. The spatial organization of brain networks — which brain regions are connected to each other
2. How these connection patterns evolve dynamically over time

Mathematically, we can represent the brain as a dynamic graph G(V, E, T), where:

• V represents brain regions
• E represents functional connections between regions
• T captures how these connections change over time

The Dataset: MyConnectome Project

Our work leverages the MyConnectome dataset — an unprecedented collection of brain scans from a single individual over 18 months, including carefully documented caffeine consumption. This dataset is particularly valuable for our research for several reasons:

1. Controlled Variables: By focusing on a single individual, we eliminate inter-subject variability, which allows us to isolate the effects of caffeine more precisely.
2. Temporal Depth: The dataset includes 72 morning scans (40 non-caffeinated, 32 caffeinated), each containing 518 time points and substantial temporal information about brain state changes.

Data Preprocessing and Graph Construction

Fig. 1. Types of graphs constructed from MyConnectome resting-state fMRI

To transform the raw neuroimaging data into a format suitable for our GNN models, we implemented several preprocessing steps:

Using resting-state fMRI data (3D T1-weighted scans+ BOLD scans), we conduct minimal processing via fMRIPrep. We regress movement out of the resulting timeseries data and use nilearn to derive functional connectivity matrices (weighted adjacency matrices) using the AAL 116 parcellation atlas. This segments the brain into 116 distinct regions of interest (ROIs).

We then constructed two complementary graph representations to capture different aspects of brain dynamics.

The first representation focuses on spatial relationships between brain regions. In this construction, each node represents one of the 116 ROIs, with the node features comprising the time series data for that region. We establish edges between regions based on their functional connectivity, specifically when the absolute value of the Pearson correlation coefficient between their respective time series exceeds a threshold of 0.1. This graph structure emphasizes the spatial relationships and functional connectivity patterns between different brain regions.

def create_spatial_graph(connectivity_matrix, threshold=1.0):
   """
   Creates a spatial graph representation where nodes are brain ROIs and edges
   represent functional connectivity between regions. Edge weights are determined
   by the input connectivity matrix.

   Parameters:
       connectivity_matrix: numpy array of shape (num_rois, num_rois)
                          Contains functional connectivity values between ROI pairs
       threshold: float, optional (default=1.0)
                 Minimum connectivity value required to create an edge

   Returns:
       edge_index: torch.tensor of shape (2, num_edges)
                  Contains indices of connected node pairs
       edge_weights: torch.tensor of shape (num_edges,)
                    Contains connectivity values for each edge

   Notes:
       - Edges are created between ROI pairs with connectivity > threshold
       - Graph is undirected, so each edge is only included once
       - Output format is compatible with PyTorch Geometric
   """
   # Create edges between ROIs with connectivity above threshold
   edge_index = []
   edge_weights = []
   num_rois = len(connectivity_matrix)

   for i in range(num_rois):
       for j in range(i + 1, num_rois):  # Upper triangle only for undirected graph
           if connectivity_matrix[i][j] > threshold:
               edge_index.append([i, j])
               edge_weights.append(connectivity_matrix[i][j])

   # Convert to PyTorch tensors in required format
   return torch.tensor(edge_index).t(), torch.tensor(edge_weights)

The second representation emphasizes temporal dynamics. Here, we construct a graph where each node represents a time point, with the node features consisting of the ROI values at that specific moment. To optimize for computational efficiency while preserving temporal dependencies, we use 100 time points (~two minutes of data) for this representation. Edges in this temporal graph connect time points whose correlation coefficients exceed 0.1.

def create_temporal_graph(connectivity_matrix, timeseries_matrix, num_time_points=100, threshold=0.1):
   """
   Creates a temporal graph representation where nodes are time points and features are ROI values. 
   Edges represent temporal correlations between time points based on the connectivity matrix.

   Parameters:
       connectivity_matrix: numpy array of shape (num_time_points, num_time_points)
                          Contains connectivity values between time points
       timeseries_matrix: numpy array of shape (num_time_points, num_rois)
                         Contains ROI values for each time point
       num_time_points: int, optional (default=100)
                       Number of time points to use for graph construction
       threshold: float, optional (default=0.1)
                 Minimum connectivity value required to create an edge

   Returns:
       x: torch.tensor of shape (num_time_points, num_rois)
          Node features containing ROI values at each time point
       edge_index: torch.tensor of shape (2, num_edges)
                  Contains indices of connected time point pairs
       edge_attr: torch.tensor of shape (num_edges,)
                 Contains connectivity values for each edge

   Notes:
       - Only uses the first num_time_points from the input matrices
       - Graph is undirected, so each edge is only included once
       - Output format is compatible with PyTorch Geometric
   """
   # Select subset of time points
   conn_subset = connectivity_matrix[:num_time_points, :num_time_points]
   time_subset = timeseries_matrix[:num_time_points]

   # Node features are ROI values at each time point
   x = torch.tensor(time_subset, dtype=torch.float)

   # Create edges between time points based on connectivity
   edge_index = []
   edge_weights = []

   for i in range(num_time_points):
       for j in range(i + 1, num_time_points):
           if conn_subset[i][j] > threshold:
               edge_index.append([i, j])
               edge_weights.append(conn_subset[i][j])

   edge_index = torch.tensor(edge_index, dtype=torch.long).t()
   edge_attr = torch.tensor(edge_weights, dtype=torch.float)

   return x, edge_index, edge_attr

These dual graph representations provide complementary views of the brain’s response to caffeine. The spatial graph (Graph 1) allows us to analyze how caffeine affects functional connectivity between brain regions, while the temporal graph (Graph 2) helps us understand how brain activity patterns evolve over time. This dual approach enables our models to capture both spatial and temporal aspects of caffeine’s effects on brain function, thus providing a more complete picture of brain state changes.

Research Objectives

Our study aims to address three key questions:

1. Can spatiotemporal GNNs effectively capture the complex dependencies between spatial and temporal dimensions in brain activity patterns?
2. How do caffeine-induced changes manifest in functional brain networks, particularly those associated with motor control and attention?
3. What is the predictive accuracy of spatiotemporal GNNs in classifying caffeinated versus non-caffeinated brain states?

The implications extend far beyond just understanding your morning coffee. This research could help develop more sensitive methods to detect subtle changes in brain states and improve our understanding of how pharmacological interventions affect the brain.

In the following sections, we’ll dive deeper into our methodology, results, and what they mean for the future of brain state decoding.

Model Architecture and Implementation

Fig. 2. Summarize model design and architecture choices for graph types

Baseline Models

To establish a performance benchmark and validate the benefits of our spatiotemporal approach, we implemented several baseline models that capture either spatial or temporal aspects independently.

Traditional Machine Learning Approach Our simplest baseline is a Multi-Layer Perceptron (MLP) that operates on the flattened adjacency matrix (116×116) of our brain region graph. This approach, while straightforward, serves as a useful baseline to understand the value added by more sophisticated architectures. The network reduces the input to 64 features through a linear layer, applies ReLU activation, and outputs the final classification through another linear layer.

Temporal-Only Baseline To assess the importance of temporal information alone, we implemented an LSTM model that processes the time series data of fMRI signals. This model takes input sequences of 518 time points, where each time point contains features from 116 brain regions. The LSTM processes this sequence and uses the final hidden state (64 features) to make the classification after passing through a ReLU activation and a final linear layer.

Spatial-Only Baseline Our Static Graph Convolutional Network (GCN) baseline focuses purely on spatial relationships between brain regions. Each node initially has 64 features (all set to 1), and the model applies two GCN layers with edge weights, incorporating ReLU activation between them. The final node features are aggregated through global mean pooling to produce the classification output.

Advanced Spatiotemporal Architectures

Building on these baselines, we developed several architectures that combine spatial and temporal processing in different ways. All these models operate on Graph 1, where nodes represent brain regions and features are time series data.

CNN-GCN Hybrid This architecture first processes temporal information through 1D convolutions before applying spatial graph convolutions:

1. Temporal Processing: — Two 1D convolution layers (window size 16, dilation 2, stride 2) — This creates a learned temporal summary for each brain region — ReLU activation between layers

2. Spatial Processing: — Either a single GCN layer for direct classification — Or two GCN layers with ReLU activation for deeper spatial processing — Global mean pooling for final prediction

class CNNGCN(nn.Module):
    def __init__(self, hidden_channels, out_channels, num_timepoints = 518, window_size=32, stride = 2, dilation = 2):
        super(CNNGCN, self).__init__()
        # 1D Convolution to process time series data for each node
        self.conv1d_1 = nn.Conv1d(
                                in_channels=1,  # Input features per time point
                                out_channels=1,  # Output features per time point
                                kernel_size=window_size,
                                stride = stride,
                                dilation = dilation,
                            )
        self.conv1d_2 = nn.Conv1d(
                                in_channels=1, # Input features per time point
                                out_channels=1, # Output features per time point
                                kernel_size=window_size,
                                stride = stride,
                                dilation = dilation,
                            )

        # Compute the output dimension of the Conv1d layers
        conv_output_dim = (num_timepoints - (window_size - 1) * dilation - 1) // stride + 1
        conv_output_dim = (conv_output_dim - (window_size - 1) * dilation - 1) // stride + 1
        in_channels = int(conv_output_dim)

        # Graph Convolution layer
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, node_features, edge_index, edge_attributes):
        """
        Forward pass for the CNNGCN model.

        Inputs:
          node_features (Tensor): Node features.
          edge_index (Tensor): Edge indices.
          edge_attributes (Tensor): Edge attributes.

        Outputs:
          output (Tensor): Output of the CNNGCN model.
        """
        x = node_features.unsqueeze(1)  # Shape: (num_nodes, 1, time_points)

        # Apply the first 1D convolution (temporal information)
        x = self.conv1d_1(x)
        x = F.relu(x)
        # Apply the second 1D convolution
        x = self.conv1d_2(x)
        x = F.relu(x)

        # Reshape the output of the 1D convolution
        x = x.squeeze(1)  # Shape: (num_nodes, time_points_new)

        # Apply GCNConv layers (spatial information)
        x = self.conv1(x, edge_index, edge_weight=edge_attributes)
        x = F.relu(x)
        x = self.conv2(x, edge_index, edge_weight=edge_attributes)

        # Apply global mean pooling
        x = global_mean_pool(x, None)

        return x

LSTM/RNN with Graph Convolutions We developed several variants combining recurrent neural networks with graph convolutions:

1. LSTM-GCN and RNN-GCN: — Process each node’s time series through shared LSTM/RNN — Extract 32-dimensional features from final hidden states — Apply two GCN layers for spatial processing — Global mean pooling for classification

class TemporalGCN(nn.Module):
    def __init__(self, hidden_channels, out_channels, temporal_layer = 'LSTM'):
        super(TemporalGCN, self).__init__()
        self.hidden_channels = hidden_channels

        if temporal_layer == 'LSTM':
            # LSTM layer to process time series data
            self.temporal_layer = nn.LSTM(input_size = 1, # 1 feature per time point
                                            hidden_size = hidden_channels, # Number of hidden units
                                            batch_first=True)
        elif temporal_layer == 'RNN':
            # RNN layer to process time series data
            self.temporal_layer = nn.RNN(input_size = 1, # 1 feature per time point
                                            hidden_size = hidden_channels, # Number of hidden units
                                            batch_first=True)

        # Graph Convolution layer
        self.conv1 = GCNConv(hidden_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, node_features, edge_index, edge_attributes):
        """
        Forward pass for the TemporalGCN model.

        Inputs:
          node_features (Tensor): Node features.
          edge_index (Tensor): Edge indices.
          edge_attributes (Tensor): Edge attributes.

        Outputs:
          output (Tensor): Output of the TemporalGCN model.
        """
        # Reshape input for batched temporal processing (num_nodes x timeseries x 1)
        x = node_features.unsqueeze(-1)

        # Pass all features through the temporal layer in a single batch
        x, _ = self.temporal_layer(x)

        # Extract the last hidden state for each node
        x = x[:, -1, :]
        x = F.relu(x)

        # Apply GCNConv layers (spatial information)
        x = self.conv1(x, edge_index, edge_weight=edge_attributes)
        x = F.relu(x)
        x = self.conv2(x, edge_index, edge_weight=edge_attributes)

        # Global mean pooling
        x = global_mean_pool(x, None)

        return x

2. LSTM-GAT and RNN-GAT: — Similar to above but uses Graph Attention Networks as replacement of the last GCNConv layer. — Final GAT layer uses 3 attention heads — Allows for dynamic, attention-based weighting of spatial relationships

Temporal Graph Models

We also explored an alternative approach using Graph 2, where nodes represent time points and features are ROI values:

1. TimeStaticGCN: — Nodes contain 116 features (ROI values) — Two GCN layers with ReLU activation — Global mean pooling for final prediction

class TimeStaticGCN(torch.nn.Module):
    def __init__(self, input_channels, hidden_channels, output_channels):
        super(TimeStaticGCN, self).__init__()

        # Define the input channels
        self.input_channels = input_channels

        # Define the first GCN layer
        self.conv1 = GCNConv(input_channels, hidden_channels)
        # Define the second GCN layer
        self.conv2 = GCNConv(hidden_channels, output_channels)

    def forward(self, node_features, edge_index, edge_attributes):
        """
        Forward pass for the TimeStaticGCN model.

        Inputs:
          node_features (Tensor): Node features.
          edge_index (Tensor): Edge indices.
          edge_attributes (Tensor): Edge attributes.

        Outputs:
          output (Tensor): Output of the TimeStaticGCN model.
        """
        x = node_features  # (num_timepoints, num_nodes)
        # Apply the first GCN layer with edge attributes
        x = self.conv1(x, edge_index, edge_weight=edge_attributes)
        x = F.relu(x)  # Apply ReLU activation

        # Apply the second GCN layer with edge attributes
        x = self.conv2(x, edge_index, edge_weight=edge_attributes)

        # Apply global mean pooling
        output = global_mean_pool(x, None)

        return output

Implementation Details

Our implementation followed several key principles to ensure robust training:

1. Training runs for 300 epochs with early stopping (patience of 10)
2. Single spatial layer to prevent overfitting
3. Edge weights incorporated in all GCN operations
4. Careful monitoring of convergence (not all models successfully converged)

This comprehensive set of models allows us to systematically evaluate the importance of different architectural choices in capturing brain state changes induced by caffeine.

def train_model(model, train_loader, optimizer, epochs=300, patience=10):
   """
   Trains a neural network model using early stopping based on training loss.

   This function implements a standard training loop with early stopping to prevent
   overfitting. It updates model parameters using backpropagation and monitors
   the total loss per epoch to determine when to stop training.

   Parameters:
       model: torch.nn.Module
           The neural network model to be trained
       train_loader: torch_geometric.loader.DataLoader
           DataLoader containing the training data in PyTorch Geometric format
       optimizer: torch.optim.Optimizer
           The optimizer to use for training
       epochs: int, optional (default=300)
           Maximum number of training epochs
       patience: int, optional (default=10)
           Number of epochs to wait for loss improvement before early stopping

   Notes:
       - Early stopping occurs when the total loss fails to improve for
         'patience' consecutive epochs
       - The model is trained in batches as provided by the train_loader
       - Each batch is expected to contain:
           * x: Node features
           * edge_index: Graph connectivity
           * edge_attr: Edge weights
           * y: Target labels

   Example:
       >>> model = TimeStaticGCN()
       >>> optimizer = torch.optim.Adam(model.parameters())
       >>> train_model(model, train_loader, optimizer)
   """
   best_loss = float('inf')
   patience_counter = 0

   for epoch in range(epochs):
       model.train()
       total_loss = 0

       # Train on batches
       for data in train_loader:
           optimizer.zero_grad()
           out = model(data.x, data.edge_index, data.edge_attr)
           loss = F.cross_entropy(out, data.y)
           loss.backward()
           optimizer.step()
           total_loss += loss.item()

       # Early stopping logic
       if total_loss < best_loss:
           best_loss = total_loss
           patience_counter = 0
       else:
           patience_counter += 1
           if patience_counter >= patience:
               print(f"Early stopping at epoch {epoch}")
               break

Results and Discussion

Table 1. Results from brain state classification (binary: caffeinated vs. non-caffeinated)

Model Performance Analysis

Our experimental results reveal several interesting patterns in the ability of different architectures to classify caffeinated versus non-caffeinated brain states. Most notably, the temporal graph-based models (TimeStaticGCN and TimeStaticGCNGAT) significantly outperformed other architectures across all metrics.

The TimeStaticGCN model achieved the best overall performance with a balanced accuracy of 0.790 and an impressive AUC score of 0.898. This model also demonstrated strong precision (0.742) and recall (0.800), resulting in an F1 score of 0.763. The GAT variant of this model showed comparable performance, suggesting that the temporal graph representation itself, rather than the specific graph neural network architecture, was the key to success.

In contrast, models operating on the spatial graph (Graph 1) showed more modest performance. The LSTMGAT architecture emerged as the best performer among these models, achieving a balanced accuracy of 0.638 and an AUC of 0.709. However, the simpler baseline MLP model surprisingly outperformed many of the more sophisticated architectures, suggesting potential overfitting issues in the more complex models.

Fig 3. Summarized metrics for model performances.

Interpretability

We also applied GNNExplainer to our best-performing model (TimeStaticGCN) to analyze insights into the neurobiological basis of caffeine’s effects.

For each fold, we use the GNNExplainer algorithm on the trained TimeStaticGCN models to identify the features of the nodes that maximize mutual information. This allows us to quantify the importance of brain ROIs in predicting the caffeine vs. non-caffeine state.

explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='common_attributes',
    edge_mask_type=None,
    model_config=dict(
        mode='multiclass_classification',
        task_level='graph',
        return_type='raw',
    ),
)

test_user_indxs = split_df.query("split == 'test' & fold == @fold")['user_indx'].to_list()
node_importances = []

# For each user on testing set, get the explanation of the node features of the model (brain ROIs)
for user_indx in test_user_indxs:
    print(f'User {user_indx}')
    explanation = explainer(dataset[user_indx].x, dataset[user_indx].edge_index, edge_attr = dataset[user_indx].edge_attr)
    node_importances.append(explanation.node_mask.numpy()[0, :])

The resulting visualization shows regions of high importance (bright yellow) that align well with our understanding of caffeine’s mechanism of action.

Fig. 4. GNNExplainer on TimeStaticGCN

Particularly noteworthy is the high importance assigned to cerebellar regions, which collectively account for approximately 6% of the model’s decision-making process. This finding is especially significant given that the cerebellum is known to have high concentrations of A1 adenosine receptors, the primary target of caffeine in the brain.

Fig. 5. GNNExplainer ROI importance rankings

The top three most important regions identified were:

1. Cerebellum_Crus2_L
2. Cerebellum_Crus1_L
3. Cerebellum_6_R

This biological plausibility in our model’s decision-making process provides additional validation of our approach and suggests that the model has indeed learned meaningful patterns rather than spurious correlations.

Limitations

Our study faced several key constraints that warrant further consideration. While our models showed promising results, we observed a notable tendency toward overfitting, particularly in models operating on spatial brain region graphs (Graph 1), where simpler architectures often outperformed more complex ones. This challenge points to potential limitations in our dataset size and suggests that the spatial relationships between brain regions might be less complex than initially hypothesized. The marked performance disparity between spatial and temporal graph representations also indicates that our initial approach may have overemphasized spatial relationships when temporal patterns were more informative for caffeine state classification. These limitations collectively suggest that future work should carefully balance model complexity with dataset size and potentially prioritize temporal dynamics over spatial relationships in brain state classification tasks.

Conclusions and Future Work

Our study demonstrates the potential of spatiotemporal graph neural networks in classifying caffeine-induced brain states from fMRI data. The temporal graph-based approach (TimeStaticGCN) achieved particularly strong results that significantly outperformed traditional spatial graph representations. The model’s interpretable focus on cerebellar regions also aligns well with known neurobiological mechanisms of caffeine, which validates our approach. This success suggests that temporal relationships between brain activity patterns may be more informative than spatial connections for brain state classification.

Looking ahead, several promising directions emerge for future research. First, investigating hybrid architectures that can better balance spatial and temporal information processing while maintaining model simplicity could yield improved performance. Second, exploring methods to increase model robustness with limited data, such as self-supervised pre-training or data augmentation techniques, could help address the overfitting challenges we observed. Finally, extending this approach to other cognitive states and pharmacological interventions could help establish the generalizability of our findings and potentially lead to broader applications in clinical neuroscience.

The insights gained from this exploratory work offer meaningful contributions to our understanding of both brain dynamics and the practical application of graph neural networks in neuroscience. As we continue to refine these methods, they may eventually serve as valuable tools for studying and monitoring brain states in additional research and clinical settings.

If you wish to see all the code, feel free to check out this Github Repository!