Adafactor: The Factored Second Moment
Advanced
The problem, from first principles
Adafactor (Shazeer & Stern, 2018) stores the second moment of a weight matrix as a row vector and a column vector instead of a full matrix. The full estimate is rebuilt as an outer product, which costs sublinear memory.
From theory to code
Implement factored_second_moment(row, col), the rank-one reconstruction of the second moment.
Constraints
- Normalize by the total of the row sums.
Hints
Hint 1
Take the outer product of the row and column vectors and divide by the sum of the row sums.
Keywords
research-papersoptimizationmemory
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.