Paper: Provable Repair of Vision Transformers: Last Layer is All You Need (ICLR-2024)

Reviewer comment (Theory):
The authors claim that their method has "provable correctness guarantees" (all images in the repair set will be classified correctly post-repair), which is of course not true because the final accuracy deepens on the repair set and the capacity of the last layer. (Take the ImageNet training set as the repair set for example, can the method guarantee a 100% accuracy? Obviously not.) This incorrect claim seems to come from a misconception regarding the algorithm, as discussed below in the next point.

Author response:
In our presentation of the LPs in Sections 2 and 3, we originally noted that an LP solver can find a solution “if it exists.” However, we also included such language in the introduction to be clear about the soundness vs. completeness, in line with prior provable repair approaches. Our approach is sound; that is, if the approach says that it has repaired the network, then the resulting network is guaranteed to satisfy the repair specification. Similar to all prior provable repair approaches, our approach is not complete; that is, it may not be able to find a network that satisfies the given repair specification. The reason for such incompleteness could either be because such a network does not exist due to, for instance, a contradictory repair specification or because the approach is restricted to modify only certain parameters of the network. (Note that Section 5 of APRNN (Tao et al., 2023) shows that the algorithm of REASSURE (Fu and Li, 2022) is not sound or complete.)

Reviewer comment (Theory):
There's a fundamental algorithmic issue. "Since FTall only terminates once all inputs are classified correctly, it is a provable repair approach." (at the end of section 2), and "If there is no solution to the LP, the loop continues. Otherwise, the repaired Vision Transformer is returned." (at the end of page 5): the authors claim that their algorithm terminates when all the images are correctly classified. The question is: Is it guarantee to terminate?

Author response:
The issue of termination is merely related to how one expresses an approach that is sound but not complete. One alternative to describing such an approach is for the algorithm to always terminate and return either a solution or $\bot$ (representing “I could not find a solution”). In this case, one would prove that if the return value is not $\bot$, then the answer is correct (sound). A second alternative is to only terminate when a solution has been found. In this second case, one would prove that any value returned by the algorithm is correct. However, the algorithm might not terminate.  We just happened to pick this latter approach for describing $FT_{all}$ in Section 2 and $PRoViT_{LP+FT}$ in Algorithm 1.  
&nbsp;&nbsp;&nbsp;&nbsp;In practice, we use a timeout to terminate the algorithm (both $FT_{all}$ and $PRoViT_{LP+FT}$), so in practice we have implemented the first alternative described above. Furthermore, in our experiments $FT_{all}$ timed out while PRoViT always terminated within the timeout limit.  
&nbsp;&nbsp;&nbsp;&nbsp;We have also clarified the presentation of $FT_{all}$ at the end of Section 2. $FT_{all}$ is a baseline, not part of PRoViT. $FT_{all}$ is just regular fine tuning across the entire Vision Transformer with the termination condition of achieving 100% accuracy on the repair set.

Reviewer comment (Evaluation):
Some other claims are also questionable. For example, "PRoViT scales to thousands of images": this depends on the number of classes present in the repair set. For example, if there's only a single class, then this claim becomes uninteresting.

Author response:
None of the prior provable repair approaches are able to repair thousands of images for ImageNet Vision Transformers with low drawdown and high generalization, even when restricted to just one class. APRNN runs out of memory, fine tuning times out or causes high drawdown, etc. Hence, we believe that the claim that PRoViT can repair thousands of images is interesting. For example, it would be incredibly useful to repair thousands of instances of a single class such as a “stop sign” for an autonomous vehicle.

Reviewer comment (Theory):
Are you sure that the optimization problems in Equations (2) and (3) are linear programs? In addition, please avoid using "Theorem" for everything. For example, Theorems 2 and 4 should be stated as remarks, in my opinion. (And perhaps the other ones should be labeled with Proposition instead of Theorem.)

Author response:
Yes, we are certain that Equations 2 and 3 are LPs. Is there notation that we have used that is confusing in our LP formulation? Please let us know so we can rectify that. We have changed the Theorems to Remarks and Propositions.

Reviewer comment (Novelty):
Most importantly, the proposed method has nothing to do with vision transformers. It is largely based on existing works, the optimization problem is the same, only the resolution has been improved by taking into account sparsity. And it also works for any models other than ViT (as long as they has a linear layer at the end, which is the case for all common architectures). It happens to work well for ViTs, but there was absolutely no explanation for this phenomenon, and ViTs were also not the motivation for the design of the method, so the title and the framing of this paper seem to use ViTs only to attract attention.

Author response:
This work was originally inspired by trying to find a provable repair method for Vision Transformers, because one does not already exist. We were pleasantly surprised to find that by only focusing on repairing the last layer, we were able to find a simple, efficient, and effective way to repair Vision Transformers. The novelty of this work is not only taking sparsity into account in the LP formulation of the repair, but also that the combination of fine tuning and LP-based repair of the last layer strikes a balance between drawdown and generalization. Our experiments on ResNet and VGG (Section 4) show that focusing on the last layer does not work as well as it does on Vision Transformers, so we still believe that this approach is ViT-specific. Our hypothesis for this phenomenon is that because the last layer of the Vision Transformer is an output vector containing the “class token” from the final encoder layer, the class information encoded in that vector is more robust than the information that is just the result of convolutions. However, further theoretical justification for this hypothesis is left to future work. We have added this note regarding future work in Section 4.