Linear layers constitutes almost half of the trainable parameters

[aditya0by0]

While inspecting the parameter distribution of the model, I noticed that ~43% of the total trainable parameters are concentrated in the linear layers alone:

Component	Parameters
gnn	1.5M
linear_layers	365K
final_layer	783K
Total	2.648M

I was wondering about a few architectural questions:

• Is this heavy reliance on linear layers intentional, or could it be simplified? It may overshadow the representational power of the GNN backbone.
• Could we experiment with a single linear head, as commonly done in standard models like those in PyG?
• Alternatively, would it make sense to shift more capacity back into the GNN module, where the inductive bias is typically stronger?

Metrics:

• train_macro-f1: 0.960
• train_micro-f1: 0.990
• val_macro-f1: 0.699
• val_micro-f1: 0.911

Metric	Train	Validation	Gap
Macro F1	0.960	0.699	0.261
Micro F1	0.990	0.911	0.079

[aditya0by0]

Just to add some context on the GNN vs. linear layer parameter balance:

The current model uses 4 convolution layers (n_conv_layers = 3), which is fairly standard for molecular property prediction. In most cases, GNNs for small molecules use 3–5 convolution layers. Going beyond 5–6 layers often leads to over-smoothing, where node embeddings become indistinguishable and the model loses expressiveness.

Because of this natural upper bound on the number of GNN layers, there's a limit to how much capacity we can assign to the GNN. When we fix the number of convolution layers (e.g. to 3), the GNN’s parameter share remains modest, especially with fixed hidden dimensions. That’s why the linear layers start to dominate the total parameter count — often approaching or exceeding 50%.

This makes a case for either simplifying/reducing the linear stack or increasing the hidden dimensionality in GNN layers.

[aditya0by0]

🔍 Layer Configuration Summary

• Input dimension: 158 (node features)

• Convolution layers:

  • Layer 1: 158 → 256
  • Layer 2: 256 → 256
  • Layer 3: 256 → 256
  • Final GNN output layer: 256 → 512

• Linear layers:

  • First: 512 + 200 → 712 ( with molecule-level properties)
  • Intermediate: 712 → 512
  • More dense layers: 512 → 512
  • Final classifier head: → ~1056 (output label size)

[aditya0by0]

GNN with Single Convolution Layer (Undirected)

For experiment, I trained a undirected GNN with only one convolution layer (n_conv_layers = 0).

• input_size = n_atom_properties
• output_size = hidden_dim

📈 Metrics @ Epoch 199:

Metric	Value
train_loss_epoch	0.0070
train_loss_step	0.0112
train_macro-f1	0.7920
train_micro-f1	0.9094
val_loss_epoch	0.0096
val_loss_step	0.0059
val_macro-f1	0.6266
val_micro-f1	0.8833
global_step	62,799

Parameter distribution (n_conv_layers = 0):

Component	Parameters	Trainable
gnn	336K	✅
linear_layers	365K	✅
final_layer	783K	✅
Total	1.484M	✅

➡️ GNN accounts for ~22.6% of total trainable parameters.

https://wandb.ai/chebai/chebai/runs/h2u0aupe/overview

[aditya0by0]

An option to set the number of linear layers is provided in #2 for this issue.