SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 5 "Fixed" Rounding via a Convex Program
Require:  $W\in \mathbb{R}^{m\times n}$ $H\in \mathbb{R}^{n\times n}$ $c > 0$ $\rho >0$    
Require: factorization  $m = p_1p_2$ $n = p_3p_4$    
draw  $U_{1}\in \mathbb{R}^{p_{1}\times p_{1}}$  uniformly from the set of orthogonal matrices using seed seed  $(U_{1})$    
draw  $U_{2}\in \mathbb{R}^{p_{2}\times p_{2}}$  uniformly from the set of orthogonal matrices using seed seed  $(U_{2})$    
draw  $U_{3}\in \mathbb{R}^{p_{3}\times p_{3}}$  uniformly from the set of orthogonal matrices using seed seed  $(U_{3})$    
draw  $U_{4}\in \mathbb{R}^{p_{4}\times p_{4}}$  uniformly from the set of orthogonal matrices using seed seed  $(U_{4})$ $W\gets (U_1\otimes U_2)W(U_3\otimes U_4)$ $H\gets (U_3^T\otimes U_4^T)H(U_3\otimes U_4)$ $W\gets \frac{2^b - 1}{2}\left(\frac{W}{\rho} +1\right)$  elementwise   
 $W\gets \mathrm{clamp}(W,\min = 0,\max = 2^b -1))$  elementwise   
use ADMM or some other solver to solve minimize: tr  $(HL^T L)$  over:  $L$  unit upper triangular subject to:  $e_i^T L^T Le_i\leq 1 + c,\forall i\in \{1,\dots ,n\}$  note that when  $c = \infty$ $L^{-1}$  is the factor from the LDL decomposition of  $H$ $\dot{U}\gets L^{-1} - I$    
for  $k\in \{1,\ldots ,n\}$  do  $\hat{W}_k\gets \mathrm{clamp}(\mathcal{Q}(W_k + (W - \hat{W})\hat{U}_k),0,2^b -1)$  ▷ round with LF  $\hat{W}\gets \rho \left(\frac{2\hat{W}}{2^b - 1} -1\right)$ $\hat{W}\gets (U_1^T\otimes U_2^T)\hat{W} (U_3^T\otimes U_4^T)$    
return  $\hat{W}$  encoded as a tuple of the integer rounded values, the scale factor  $\rho$  , and the seeds
Lemma 11. Suppose that for positive definite  $\mu$ -incoherent matrix  $H \in \mathbb{R}^{n \times n}$  and scalar  $c > 0$ ,  $L$  is the solution to the optimization problem
$$
\text {m i n i m i z e :} \operatorname {t r} \left(H L ^ {T} L\right)
$$
over:  $L$  unit upper triangular
$$
\text {s u b j e c t} \colon e _ {i} ^ {T} L ^ {T} L e _ {i} \leq 1 + c, \forall i \in \{1, \dots , n \}.
$$
Then the solution satisfies
$$
\operatorname {t r} \left(H L ^ {T} L\right) = \frac {\mu^ {2}}{n \cdot \operatorname* {m i n} (1 , c)} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2}.
$$
Proof. Let  $\eta \in \mathbb{R}^{1\times n}$  be a random standard Gaussian variable as a row vector, let  $A$  be a matrix, and consider the recurrence relation over  $x_{t}\in \mathbb{R}^{1\times n}$  given by  $x_0 = 0$  and
$$
x _ {t} = x _ {t - 1} - x _ {t - 1} A e _ {i} e _ {i} ^ {T} + \eta e _ {i} e _ {i} ^ {T}
$$
We first note that since  $x_{t}$  is supported only on  $\{1, \dots, t\}$ , if  $M$  denotes the strictly upper triangular mask, this update step is equivalent to
$$
x _ {t} = x _ {t - 1} - x _ {t - 1} (A \odot M) e _ {i} e _ {i} ^ {T} + \eta e _ {i} e _ {i} ^ {T}.
$$
From here, it's fairly easy to see by induction that
$$
x _ {n} = - x _ {n} (A \odot M) + \eta ,
$$
$$
x _ {n} (I + A \odot M) = \eta ,
$$
$$
x _ {n} = \eta (I + A \odot M) ^ {- 1}.
$$
Now, since  $I + A \odot M$  is a unit upper triangular matrix, its inverse is also a unit upper triangular matrix. If we let  $L = (I + A \odot M)^{-1}$ , then  $L$  is a unit upper triangular matrix and
$$
\mathbf {E} \left[ x _ {n} ^ {T} x _ {n} \right] = L ^ {T} L.
$$
We are going to choose  $A$  such that  $L$  is a feasible solution to our optimization problem and has the desired objective. Next, let  $\Sigma_t = \mathbf{E}\left[x_t^T x_t\right]$ , and observe that
$$
\Sigma_ {t} = \left(I - A e _ {i} e _ {i} ^ {T}\right) ^ {T} \Sigma_ {t - 1} \left(I - A e _ {i} e _ {i} ^ {T}\right) + e _ {i} e _ {i} ^ {T}.
$$
Let  $\alpha > 0$  be some constant to be set later, and set  $A = \alpha H^{1/2}$ . Suppose by way of induction that for some constant  $\beta > 0$  to be set later,  $\Sigma_t \preceq \beta H^{-1/2}$ . The base case clearly holds since  $\Sigma_0 = 0$ . For the inductive step,
$$
\begin{array}{l} \Sigma_ {t} \preceq \beta \left(I - \alpha H ^ {1 / 2} e _ {i} e _ {i} ^ {T}\right) ^ {T} H ^ {- 1 / 2} \left(I - \alpha H ^ {1 / 2} e _ {i} e _ {i} ^ {T}\right) + e _ {i} e _ {i} ^ {T} \\ = \beta H ^ {- 1 / 2} - 2 \alpha \beta e _ {i} e _ {i} ^ {T} + \alpha^ {2} \beta e _ {i} e _ {i} ^ {T} H ^ {1 / 2} e _ {i} e _ {i} ^ {T} + e _ {i} e _ {i} ^ {T}. \\ \end{array}
$$
This inductive step will hold if, letting  $h = \max_{i} e_{i}^{T} H^{1/2} e_{i}$ ,
$$
2 \alpha \beta \geq 1 + \alpha^ {2} \beta h
$$
$$
\begin{array}{l} e _ {i} ^ {T} L ^ {T} L e _ {i} = \mathbf {E} \left[ (x _ {n} e _ {i}) ^ {2} \right] \\ = \mathbf {E} \left[ \left(- x _ {i - 1} A e _ {i} + \eta e _ {i}\right) ^ {2} \right] \\ = \mathbf {E} \left[ (- x _ {i - 1} A e _ {i}) ^ {2} \right] + 1 \\ = e _ {i} ^ {T} A ^ {T} \Sigma_ {i - 1} A e _ {i} + 1 \\ = \alpha^ {2} e _ {i} ^ {T} H ^ {1 / 2} \Sigma_ {i - 1} H ^ {1 / 2} e _ {i} + 1 \\ \leq \alpha^ {2} \beta e _ {i} ^ {T} H ^ {1 / 2} H ^ {- 1 / 2} H ^ {1 / 2} e _ {i} + 1 \\ \leq \alpha^ {2} \beta e _ {i} ^ {T} H ^ {1 / 2} e _ {i} + 1. \\ \end{array}
$$
So the constraint of our optimization problem will be satisfied if
$$
\alpha^ {2} \beta h \leq c.
$$
To satisfy these constraints, set  $\beta = \max (h,h / c)$  and  $\alpha = \beta^{-1}$ . Then
$$
2 \max  (h, h / c) ^ {- 1} \cdot \max  (h, h / c) \geq 1 + \max  (h, h / c) ^ {- 2} \cdot \max  (h, h / c) \cdot h,
$$
$$
\max  (h, h / c) ^ {- 2} \cdot \max  (h, h / c) \cdot h \leq c.
$$
Also, the objective will be bounded by
$$
\operatorname {t r} \left(H L ^ {T} L\right) = \operatorname {t r} \left(H \Sigma_ {n}\right) \leq \beta \operatorname {t r} \left(H ^ {1 / 2}\right) = \max  \left(1, e ^ {- 1}\right) \cdot h \cdot \operatorname {t r} \left(H ^ {1 / 2}\right).
$$
Now, applying incoherence to bound  $h$ , where  $H = U\Lambda U^T$  is the eigendecomposition of  $H$ ,
$$
e _ {i} ^ {T} H ^ {1 / 2} e _ {i} = \sum_ {j = 1} ^ {n} \lambda_ {j} ^ {1 / 2} (e _ {i} ^ {T} U e _ {j}) ^ {2} \leq \sum_ {j = 1} ^ {n} \lambda_ {j} ^ {1 / 2} \frac {\mu^ {2}}{n} = \frac {\mu^ {2}}{n} \operatorname {t r} \left(H ^ {1 / 2}\right).
$$
So this yields a whole bound of
Lemma 12. Suppose that we quantize the row vector  $w \in \mathbb{R}^{1 \times n}$  using  $L$  the solution to the optimization problem
$$
s u b j e c t \quad t o: e _ {i} ^ {T} L ^ {T} L e _ {i} \leq 1 + c, \forall i \in \{1, \dots , n \}
$$
$$
\hat {w} = \mathcal {Q} _ {\mathrm {s t o c h}} \left(w - (\hat {w} - w) (L ^ {- 1} - I)\right),
$$
where  $\mathcal{Q}_{\mathrm{stoch}}$  denotes elementwise unbiased stochastic rounding. Then for any  $u \in \mathbb{R}^n$  and any  $\delta > 0$
$$
\mathbf {P} \left(\left| (\hat {w} - w) u \right| \geq \| L u \| \sqrt {\frac {1}{2} \log \left(\frac {2}{\delta}\right)}\right) \leq \delta .
$$
$$
\mathbf {P} \left(\left| (\hat {w} - w) (L ^ {- 1} - I) e _ {i} \right| \geq \sqrt {\frac {c}{2} \log \left(\frac {2}{\delta}\right)}\right) \leq \delta .
$$
Proof. Let  $\eta$  be the error of stochastic rounding, and observe that each entry is, conditioned on earlier steps, zero mean and supported on two values that differ by 1. Also observe that
$$
\hat {w} = \big (w - (\hat {w} - w) (L ^ {- 1} - I) \big) + \eta ,
$$
$$
\hat {w} - w = \eta L
$$
$$
\mathbf {E} \left[ \exp \left((\hat {w} - w) u\right) \right] = \mathbf {E} \left[ \exp \left(\eta L u\right) \right].
$$
From a repeated application of Hoeffding's lemma, we get
$$
\mathbf {E} \left[ \exp \left((\hat {w} - w) u\right) \right] \leq \exp \left(\frac {1}{8} \| L u \| ^ {2}\right).
$$
Setting  $u\mapsto \gamma u$  for  $\gamma >0$
$$
\mathbf {E} \left[ \exp \left(\gamma (\hat {w} - w) u\right) \right] \leq \exp \left(\frac {\gamma^ {2}}{8} \| L u \| ^ {2}\right).
$$
And by Markov's inequality,
$$
\mathbf {P} \left(\exp \left(\gamma (\hat {w} - w) u\right) \geq \exp (\gamma R)\right) \leq \exp (- \gamma R) \exp \left(\frac {\gamma^ {2}}{8} \| L u \| ^ {2}\right),
$$
$$
\mathbf {P} \left((\hat {w} - w) u \geq R\right) \leq \exp \left(- \gamma R + \frac {\gamma^ {2}}{8} \| L u \| ^ {2}\right).
$$
Minimizing the right side over  $\gamma$  yields  $\gamma = 4R\|Lu\|^{-2}$  and
$$
\mathbf {P} \left((\hat {w} - w) u \geq R\right) \leq \exp \left(- 2 R ^ {2} \| L u \| ^ {- 2}\right).
$$
By a union bound,
$$
\mathbf {P} \left(\left| (\hat {w} - w) u \right| \geq R\right) \leq 2 \exp \left(- 2 R ^ {2} \| L u \| ^ {- 2}\right).
$$
Now setting the right side equal to  $\delta$
This is what we wanted to show. The second statement follows from the fact that
$$
\left\| L (L ^ {- 1} - I) e _ {i} \right\| ^ {2} = \left\| e _ {i} - L e _ {i} \right\| ^ {2} = e _ {i} ^ {T} e _ {i} - e _ {i} ^ {T} L e _ {i} - e _ {i} ^ {T} L ^ {T} e _ {i} + e _ {i} ^ {T} L ^ {T} L e _ {i} \leq 1 - 1 - 1 + (1 + c) = c.
$$
Lemma 13. Suppose that we quantize the row vector  $w \in \mathbb{R}^{1 \times n}$  using  $L$  the solution to the optimization problem
$$
m i n i m i z e: \operatorname {t r} \left(H L ^ {T} L\right)
$$
$$
\hat {w} = \mathcal {Q} _ {\text {s t o c h}} (w - (\hat {w} - w) (L ^ {- 1} - I)),
$$
where  $\mathcal{Q}_{\mathrm{stoch}}$  denotes elementwise unbiased stochastic rounding. Suppose that for some integer  $b$ ,  $1 \leq w_{ij} \leq 2^b - 2$ . Then if we set
$$
c = 2 \left(\log \left(\frac {4 m n}{\delta}\right)\right) ^ {- 1},
$$
then with probability at least  $1 - \delta$ ,  $0 \leq \hat{w}_{ij} \leq 2^b - 1$  and
$$
\operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {\mu^ {2} m}{4 n} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \left(\log \left(\frac {4 m n}{\delta}\right) ^ {2}\right).
$$
Proof. First, from the previous lemmas, if  $Ue_{i}$  is the  $i$ th eigenvector of  $H$ , with eigenvalue  $\lambda_{i}$  since
$$
\mathbf {P} \left(\lambda_ {i} \left(e _ {j} ^ {T} (\hat {w} - w) U e _ {i}\right) ^ {2} \geq \lambda_ {i} \| L U e _ {i} \| ^ {2} \cdot \frac {1}{2} \log \left(\frac {2}{\delta}\right)\right) \leq \delta .
$$
By the union bound,
$$
\mathbf {P} \left(\exists i, j, \lambda_ {i} (e _ {j} ^ {T} (\hat {w} - w) U e _ {i}) ^ {2} \geq \lambda_ {i} \| L U e _ {i} \| ^ {2} \cdot \frac {1}{2} \log \left(\frac {2 m n}{\delta}\right)\right) \leq \delta .
$$
$$
\mathbf {P} \left(\sum_ {i, j} \lambda_ {i} (e _ {j} ^ {T} (\hat {w} - w) U e _ {i}) ^ {2} \geq \sum_ {i, j} \lambda_ {i} \| L U e _ {i} \| ^ {2} \cdot \frac {1}{2} \log \left(\frac {2 m n}{\delta}\right)\right) \leq \delta ,
$$
$$
\mathbf {P} \left(\operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \geq m \operatorname {t r} \left(H L ^ {T} L\right) \cdot \frac {1}{2} \log \left(\frac {2 m n}{\delta}\right)\right) \leq \delta .
$$
Now applying the other lemma,
$$
\mathbf {P} \left(\mathrm {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \geq \frac {\mu^ {2} m}{2 n \cdot \min (1 , c)} \mathrm {t r} \left(H ^ {1 / 2}\right) ^ {2} \log \left(\frac {2 m n}{\delta}\right)\right) \leq \delta .
$$
And substituting  $\delta \mapsto \delta / 2$ ,
$$
\mathbf {P} \left(\mathrm {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \geq \frac {\mu^ {2} m}{2 n \cdot \min (1 , c)} \mathrm {t r} \left(H ^ {1 / 2}\right) ^ {2} \log \left(\frac {4 m n}{\delta}\right)\right) \leq \frac {\delta}{2}.
$$
On the other hand, again by a union bound from the previous lemma,
$$
\mathbf {P} \left(\exists i, j, \left| e _ {j} ^ {T} (\hat {w} - w) (L ^ {- 1} - I) e _ {i} \right| \geq \sqrt {\frac {c}{2} \log \left(\frac {4 m n}{\delta}\right)}\right) \leq \frac {\delta}{2}.
$$
$$
c = 2 \left(\log \left(\frac {4 m n}{\delta}\right)\right) ^ {- 1}
$$
$$
\mathbf {P} \left(\exists i, j, \left| e _ {j} ^ {T} (\hat {w} - w) (L ^ {- 1} - I) e _ {i} \right| \geq 1\right) \leq \frac {\delta}{2}.
$$
And so by another union bound, the probability that
$$
\operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {\mu^ {2} m}{4 n} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \left(\log \left(\frac {4 m n}{\delta}\right)\right) ^ {2}
$$
$$
\max _ {i, j} \left| e _ {j} ^ {T} (\hat {w} - w) (L ^ {- 1} - I) e _ {i} \right| \leq 1
$$
is no less than  $1 - \delta$ . It's clear that if this second inequality holds, the value we pass in to the stochastic quantizer will be in range, and thus so will the output. This proves what we want.
Theorem 14. Suppose that we are given an input matrix  $w$  with bounded maximum entry magnitude  $\| w \|_{\infty}$  and we want to quantize it using  $b$  bits. Suppose that we first re-scale the entries of  $w$  by mapping
$$
w _ {i j} \mapsto \frac {2 ^ {b} - 3}{2} \left(\frac {w _ {i j}}{\| w \| _ {\infty}} + 1\right) + 1;
$$
this guarantees that  $1 \leq w_{ij} \leq 2^b - 2$ . Then, suppose we quantize using the procedure described in the previous lemma. Finally, we undo the scaling. Then then with probability at least  $1 - \delta$ , all the quantized weights will be in range (no overflow or need for clipping) and
$$
\operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {\mu^ {2} m}{n (2 ^ {b} - 3) ^ {2}} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \| w \| _ {\infty} ^ {2} \left(\log \left(\frac {4 m n}{\delta}\right) ^ {2}\right).
$$
Proof. This is a straightforward consequence of the previous lemma.
Theorem 15. Suppose that we are given an input matrix  $w$  with bounded  $\| w \|_F$  and we want to quantize it using  $b$  bits. Suppose that we first multiply by two-factor orthogonal matrices, and then we re-scale the entries of  $w$  by mapping
$$
w _ {i j} \mapsto \frac {2 ^ {b} - 3}{2} \left(\frac {w _ {i j}}{\| w \| _ {F} \sqrt {\frac {A ^ {2}}{m n} \log \left(\frac {2 C m n}{\delta}\right) ^ {2}}} + 1\right) + 1;
$$
this guarantees that  $1 \leq w_{ij} \leq 2^b - 2$ . Then, suppose we quantize using the procedure described in the previous lemma. Finally, we undo the scaling and multiplication. Then then with probability at least  $1 - \delta$ , all the quantized weights will be in range (no overflow or need for clipping) and
$$
\begin{array}{l} \operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {A ^ {4}}{n ^ {2} (2 ^ {b} - 3) ^ {2}} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \| w \| _ {F} ^ {2} \left(\log \left(\frac {1 2 C m n ^ {2}}{\delta}\right)\right) ^ {6} \\ = \tilde {\mathcal {O}} \left(\frac {1}{n ^ {2} 4 ^ {b}} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \| w \| _ {F} ^ {2}\right). \\ \end{array}
$$
Proof. It is a straightforward consequence of Lemma 5, that unioning over the three bounds on the infinity norm of  $w$ , the incoherence of  $H$ , and the stochastic rounding, with probability at least  $1 - 3\delta$
$$
\begin{array}{l} \operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {m}{n (2 ^ {b} - 3) ^ {2}} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \| w \| _ {F} ^ {2} \left(\log \left(\frac {4 m n}{\delta}\right) ^ {2}\right) \\ \cdot A ^ {2} \log \left(\frac {2 C n ^ {2}}{\delta}\right) ^ {2} \cdot \frac {A ^ {2}}{m n} \log \left(\frac {2 C n}{\delta}\right) ^ {2}. \\ \end{array}
$$
Substituting  $\delta \mapsto \delta /3$
$$
\begin{array}{l} \operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {1}{n (2 ^ {b} - 3) ^ {2}} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \| w \| _ {F} ^ {2} \left(\log \left(\frac {1 2 m n}{\delta}\right) ^ {2}\right) \\ \cdot A ^ {2} \log \left(\frac {6 C n ^ {2}}{\delta}\right) ^ {2} \cdot \frac {A ^ {2}}{n} \log \left(\frac {6 C n}{\delta}\right) ^ {2}. \\ \end{array}
$$
And this right side is clearly less than
$$
\operatorname {t r} \left((\hat {w} - w) H (\hat {w} - w) ^ {T}\right) \leq \frac {A ^ {4}}{n ^ {2} (2 ^ {b} - 3) ^ {2}} \operatorname {t r} \left(H ^ {1 / 2}\right) ^ {2} \| w \| _ {F} ^ {2} \left(\log \left(\frac {1 2 C m n ^ {2}}{\delta}\right)\right) ^ {6}.
$$
Proof. This follows directly from the previous theorem, which says explicitly what the hyperparameter assignments should be.  $\square$
References for the Appendix
[1] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Optq: Accurate quantization for generative pre-trained transformers. In International Conference on Learning Representations, 2023.
[2] Steve Lalley. Lecture notes on measure-theoretic probability 2. http://galton.uchicago.edu/~lalley/Courses/383/Concentration.pdf, 2018.
[3] Markus Nagel, Rana Ali Amjad, Mart Van Baalen, Christos Louizos, and Tijmen Blankevoort. Up or down? adaptive rounding for post-training quantization. In International Conference on Machine Learning, pages 7197-7206. PMLR, 2020.