Analyse des contraintes et des optimisations pour une exécution locale

L'idée de faire fonctionner un modèle Transformer sur un module Wi-Fi à bas coût semble initialement contre-intuitive. Pourtant, une implémentation sur la puce ESP32-S3 est réalisable, avec des temps d'inférence autour de 200 ms et un modèle quantifié tenant dans moins de 100 Ko de mémoire. Cela permet des tâches de taritement du langage léger, comme la détection de mots-clés, sans aucune dépendance réseau.

Capacités matérielles de l'ESP32-S3

Contrairement aux perceptions courantes, l'ESP32-S3 offre des foncsionnalités avancées pour l'IA embarquée. Son architecture Xtensa LX7 double cœur tourne à 240 MHz et inclut des extensions vectorielles SIMD 128 bits. Cela accélère significativement les opérations de base des réseaux neuronaux, comme les multiplications matricielles, en traitant plusieurs données en une seule instruction via l'instruction VDOT.

Comparé aux solutions Cortex-M, l'ESP32-S3 présente des avantages clés pour le calcul neuronal à la périphérie : une largeur de vecteur supérieure, une fréquence d'horloge plus élevée, une connectivité sans fil intégrée et une bibliothèque d'optimisation neuronale (ESP-NN) fournie par le constructeur.

Architecture du modèle : simplification extrême

Pour exécuter un Transformer sur une puce avec environ 320 Ko de SRAM utilisable, une réduction drastique de l'architecture est nécessaire. Les paramètres recommandés sont :

• Profondeur ≤ 2 couches
• Dimension cachée (d_model) ≤ 64
• Têtes d'attention = 1 ou 2
• Longueur de séquence ≤ 32

Un tel modèle, après quantification, peut occuper moins de 100 Ko. L'utilisation d'une quantification en entiers 8 bits (int8) est obligatoire pour réduire l'empreinte mémoire et aligner les données sur les instructions SIMD optimisées.

Stratégies d'optimisation logicielle

Plusieurs transformations du modèle et de l'inférence sont requises pour une exécution efficace :

1. Remplacement des opérateurs coûteux :

2. Décomposition des opérations composites :

L'interpréteur TFLite Micro ne supporte pas nativement les blocs comme MultiHeadAttention. Ils doivent être décomposés en opérations basiques supportées (FULLY_CONNECTED, MATMUL, SOFTMAX, RESHAPE) avant la conversion du modèle. La bibliothèque ESP-NN fournit des versions accélérées de ces opérateurs pour int8.

3. Gestion de la mémoire :

Les poids du modèle sont stockés en Flash et mappés directement en mémoire. Pour les séquences longues, un traitement par morceaux (chunked inference) peut être nécessaire. Il est crucial d'allouer suffisamment de pile pour la tâche d'inférence (8-16 Ko) et d'utiliser la mémoire dynamique pour les grands tableaux.

Exemple de code optimisé

Voici des extraits illustrant l'utilisation des routines optimisées et une attention simplifiée.

Matrice multipliée accélérée en int8 :

#include "esp_nn.h"

void multiply_matrices_int8(const int8_t *input_vector, const int8_t *weight_matrix,
                            int8_t *output_vector, int input_depth, int output_depth) {
    const int output_multiplier = 1;
    const int right_shift = 15;
    const int32_t *bias = NULL; // À fournir si nécessaire
    int32_t *scratch_buf = NULL;

    esp_nn_mat_mult_s8(input_vector, weight_matrix, bias, output_vector,
                       scratch_buf, input_depth, output_depth, 1,
                       output_multiplier, right_shift);
}

Cette fonction appelle des instructions vectorielles sous le capot pour un calcul parallèle.

Implémentation d'une auto-attention simplifiée :

void compute_self_attention(const int8_t *input_seq, const int8_t *Wq, const int8_t *Wk,
                            const int8_t *Wv, int8_t *output, int seq_len, int depth) {
    // 1. Projection des requêtes, clés et valeurs
    int8_t *Q = (int8_t *)malloc(seq_len * depth);
    int8_t *K = (int8_t *)malloc(seq_len * depth);
    int8_t *V = (int8_t *)malloc(seq_len * depth);

    for (int i = 0; i < seq_len; ++i) {
        multiply_matrices_int8(&input_seq[i*depth], Wq, &Q[i*depth], depth, depth);
        multiply_matrices_int8(&input_seq[i*depth], Wk, &K[i*depth], depth, depth);
        multiply_matrices_int8(&input_seq[i*depth], Wv, &V[i*depth], depth, depth);
    }

    // 2. Calcul des scores d'attention (Q * K^T) avec produit scalaire optimisé
    int16_t *scores = (int16_t *)malloc(seq_len * seq_len * sizeof(int16_t));
    for (int i = 0; i < seq_len; ++i) {
        for (int j = 0; j < seq_len; ++j) {
            int32_t dot_product;
            esp_nn_dot_prod_s8(&Q[i*depth], &K[j*depth], depth, &dot_product);
            scores[i*seq_len + j] = (int16_t)(dot_product >> 3); // Échelle approximative pour sqrt(depth)
        }
    }

    // 3. Softmax sur les scores (version int8 à base de table)
    int8_t *attn_weights = (int8_t *)malloc(seq_len * seq_len);
    esp_nn_softmax_s8(scores, seq_len * seq_len, attn_weights);
    free(scores);

    // 4. Pondération des valeurs (attn_weights * V)
    int8_t *context = (int8_t *)malloc(seq_len * depth);
    for (int i = 0; i < seq_len; ++i) {
        multiply_matrices_int8(&attn_weights[i*seq_len], V, &context[i*depth], seq_len, depth);
    }

    // 5. Copie vers la sortie et libération mémoire
    memcpy(output, context, seq_len * depth);
    free(Q); free(K); free(V); free(attn_weights); free(context);
}

Performances et cas d'usage

Pour un modèle de reconnaissance de mots-clés (2 couches, d_model=64, seq_len=32, quantifié en int8), les temps de traitement typiques sur ESP32-S3-DevKitC-1 sont :

• Extraction des caractéristiques audio (MFCC) : ~45 ms
• Inférence du modèle Transformer : ~192 ms
• Traitement post-softmax : ~8 ms
• Latence totale : ~250 ms

La consommation pendant l'inférence est d'environ 80 mA à 3.3V. L'utilisation du coprocesseur ULP pour une détection préalable de signal audio peut réduire la consommation moyenne à moins de 10 mA.

Limites et perspectives

Cette solution est adaptée à des tâches de classification légère avec peu de catégories et des entrées courtes : détection de mots-clés, classification d'événements sensoriels simples, reconnaissance de commandes basiques. Elle n'est pas viable pour la génération de texte, la traduction ou la reconnaissance vocale continue.

L'optimisation se poursuit côté constructeur, avec des travaux sur des modules d'attention natifs dans ESP-DL et l'exploration d'architectures d'attention à complexité linéaire plus adaptées aux contraintes matérielles. De futures puces comme l'ESP32-P4 pourraient élargir encore le champ des possibles avec des capacités SIMD améliorées.