Givens orthogonal layer#57
Conversation
|
I somehow broke @kevinchern's tests, what the hell... |
@VolodyaCO which tests? I'm seeing |
I forgot to update my tests to float64 precision. Now that I've done it, it's weird that all of the current failing tests are failing on File "/Users/distiller/project/tests/test_nn.py", line 144, in test_LinearBlock
self.assertTrue(model_probably_good(model, (din,), (dout,))) |
Ahhhhhh. OK Theo also flagged this at #50 . It's a poorly-written test.. you can ignore it. |
| angles, blocks, Ufwd_saved = ctx.saved_tensors | ||
| Ufwd = Ufwd_saved.clone() | ||
| M = grad_output.t() # dL/dU, i.e., grad_output is of shape (n, n) | ||
| n = M.size(1) | ||
| block_size = n // 2 | ||
| A = torch.zeros((block_size, n), device=angles.device, dtype=angles.dtype) |
There was a problem hiding this comment.
Same here re lowercase for Ufwd, M, and A. Avoids incorrect colour highlighting in themes.
There was a problem hiding this comment.
Hmmm, I didn't read this about the incorrect colour highlighting before I made my previous comment. I still think that it is easier to read the algorithm alongside the code if the use of lower/upper case match. For example, lower case m is usually used for an integer variable, not a tensor.
VolodyaCO
force-pushed
the
givens-rotation
branch
from
f18b476 to
46062e0
Compare
|
After a bit of git wrangling, I was able to clean my whole mess of merge commits 😆. |
anahitamansouri
left a comment
anahitamansouri
left a comment
There was a problem hiding this comment.
This is a nice PR Vlad. It took me a while to go over the paper and this PR :) The only thing is the tests that are failing. Thanks for the great work.
VolodyaCO
force-pushed
the
givens-rotation
branch
from
46062e0 to
7f76571
Compare
kevinchern
left a comment
kevinchern
left a comment
There was a problem hiding this comment.
Did a quick pass to provide some feedback before taking some time to take a deep dive into the paper.
| for _ in range(n - 1): | ||
| pairs = circle_method(sequence) | ||
| if is_odd: | ||
| # Remove pairs involving the dummy dimension: |
There was a problem hiding this comment.
| # Remove pairs involving the dummy dimension: | |
| # Remove pairs involving the dummy dimension |
There was a problem hiding this comment.
I was gonna ask why remove the colon?
VolodyaCO
force-pushed
the
givens-rotation
branch
3 times, most recently
from
1146ebb to
df400fb
Compare
|
@VolodyaCO can you rebase on main? |
VolodyaCO
force-pushed
the
givens-rotation
branch
from
df400fb to
8a00932
Compare
VolodyaCO
force-pushed
the
givens-rotation
branch
3 times, most recently
from
6587b89 to
97c2b8e
Compare
VolodyaCO
force-pushed
the
givens-rotation
branch
from
97c2b8e to
36a15c8
Compare
VolodyaCO
left a comment
VolodyaCO
left a comment
There was a problem hiding this comment.
@kevinchern I have addressed all comments. Should be ready to merge now.
| __all__ = ["GivensRotation"] | ||
|
|
||
|
|
||
| class _RoundRobinGivens(torch.autograd.Function): |
There was a problem hiding this comment.
The documentation is quite thorough for a hidden class. Would this perhaps make sense moving to a nn.functions namespace and removing the underscore? @kevinchern
| # Initialize U^fwd from forward pass output U. Mathematically, U^fwd represents U^{1:k-1} | ||
| # at block k, defined in equation (11). It is post-multiplied by G_bk^T at each block | ||
| # iteration to "remove the effect of the block's rotations" (Section 4.1, paragraph before | ||
| # equation (12)). This corresponds to the update from U^{1:k} back to U^{1:k-1}. |
There was a problem hiding this comment.
I haven't looked at the reference link, but these comments seem a bit messy (for lack of a better term). If the algorithm in the link describes the process in a clear way, I'd perhaps shorten these comments a bit, although references to equations are always good.
Also, separation by empty lines always makes things look a bit tidier.
There was a problem hiding this comment.
I have added separation by empty lines to make things look tidier. I added these inline comments at @kevinchern 's request. I do agree that it is enough for the interested reader to read the paper. I also think it's helpful to have the inline comments too.
There was a problem hiding this comment.
shorten these comments a bit
I second this opinion
kevin's request
To clarify my suggestion: it will be very helpful for future maintenance to include references to "where" these equations appear. The idea is for the developer to be able to easily reference the material and not necessarily understand the method from comments.
There was a problem hiding this comment.
I see your point. Comments have now been shortened.
kevinchern
left a comment
kevinchern
left a comment
There was a problem hiding this comment.
Unit tests hinge on NaiveGivensRotationLayer being correct but there aren't correctness tests for it. It checks for serial versus parallel implementations.
We need a unit test to check the forward and backward methods are correct, e.g., work out the expected forward/backward values a small input (n=3?)
| Sequentially applies all Givens rotations to implement an orthogonal transformation in an order | ||
| provided by blocks, which are of shape (n_blocks, n/2, 2), and where usually each block contains | ||
| pairs of indices such that no index appears more than once in a block. However, this | ||
| implementation does not rely on that assumption, so that indeces can appear multiple times in a |
| from tests.helper_functions import model_probably_good | ||
|
|
||
|
|
||
| class NaiveGivensRotationLayer(nn.Module): |
There was a problem hiding this comment.
Why is this class in tests?
My understanding is its purpose is to test the parallel implementation, which is error-prone in implementation.
If that's correct, I would note its purpose in the docstring.
There was a problem hiding this comment.
I updated the docstring to note its purpose.
|
|
||
| import torch | ||
| import torch.nn as nn | ||
| from einops import einsum |
There was a problem hiding this comment.
why use einops.einsum instead of torch.einsum?
There was a problem hiding this comment.
did a quick google search and it suggests torch.einsum has the advantage of being torch-native, so it benefits from torch.compile whereas einops is more general but less optimized for torch.
einops is generally useful and i forget if it's used elsewhere in the package---if you revert this to torch.einsum (I think we should), remove einops from requirements.txt
There was a problem hiding this comment.
I vaguely recall Theo commented on einops as a requirement when this PR first came up, and I mentioned I also use it (einops.rearrange). Can't seem to find that discussion...
There was a problem hiding this comment.
https://github.com/arogozhnikov/einops/wiki/Using-torch.compile-with-einops
They are compatible. einops is used in tests, repeating and rearranging and indeed better written in einops than pytorch. It's just a clearer syntax—easier to read.
| # Test orthogonality: | ||
| I = torch.eye(n, dtype=U_parallel.dtype) | ||
| UU_T = U_parallel @ U_parallel.T | ||
| self.assertTrue(torch.allclose(I, UU_T, atol=1e-6)) |
| U_parallel = layer._create_rotation_matrix() | ||
|
|
||
| # Test that the matrices are close | ||
| self.assertTrue(torch.allclose(U_naive, U_parallel, atol=1e-6)) |
| U_parallel = layer._create_rotation_matrix() | ||
|
|
||
| # Test that the matrices are close | ||
| self.assertTrue(torch.allclose(U_naive, U_parallel, atol=1e-6)) |
There was a problem hiding this comment.
For what it's worth, if I understood correctly, this test is redundant given the next test (identity test).
Good to keep though!
There was a problem hiding this comment.
They are not equivalent tests. One tests orthogonality, the other tests that both non-parallel and parallel implementations provide the same matrix.
| angles_in_block = angles[idx_block + b * block_size] # shape (n/2,) | ||
| c = torch.cos(angles_in_block).unsqueeze(0) | ||
| s = torch.sin(angles_in_block).unsqueeze(0) | ||
| i_idx = block[:, 0] | ||
| j_idx = block[:, 1] | ||
| r_i = c * U[:, i_idx] + s * U[:, j_idx] | ||
| r_j = -s * U[:, i_idx] + c * U[:, j_idx] | ||
| U[:, i_idx] = r_i | ||
| U[:, j_idx] = r_j |
There was a problem hiding this comment.
There was a problem hiding this comment.
I am stating in the docstring that this refers to Algorithm 2, which is where all of these instructions are written. I think this is sufficient for the reader to establish a way of reading the paper that corresponds to the written code.
| # Initialize U^fwd from forward pass output U. Mathematically, U^fwd represents U^{1:k-1} | ||
| # at block k, defined in equation (11). It is post-multiplied by G_bk^T at each block | ||
| # iteration to "remove the effect of the block's rotations" (Section 4.1, paragraph before | ||
| # equation (12)). This corresponds to the update from U^{1:k} back to U^{1:k-1}. |
There was a problem hiding this comment.
shorten these comments a bit
I second this opinion
kevin's request
To clarify my suggestion: it will be very helpful for future maintenance to include references to "where" these equations appear. The idea is for the developer to be able to easily reference the material and not necessarily understand the method from comments.
| # M represents the product U^bck @ Γ^T. From equation (15): M ≡ U^bck @ Γ^T, where | ||
| # U^bck = U^{k:n-1} (the product of remaining blocks, defined in equation (11)). | ||
| # At each block k, M is pre-multiplied by G_bk to advance from U^{k+1:n-1} to U^{k:n-1}. | ||
| M = grad_output.t() |
There was a problem hiding this comment.
Any reason to use tensor.t() instead of tensor.T? I suggest picking one and using them consistently throughout.
There was a problem hiding this comment.
Not really. Going for .T
| self.n = n | ||
| self.n_angles = n * (n - 1) // 2 | ||
| self.angles = nn.Parameter(torch.randn(self.n_angles)) | ||
| blocks_edges = self._get_blocks_edges(n) | ||
| self.register_buffer("blocks", blocks_edges) | ||
| if bias: | ||
| self.bias = nn.Parameter(torch.zeros(n)) | ||
| else: | ||
| self.register_parameter("bias", None) |
There was a problem hiding this comment.
Should some of these be made private and, if needed, accessed via @property?
cc @thisac
There was a problem hiding this comment.
Could do that, to signal that they shouldn't be changed by user. 👍
Address orthogonal module PR feedback.
VolodyaCO
force-pushed
the
givens-rotation
branch
from
9f31f15 to
c0cd080
Compare
4 participants
This PR adds an orthogonal layer given by Givens rotations, using the parallel algorithm described by Firas in https://arxiv.org/abs/2106.00003, which gives a forward complexity of O(n) and backward complexity of O(n log(n)), even though there are O(n^2) rotations.
This PR still is in draft. I wrote it for even n. Probably some more unit tests are to be done, but I am quite lazy (will do it after all math is checked for odd n).