<jupyter_start><jupyter_text>Variance reductionBecause the sorting operation should be invariant to permutations, the error can be reduced very slightlyby running $m$ parallel sorting networks, with each input preceded by a fixed random permutation matrix, and averaging the results. This reduces the error (very) marginally.```python            def random_permutations(n, m):        return [np.eye(n)[np.random.permutation(np.arange(n))] for i in range(m)]    variance_reduced = np.mean([diff_sort(matrices, p @ x) for p in random_permutations(n, m)], axis=0)    ```<jupyter_code>np.random.seed(2019)
x = np.random.uniform(-200, 200, 16)
matrices = bitonic_woven_matrices(16)
def random_permutations(n, m):
        return [np.eye(n)[np.random.permutation(np.arange(n))] for i in range(m)]

def variance_reduced(matrices, x, m = 5):
    n = len(x)
    return np.mean([diff_sort_weave(matrices, p @ x) for p in random_permutations(n, m)], axis=0) 
def show_errors(xs, sort_fn):
    truth = np.sort(xs, axis=1)
    sorts = np.stack([sort_fn(x) for x in xs])
    
    plt.imshow(truth-sorts, cmap='RdBu', vmin=-1, vmax=1)
    plt.colorbar()
    plt.axis("off")
    print("RMS error", np.sqrt(np.mean((truth-sorts)**2)))<jupyter_output><empty_output><jupyter_text>Original<jupyter_code>xs = np.random.uniform(-250, 250, (15, 16))
show_errors(xs, lambda x:diff_sort_weave(matrices, x))<jupyter_output>RMS error 0.18603505359785322<jupyter_text>Averaged over 5 permutations<jupyter_code>show_errors(xs, lambda x:variance_reduced(matrices, x, m=5))<jupyter_output>RMS error 0.17893799497429339<jupyter_text>TemperingAs well adjustable $\alpha$ parameters in the smoothmax/logsumexp functions, we can also interpolate between sorting and no-sorting, by replacing each permutation matrix $A_i$ with $\beta I + (1-\beta A_i)$, for some small $\beta$.<jupyter_code>x = np.random.uniform(-250, 250, 16)
matrices = bitonic_woven_matrices(16)

def show_interpolations(matrices, x):
   
    alphas = [2**-alpha-2**-8 for alpha in np.linspace(-1,8,16)]
    sorts =  [[diff_sort_weave(matrices, x, softmax=lambda a,b:smoothmax(a,b,alpha=alpha), 
                                  beta=beta) for beta in np.linspace(0,1,16)] for alpha in alphas]
    
    fig = plt.figure(figsize=(12,12))
    for i,(alpha,alpha_block) in enumerate(zip(alphas, sorts)):
        ax = fig.add_subplot(4,4,i+1)
        ax.imshow(sorts[i], cmap='magma', vmin=-250, vmax=250)
        ax.text(0, 1, f"$\\alpha={alpha:.2f}$", color='white')
        if i%4==0:
            ax.text(-5, 0.5, "$\\beta=0.0$", color='black')
            ax.text(-5, 16-0.5, "$\\beta=1.0$", color='black')

        ax.axis("image")
        ax.axis("off")
show_interpolations(matrices, x)<jupyter_output><empty_output><jupyter_text>TransformationsAny invertible transform can be applied to values before sorting and inverted afterwards. Typically this would be to standardise the variables, and scale the z-scores by a factor of 150 or so:<jupyter_code>xs = np.random.uniform(-10, 10, (16, 16))
show_errors(xs, lambda x:diff_sort_weave(matrices, x))
std = (xs - np.mean(xs,axis=1))/ np.std(xs,axis=1)
show_errors(std, lambda x:(diff_sort_weave(matrices, x*150)/150))<jupyter_output>RMS error 0.0008208592869724583
{'fraction_non_alphanumeric': 0.14444102248229135, 'fraction_numerical': 0.04927625500461965, 'mean_word_length': 4.580756013745704, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '486107', 'license': 'MIT', 'n_tokens_mistral': 1299, 'n_tokens_neox': 1176, 'n_words': 273, 'path': 'notebooks/Variations.ipynb', 'repo_name': 'asaran/differentiable_sorting', 'size': '3247'}