DExTra

DExTra, or Deep and Light-weight Expand-reduce Transformation, is a light-weight expand-reduce transformation that enables learning wider representations efficiently.

DExTra maps a dₘ dimensional input vector into a high dimensional space (expansion) and then reduces it down to a dₒ dimensional output vector (reduction) using N layers of group transformations. During these expansion and reduction phases, DExTra uses group linear transformations because they learn local representations by deriving the output from a specific part of the input and are more efficient than linear transformations. To learn global representations, DExTra shares information between different groups in the group linear transformation using feature shuffling

Formally, the DExTra transformation is controlled by five configuration parameters: (1) depth N, (2) width multiplier m_w, (3) input dimension dₘ, (4) output dimension dₒ, and (5) maximum groups gₘₐₓ in a group linear transformation. In the expansion phase, DExTra projects the dₘ-dimensional input to a high-dimensional space, dₘₐₓ = m_wdₘ, linearly using ceil(N/2) layers. In the reduction phase, DExTra projects the dₘₐₓ-dimensional vector to a dₒ-dimensional space using the remaining N -ceil(N/2) layers. Mathematically, we define the output Y at each layer l as:

𝐘ₗ = ℱ(𝐗, 𝐖ˡ, 𝐛ˡ, gˡ) if l=1 𝐘ₗ = ℱ(ℋ(𝐗, 𝐘ˡ⁻¹), 𝐖ˡ, 𝐛ˡ, gˡ) Otherwise

where the number of groups at each layer l are computed as:

gˡ = min(2ˡ⁻¹, gₘₐₓ), 1 ≤l ≤ceil(N/2) gᴺ⁻ˡ, Otherwise

In the above equations, ℱ is a group linear transformation function. The function ℱ takes the input (𝐗 or ℋ(𝐗, 𝐘ˡ⁻¹) ), splits it into gˡ groups, and then applies a linear transformation with learnable parameters 𝐖ˡ and bias 𝐛ˡ to each group independently. The outputs of each group are then concatenated to produce the final output 𝐘ˡ. The function ℋ first shuffles the output of each group in 𝐘ˡ⁻¹ and then combines it with the input 𝐗 using an input mixer connection.

In the authors' experiments, they use gₘₐₓ = ceil(dₘ/32) so that each group has at least 32 input elements. Note that (i) group linear transformations reduce to linear transformations when gˡ = 1, and (ii) DExTra is equivalent to a multi-layer perceptron when gₘₐₓ = 1.