Clarify and Fix Usage of Input Channels in Convolution Layers

[aditya0by0]

Issue Description:

There is currently a mismatch between the intended and actual usage of input and hidden channels in the ResGatedGraphConvNetBase class implementation.

🧩 Problem Description

1. Why are input channels needed if we have n_atom_properties?
   In our current setup, n_atom_properties is effectively acting as the number of input channels for the first convolution layer. This raises the question of whether we need an explicit in_length parameter at all if it’s not used directly. The only place where it appears meaningful is in the code, where it is incorrectly applied to multiple convolution layers.

2. Incorrect channel assignment in previous version:
   In current implementation, all convolution layers used in_length as the input channel size, and only the final layer transitioned to hidden_length. This does not follow conventional practice as used in PyG’s BasicGNN, the correct pattern should be:

   • First layer: input_channels (n_atom_properties in our case) → hidden_channels
   • Hidden layers: hidden_channels → hidden_channels
   • (Optionally) Final layer: hidden_channels → output_channels or similar

3. Off-by-one error in number of convolution layers:
   If the user specifies n_conv_layers=3, the previous code creates 4 layers due to an extra one outside the loop (n + 1 instead of n). This is misleading and could result in unintended architectural depth.

✅ Proposed Fix

• Use n_atom_properties as the number of input channels.
• Remove in_length if it's redundant.
• Ensure the first convolution layer maps from n_atom_properties → hidden_length, and all subsequent layers use hidden_length → hidden_length.
• Fix layer creation logic to respect n_conv_layers accurately (i.e., create exactly n layers).

[aditya0by0]

I trained the model with the change implemented in associated PR. The final results are same, but the number of parameters for GNN is more, which is expected as we are now using hidden_channels (512) in rest of convolution layers rather than input_channels (256). Also this implementation uses exactly 3 convolution layers, in contrast to the current implementation where specifying 3 layers results in 4 being used.

https://wandb.ai/chebai/chebai/runs/dlbf4d5s/logs

[aditya0by0]

https://wandb.ai/chebai/chebai/runs/sn3r6vcm/logs

Trained the model with implemented change in associated PR and now with hidden length as 256. The model almost equal performance with half the number of trainable parameters.

Also this implementation uses exactly 3 convolution layers, in contrast to the current implementation where specifying 3 layers results in 4 being used.

[sfluegel05]

You are right on the input channels and number of layers. Those are not named correctly.

The runs you did are an interesting result. But I would not say that the performance is almost equal. The run with less parameters has a macro-f1 of 0.62 and the run with more parameters has 0.67 - that's a noticeable difference. What is likely happening here is that the model with less parameters converges much slower (in terms of number of epochs). If you compare the train- and validation macro-f1 (as well as the other metrics) for both runs, epoch 200 of the less-parameter run matches epoch ~35-40 of the more-parameter run.

I would suggest you train the less-parameter model for more epochs (600+). Then, either the metrics will converge to the level of the more-parameter run, or (hopefully), the macro-f1 will be a bit higher because the smaller model avoids overfitting.

[aditya0by0]

Title	Epoch	Global Step	Hidden Length	In Length	#Conv Layers	GNN	Linear Layers	Final Layer	Train Loss (Epoch / Step)	Train Macro-F1	Train Micro-F1	Val Loss (Epoch / Step)	Val Macro-F1	Val Micro-F1
Undirected w/o any changes from this PR	199	62,799	512	256	4	1.5M	365K	783K	0.00140 / 0.00218	0.9274	0.9830	0.01783 / 0.00697	0.6635	0.9067
Undirected with changes	599	188,399	256	N/A	3	705K	116K	392K	0.00190 / 0.00082	0.9171	0.9770	0.02097 / 0.00895	0.6406	0.8956
Undirected with changes	199	62,799	256	N/A	3	705K	116K	392K	0.00333 / 0.00389	0.8652	0.9593	0.01248 / 0.00412	0.6480	0.8982
Undirected with changes	199	62,799	256	N/A	4	973K	116K	392K	0.00240 / 0.00316	0.8867	0.9708	0.01461 / 0.01121	0.6436	0.9017
Undirected with changes	599	188,399	256	N/A	4	973K	116K	392K	0.00140 / 0.00098	0.9287	0.9828	0.02156 / 0.00914	0.6490	0.9010

[aditya0by0]

3. Off-by-one error in number of convolution layers:
If the user specifies n_conv_layers=3, the previous code creates 4 layers due to an extra one outside the loop (n + 1 instead of n). This is misleading and could result in unintended architectural depth.

Just a small clarification on the previous note about convolution layers:

If n_conv_layers=3 is specified, the current code actually results in 5 convolution layers (n + 2), not 4.

This happens due to the structure of the loop and the final_conv layer added outside it:

self.convs = torch.nn.ModuleList([])
for i in range(self.n_conv_layers):
    if i == 0:
        self.convs.append(
            tgnn.ResGatedGraphConv(
                self.n_atom_properties,
                self.in_length,
                edge_dim=self.n_bond_properties,
            )
        )
    self.convs.append(
        tgnn.ResGatedGraphConv(
            self.in_length, self.in_length, edge_dim=self.n_bond_properties
        )
    )
self.final_conv = tgnn.ResGatedGraphConv(
    self.in_length, self.hidden_length, edge_dim=self.n_bond_properties
)

🔍 Breakdown:

• The first iteration (i == 0) adds 2 layers (input + one hidden).
• Each subsequent iteration adds 1 layer.
• final_conv adds 1 more layer at the end.

So:

• n_conv_layers = 3 → 2 (from i = 0) + 2 (from i = 1, 2) + 1 (final_conv) = 5 layers total.

🛠️ Suggestion: If the intention is to have exactly n_conv_layers in total, the loop structure and final layer handling may need adjustment.

[aditya0by0]

I have implemented PyG resgated implementation in commit dfdd810 which address this entire issue, without changing the original architecture of the model.
To be 100% sure, I trained gnn with the pyg implementation and with the original implementation and the results are identical down to last decimals.

Run with pyg implementation: https://wandb.ai/chebai/chebai/runs/x6bvjr10/overview
Run with original implementation: https://wandb.ai/chebai/chebai/runs/5ech9dfp/overview

[sfluegel05]

Thats funny. I had never noticed that the first iteration adds two layers, but you are right on that.