2025 USA-NA-AIO Round 1, Problem 2, Part 9

AI Olympiads 2025 USA-NA-AIO Round 1
1

Part 9 (5 points, coding task)

This question follows Part 6.

In this part, you are asked to program with PyTorch, not NumPy.

Build a deep neural network class named as Symmetric_Linear_Model that meets the modifications imposed in Part 6.

2 
### WRITE YOUR SOLUTION HERE ###

class Symmetric_Linear_Model(nn.Module):
    def __init__(self, in_features, hidden_features):
        super().__init__()
        self.in_features = in_features
        self.hidden_features = hidden_features
        self.linear0 = nn.Linear(in_features, hidden_features, bias = False)

    def forward(self, x):
        x = self.linear0(x)
        x = torch.sum(self.linear0.weight * x.reshape(*x.shape, 1), dim = -2)
        return x


""" END OF THIS PART """
3

This is my solution, does it also work?

class Symmetric_Linear_Model(nn.Module):
  def __init__(self, in_features, hidden_features):
    super().__init__()
    self.weights = nn.Parameter(torch.randn(hidden_features, in_features))

  def forward(self, x):
    return self.weights.T @ self.weights @ x
4

It looks fine, maybe x should premultiply self.weights.T@self.weights? I think that allows it to handle batches correctly.

5

You are right! To support batch it should be x @ self.weights.T @ self.weights. Thanks!