Dimension Mismatch Issue in B-Spline Calculation within KANLinear Class

mca183: <p>I’m working on a project where I apply Kolmogorov-Arnold Networks (KAN) on the Mixer architecture (<a href="https://github.com/engichang1467/KAN-Mixer" rel="noopener nofollow ugc">KAN-Mixer</a>). You can check the issue <a href="https://github.com/engichang1467/KAN-Mixer/issues/1" rel="noopener nofollow ugc">here</a></p>
<p>This is the code that I tried to run</p>
<pre><code class="lang-auto">X = torch.rand(64, 1, 28, 28) # MNIST (64, 1, 28, 28), CIFAR10 (64, 3, 32, 32)

pe = PatchEmbedding(1, 128, 4) # MNIST (1, 128, 4), CIFAR10 (3, 128, 4) 
t1 = Transformation1()
t2 = Transformation2()
ml = MixerLayer(128, 49, 256, 256) # MNIST (128, 49, 256, 256), CIFAR10 (128, 64, 256, 256)

print(f"Input shape: {X.shape}")
y1 = pe(X)
print(f"Patch Embedding output shape: {y1.shape}")
y2 = t2(y1)
print(f"T2 transformation output shape: {y2.shape}")
y3 = t1(y2)
print(f"T1 transformation output shape: {y3.shape}")
y4 = t1(y3)
print(f"T1 transformation output shape: {y4.shape}")
y5 = ml(y4)    # &lt;--- Problem starts here
print(f"Mixer Layer output shape: {y5.shape}")
</code></pre>
<p>This is the Error message</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[2], line 6
      4 t1 = Transformation1()
      5 t2 = Transformation2()
----&gt; 6 ml = MixerLayer(128, 49, 256, 256) # MNIST (128, 49, 256, 256), CIFAR10 (128, 64, 256, 256)
      8 print(f"Input shape: {X.shape}")
      9 y1 = pe(X)

File ~/model.py:265, in MixerLayer.__init__(self, embedding_dim, num_patch, token_intermediate_dim, channel_intermediate_dim, dropout)
    259 def __init__(self, embedding_dim, num_patch, token_intermediate_dim, channel_intermediate_dim, dropout=0.):
    260     super().__init__()
    262     self.token_mixer = nn.Sequential(
    263         nn.LayerNorm(embedding_dim),
    264         Transformation1(),
--&gt; 265         KAN(num_patch, token_intermediate_dim, dropout),
    266         Transformation1()
    267     )
    269     self.channel_mixer = nn.Sequential(
    270         nn.LayerNorm(embedding_dim),
    271         KAN(embedding_dim, channel_intermediate_dim, dropout),
    272     )

File ~/model.py:223, in KAN.__init__(self, dim, intermediate_dim, dropout)
    220 def __init__(self, dim, intermediate_dim, dropout = 0.):
    221     super().__init__()
    222     self.kan = nn.Sequential(
--&gt; 223         KANLinear(dim, intermediate_dim),
    224         KANLinear(intermediate_dim, dim),
    225     )

File ~/model.py:60, in KANLinear.__init__(self, in_features, out_features, grid_size, spline_order)
     56 self.spline_scaler = nn.Parameter(torch.Tensor(out_features, in_features))
     58 self.base_activation = nn.SiLU()
---&gt; 60 self.reset_parameters()

File ~/model.py:73, in KANLinear.reset_parameters(self)
     71 # Compute the spline weight coefficients from the random noise
     72 grid_points = self.grid.T[self.spline_order : -self.spline_order]
---&gt; 73 spline_coefficients = self.curve2coeff(grid_points, random_noise)
     75 # Copy the computed coefficients to the spline weight tensor
     76 self.spline_weight.data.copy_(spline_coefficients)

File ~/model.py:177, in KANLinear.curve2coeff(self, input_tensor, output_tensor)
    166 """
    167 Compute the coefficients of the curve that interpolates the given points.
    168 
   (...)
    174     torch.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order).
    175 """
    176 # Compute the B-spline bases for the input tensor
--&gt; 177 b_splines_bases = self.b_splines(
    178     input_tensor
    179 )  # (batch_size, input_dim, grid_size + spline_order)
    181 # Transpose the B-spline bases and output tensor for matrix multiplication
    182 transposed_bases = b_splines_bases.transpose(
    183     0, 1
    184 )  # (input_dim, batch_size, grid_size + spline_order)

File ~/model.py:149, in KANLinear.b_splines(self, x)
    147 test = (input_tensor_expanded - expanded_grid[:, :, : -order - 1]) / (expanded_grid[:, :, order:-1] - expanded_grid[:, :, : -order - 1])
    148 print(f"left_term: {test.shape}")
--&gt; 149 test2 = test * bases[:, :, :-1]
    151 left_term = (
    152     (input_tensor_expanded - expanded_grid[:, :, : -order - 1])
    153     / (expanded_grid[:, :, order:-1] - expanded_grid[:, :, : -order - 1])
    154 ) * bases[:, :, :-1]
    156 right_term = (
    157     (expanded_grid[:, :, order + 1 :] - input_tensor_expanded)
    158     / (expanded_grid[:, :, order + 1 :] - expanded_grid[:, :, 1:-order])
    159 ) * bases[:, :, 1:]

RuntimeError: The size of tensor a (10) must match the size of tensor b (11) at non-singleton dimension 2
</code></pre>
<h2><a name="problem-starts-here-1" class="anchor" href="#problem-starts-here-1"></a>Problem starts here</h2>
<pre><code class="lang-auto">class KANLinear(nn.Module):
    # ...
    def b_splines(self, x: torch.Tensor):
        # ...

        # Compute the B-spline bases recursively
        for order in range(1, self.spline_order + 1):
            input_tensor_expanded = input_tensor_expanded.repeat(1, 1, expanded_grid[:, :, : -order - 1].shape[-1])
            test = (input_tensor_expanded - expanded_grid[:, :, : -order - 1]) / (expanded_grid[:, :, order:-1] - expanded_grid[:, :, : -order - 1])
            print(f"left_term: {test.shape}")
            test2 = test * bases[:, :, :-1] # &lt;--- Problem starts here

            left_term = (
                (input_tensor_expanded - expanded_grid[:, :, : -order - 1])
                / (expanded_grid[:, :, order:-1] - expanded_grid[:, :, : -order - 1])
            ) * bases[:, :, :-1]

            # ...

        return bases.contiguous()

    # ...
</code></pre>
<h2><a name="datas-dimension-2" class="anchor" href="#datas-dimension-2"></a>Data’s dimension</h2>
<ul>
<li><code>expanded_grid</code> and <code>bases</code>: <code>torch.Size([6, 49, 12])</code></li>
<li><code>input_tensor_expanded</code>: <code>torch.Size([6, 49])</code></li>
</ul>
<p>After computing the B-Spline bases</p>
<p><code>input_tensor_expanded</code>, <code>expanded_grid</code>, and <code>test</code>: <code>torch.Size([6, 49, 10])</code><br>
<code>bases</code>: <code>torch.Size([6, 49, 12])</code></p>
<h2><a name="problem-3" class="anchor" href="#problem-3"></a>Problem</h2>
<p>I need to ensure that the dimensions of <code>bases</code> match those of <code>input_tensor_expanded</code> and <code>expanded_grid</code> to avoid the error I encountered. Currently, the dimension of <code>test</code> is <code>torch.Size([6, 49, 10])</code> and the dimension of bases is <code>torch.Size([6, 49, 12])</code>. I need to figure out how to adjust the dimension of <code>bases</code> to match the others.</p>

mca183: <p>I have finally resolved the tensor dimension error! Check out the PR <a href="https://github.com/engichang1467/KAN-Mixer/pull/2" rel="noopener nofollow ugc">here</a></p>
<h2><a name="problem-1" class="anchor" href="#problem-1"></a>Problem</h2>
<p>The tensor dimension after the transformation is not compatible with the original KANLinear format. The KANLinear class expects a 2D input, but I want to modify it to accept a 3D tensor.</p>
<h2><a name="solution-2" class="anchor" href="#solution-2"></a>Solution</h2>
<p>I don’t have to change anything with KAN. I just need to modify the forward method to flatten the last two dimensions of the input, and then reshape the output to have the same shape as the input.</p>
<pre data-code-wrap="python"><code class="lang-python">class KANLinear(nn.Module):
    # ... (keep the existing code here)

    def forward(self, x: torch.Tensor):
        # Save the original shape
        original_shape = x.shape

        # Flatten the last two dimensions of the input
        x = x.contiguous().view(-1, self.in_features)

        base_output = F.linear(self.base_activation(x), self.base_weight)
        spline_output = F.linear(
            self.b_splines(x).view(x.size(0), -1),
            self.spline_weight.view(self.out_features, -1),
        )

        # Apply the linear transformation
        output = base_output + spline_output

        # Reshape the output to have the same shape as the input
        output = output.view(*original_shape[:-1], -1)

        return output
</code></pre>