Quantize a Factorized Linear Layer

snakers41: <p>Hi,</p>
<p>I have created a small layer to make my networks smaller as a drop-in replacement for a <code>Linear</code> layer:</p>
<pre><code class="lang-auto">class FactorizedLinear(nn.Module):
    def __init__(self,
                 or_linear,
                 dim_ratio=1.0,
                 random_init=False):
        super().__init__()
        self.bias = nn.parameter.Parameter(or_linear.bias.data)
        if random_init:
            u, vh = self.random_init(or_linear.weight.data, dim_ratio=dim_ratio)
            print(f'Doing zero init of tensor {or_linear.weight.shape}, U: {u.shape}, Vh: {vh.shape}')
        else:
            u, vh = self.spectral_init(or_linear.weight.data, dim_ratio=dim_ratio)
            print(f'Doing SVD of tensor {or_linear.weight.shape}, U: {u.shape}, Vh: {vh.shape}')
        self.u = nn.parameter.Parameter(u)
        self.vh = nn.parameter.Parameter(vh)
        self.dim_ratio = dim_ratio
        self.in_features = u.size(0)
        self.out_features = vh.size(1)

    @staticmethod
    @torch.jit.ignore
    def spectral_init(m,
                      dim_ratio=1):
        u, s, vh = torch.linalg.svd(m, full_matrices=False)
        u = u @ torch.diag(torch.sqrt(s))
        vh = torch.diag(torch.sqrt(s)) @ vh
        if dim_ratio &lt; 1:
            dims = int(u.size(1) * dim_ratio)
            u = u[:, :dims]
            vh = vh[:dims, :]
            s_share = s[:dims].sum() / s.sum() * 100
            print(f'SVD eigenvalue share {s_share:.2f}%')
        return u, vh

    @staticmethod
    @torch.jit.ignore
    def random_init(m,
                    dim_ratio=1):
        bottleneck = int(m.size(1) * dim_ratio)
        u = torch.zeros(m.size(0), bottleneck)
        vh = torch.zeros(bottleneck, m.size(1))
        return u, vh

    def extra_repr(self) -&gt; str:
        return (f'in_features={self.in_features}, '
                f'out_features={self.out_features}, '
                f'bias=True, dim_ratio={self.dim_ratio}')

    def forward(self, x):
        return x @ (self.u @ self.vh).transpose(0, 1) + self.bias
</code></pre>
<p>In practice when I use <code>dim_ratio=0.25</code> I can achieve a 50% smaller network with slightly worse performance (~10%).</p>
<p>I took a look at this <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/quantized/modules/linear.py#L103" rel="noopener nofollow ugc">module</a> and this line more or less tells me that you cannot easily extend this class with my logic:</p>
<pre><code class="lang-auto">    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        return torch.ops.quantized.linear(
            x, self._packed_params._packed_params, self.scale, self.zero_point)
</code></pre>
<p>I can of course use this <code>self._packed_params = torch.ops.quantized.linear_prepack(weight, bias)</code> to  distribute the “small” package, and then just pre-calculate the <code>self.u @ self.vh</code> part and then essentially create the quantized network on-the-fly, but this would make model packaging / distribution very complicated compared to just loading a <code>jit</code> file.</p>
<p>All of this is very experimental and so far I have played only with <code>dim_ratio=0.25</code>, which makes the network itself 2x smaller. But quantization (in my case) usually also provides about 2-3x total module size reduction, so  I would need to compress with <code>dim_ratio=0.1</code> or lower to produce really small networks without quantization.</p>
<p>Maybe I am still missing an elephant in the room, idk. Hope someone will find this useful or have any ideas.</p>
<p>Combining quantization (up to 4x smaller networks) with factorization (2-4x smaller) looks like a killer feature to me.</p>

jerryzh168: <p>could you precompute <code>self.u @ self.vh</code> and then transform FactorizedLinear to Linear before quantization? so that it can be quantized to QuantizedLinear and get the size reduction?</p>

snakers41: <p>Yeah, I thought about this. But in our particular case this would make model distribution much more complicated (or it would negate the effects of factorization).</p>
<p>But I fould an elegant solution. I forgot that <code>bias=True</code> is not mandatory in PyTorch. You can just do this and it solves the problem. You can pass the above class to the below class and it will just quantize:</p>
<pre><code class="lang-auto">class FactorizedQLinear(nn.Module):
    def __init__(self,
                 f_linear):
        super().__init__()
        self.in_features = f_linear.in_features
        self.out_features = f_linear.out_features
        self.dim_ratio = f_linear.dim_ratio
        self.u_linear = nn.Linear(in_features=f_linear.u.data.size(0),
                                  out_features=f_linear.u.data.size(1),
                                  bias=True)
        self.vh_linear = nn.Linear(in_features=f_linear.vh.data.size(0),
                                   out_features=f_linear.vh.data.size(1),
                                   bias=False)
        self.u_linear.weight.data = f_linear.u.data
        self.u_linear.bias.data = f_linear.bias
        self.vh_linear.weight.data = f_linear.vh.data

    def extra_repr(self) -&gt; str:
        return (f'in_features={self.in_features}, '
                f'out_features={self.out_features}, '
                f'bias=True, dim_ratio={self.dim_ratio}')

    def forward(self, x):
        return self.u_linear(self.vh_linear(x))
</code></pre>

snakers41: <p>Also, this method really works. There are some sacrifices in quality (we do s2s), but probably in plain classification it will work just perfectly.</p>
<p>Also quantization used together with factorization can really reduce your model size ~10x, which is nice!</p>
<p>An <code>xsmall_q</code> model here - <a href="https://github.com/snakers4/silero-models/wiki/Quality-Benchmarks#en-v3" rel="noopener nofollow ugc">Quality Benchmarks · snakers4/silero-models Wiki (github.com)</a> - is 30x smaller that <code>large</code> models and it is created by a combination of quantization / minification / factorization.</p>