Repository: facebookresearch/DomainBed
Issue: RSC shows poor performance OfficeHome
Root cause: Algorithmic Error

@SirRob1997: Preliminar runs show a very poor performance (39%) at OfficeHome. The paper reports 63%. Could you investigate how to get a higher number in this dataset with default hyper-parameters?

_Originally posted by @lopezpaz in https://github.com/facebookresearch/DomainBed/pull/12#issuecomment-697476956_

Reimplementing RSC following the Featurizer/Classifier abstractions used for the rest of the algorithms.

Edit: [similar discussion for SagNets](https://github.com/facebookresearch/DomainBed/blob/master/domainbed/algorithms.py#L667-L669), where we chose simplicity to make all baselines use the same featurizer and classifier architectures. 

Yeah the algorithm has some flaws, also in their original implementation it seems like the performance varies a lot (also on PACS). Further, they also adapted their implementation with downsampling of middle layers for ResNet which is not mentioned in the paper. I‘m currently running experiments to get performance up, let’s see   how these turn out. 

Minimal implementation, will investigate from here:

```python
class RSC(ERM):
    """
    Representation Self-Challenging,
    Algorithm 1 from: https://arxiv.org/abs/2007.02454
    """
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(RSC, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)
        self.percentile = (1 - hparams['rsc_f_drop_factor']) * 100
        self.num_classes = num_classes

    def update(self, minibatches):
        # inputs
        all_x = torch.cat([x for x,y in minibatches])
        # labels
        all_y = torch.cat([y for _,y in minibatches])
        # one-hot labels
        all_o = torch.nn.functional.one_hot(all_y, self.num_classes)
        # features
        all_f = self.featurizer(all_x) 
        # predictions
        all_p = self.classifier(all_f)

        # Equation (1): compute gradients with respect to representation
        all_g = autograd.grad((all_p * all_o).sum(), all_f)[0]
       
        # Equation (2): compute top-gradient-percentile mask 
        percentiles = np.percentile(all_g.cpu(), self.percentile, axis=1)
        percentiles = torch.Tensor(percentiles)
        percentiles = percentiles.unsqueeze(1).repeat(1, all_g.size(1))
        mask = all_g.lt(percentiles.cuda()).float()

        # Equation (3): mute top-gradient-percentile activations
        all_muted = all_f * mask

        # Equation (4): compute muted predictions
        muted_predictions = self.classifier(all_muted)

        # Equation (5): update
        loss = F.cross_entropy(muted_predictions, all_y)
        self.optimizer.zero_grad()
        loss.backward() 
        self.optimizer.step()

        return {'loss': loss.item()}
```

Will be interesting to see how this will change the performance as this operates on features after average pooling instead of the suggested conv4 output. Also, it disregards the whole Spatial / Channel averaging and applying the masking only for certain samples inside each batch. One thing which might be the cause of the degrading performance in the initial implementation is that I was missing the following lines after [this line](https://github.com/facebookresearch/DomainBed/blob/master/domainbed/algorithms.py#L884).

```
inverted_mask = ~drop_indeces
inverted_mask_indeces = inverted_mask.nonzero()[:,0]
mask[inverted_mask_indeces, :] = 1
```

This recovers the indices and without it, the masking gets dropped for the wrong samples.

The above minimal implementation achieves:

```
-------- Averages, model selection method: training-domain validation set
Algorithm                  ColoredMNIST               OfficeHome                 PACS                       RotatedMNIST               TerraIncognita             VLCS                      
RSC                        51.6 +/- 0.1               65.7 +/- 0.1               84.4 +/- 0.3               97.9 +/- 0.1               47.1 +/- 1.0               77.0 +/- 0.8
```

I will run your modification (adding those three lines) as soon as I get some bandwidth.

Looks quite good actually, didn't expect that. Applying the mask only for top-p percentile samples in the batch should boost performance quite a bit. Their implementation does it like this: https://github.com/DeLightCMU/RSC/blob/master/models/resnet.py#L118

I see, I had missed this option by directly implementing Algorithm 1 as described in Section 3.1. of the paper. Will do!

Algorithm 1 misses some of the key aspects of their implementation. Take a look at Section 3.3 where they introduce Spatial-Wise and Channel-Wise RSC and the stuff for applying them only for certain samples inside each batch. My initial implementation should have these features (with the small fix mentioned above) but can probably be written a lot cleaner.

By not stripping off the avgpool of ResNet, I believe that you inherently always apply Channel-Wise RSC with average pooling on the spatial dimension. 

The minimal version above shows good performance. Next, I will try this version including batch dropping:

```python
class RSC(ERM):
    def __init__(self, input_shape, num_classes, num_domains, hparams):
        super(RSC, self).__init__(input_shape, num_classes, num_domains,
                                   hparams)
        self.drop_f = (1 - hparams['rsc_f_drop_factor']) * 100
        self.drop_b = (1 - hparams['rsc_b_drop_factor']) * 100
        self.num_classes = num_classes

    def update(self, minibatches):
        # inputs
        all_x = torch.cat([x for x, y in minibatches])
        # labels
        all_y = torch.cat([y for _, y in minibatches])
        # one-hot labels
        all_o = torch.nn.functional.one_hot(all_y, self.num_classes)
        # features
        all_f = self.featurizer(all_x)
        # predictions
        all_p = self.classifier(all_f)

        # Equation (1): compute gradients with respect to representation
        all_g = autograd.grad((all_p * all_o).sum(), all_f)[0]

        # Equation (2): compute top-gradient-percentile mask
        percentiles = np.percentile(all_g.cpu(), self.drop_f, axis=1)
        percentiles = torch.Tensor(percentiles)
        percentiles = percentiles.unsqueeze(1).repeat(1, all_g.size(1))
        mask_f = all_g.lt(percentiles.cuda()).float()

        # Equation (3): mute top-gradient-percentile activations
        all_f_muted = all_f * mask_f

        # Equation (4): compute muted predictions
        all_p_muted = self.classifier(all_f_muted)

        # Section 3.3: Batch Percentage
        all_s = F.softmax(all_p, dim=1)
        all_s_muted = F.softmax(all_p_muted, dim=1)
        changes = (all_s * all_o).sum(1) - (all_s_muted * all_o).sum(1)
        percentile = np.percentile(changes.detach().cpu(), self.drop_b)
        mask_b = changes.lt(percentile).float().view(-1, 1)
        mask = torch.logical_or(mask_f, mask_b).float()

        # Equations (3) and (4) again, this time mutting over examples
        all_p_muted_again = self.classifier(all_f * mask)

        # Equation (5): update
        loss = F.cross_entropy(all_p_muted_again, all_y)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item()}
```
I haven't implemented spatial dropping yet, because currently all algorithms in DomainBed perform regularization in a `(n_examples, n_features)` bottleneck followed by a linear classifier.

Thank you for your help!