Cách Sophos đào tạo một trình phát hiện phần mềm độc hại PDF nhẹ, mạnh mẽ ở quy mô cực lớn với Amazon SageMaker

Bài đăng này được đồng tác giả bởi Salma Taoufiq và Harini Kannan từ Sophos.

Là người đi đầu trong lĩnh vực an ninh mạng thế hệ tiếp theo, Sophos cố gắng bảo vệ hơn 500,000 tổ chức và hàng triệu khách hàng trên 150 quốc gia trước các mối đe dọa đang phát triển. Được hỗ trợ bởi trí thông minh về mối đe dọa, máy học (ML) và trí tuệ nhân tạo từ Sophos X-Ops, Sophos cung cấp một danh mục sản phẩm và dịch vụ tiên tiến đa dạng và phong phú để bảo mật và bảo vệ người dùng, mạng và thiết bị đầu cuối chống lại lừa đảo, ransomware, phần mềm độc hại và một loạt các cuộc tấn công mạng ngoài đó.

Sản phẩm Nhóm trí tuệ nhân tạo (AI) Sophos (SophosAI) giám sát việc phát triển và duy trì công nghệ bảo mật ML chính của Sophos. Bảo mật là một vấn đề dữ liệu lớn. Để tránh bị phát hiện, tội phạm mạng liên tục tạo ra các cuộc tấn công mới lạ. Điều này chuyển thành bộ dữ liệu về mối đe dọa khổng lồ mà nhóm phải làm việc với để bảo vệ khách hàng tốt nhất. Một ví dụ đáng chú ý là việc phát hiện và loại bỏ các tệp bị dính phần mềm độc hại một cách xảo quyệt, trong đó các tập dữ liệu có kích thước bằng terabyte.

Trong bài đăng này, chúng tôi tập trung vào hệ thống phát hiện phần mềm độc hại của Sophos dành riêng cho định dạng tệp PDF. Chúng tôi giới thiệu cách SophosAI sử dụng Amazon SageMaker đào tạo phân tán với hàng terabyte dữ liệu để đào tạo mô hình XGBoost (Extreme Gradient Bo boost) nhẹ mạnh mẽ. Điều này cho phép nhóm của họ lặp lại dữ liệu đào tạo lớn nhanh hơn với tính năng điều chỉnh siêu thông số tự động và không cần quản lý cơ sở hạ tầng đào tạo bên dưới.

Giải pháp hiện được tích hợp liền mạch vào quy trình đào tạo sản xuất và mô hình được triển khai trên hàng triệu điểm cuối của người dùng thông qua Dịch vụ điểm cuối Sophos.

Ngữ cảnh ca sử dụng

Cho dù bạn muốn chia sẻ một hợp đồng quan trọng hay giữ nguyên thiết kế đẹp mắt của CV, thì định dạng PDF là lựa chọn phổ biến nhất. Việc sử dụng rộng rãi và nhận thức chung rằng các tài liệu như vậy là kín gió và tĩnh đã khiến người dùng rơi vào cảm giác an toàn sai lầm. Do đó, PDF đã trở thành một vector lây nhiễm được lựa chọn trong kho vũ khí của những kẻ tấn công. Các hành động độc hại khi sử dụng tệp PDF thường xảy ra thông qua việc nhúng một tải trọng JavaScript được chạy bởi trình đọc PDF để tải xuống vi-rút từ URI, phá hoại máy của người dùng hoặc lấy cắp thông tin nhạy cảm.

Sophos phát hiện các tệp PDF độc hại tại các điểm khác nhau của cuộc tấn công bằng cách sử dụng một tập hợp các mô hình xác định và ML. Một trong những cách tiếp cận như vậy được minh họa trong sơ đồ sau, nơi tệp PDF độc hại được gửi qua email. Ngay sau khi nỗ lực tải xuống được thực hiện, nó sẽ kích hoạt tập lệnh thực thi độc hại để kết nối với máy chủ Chỉ huy và Điều khiển của kẻ tấn công. Trình phát hiện PDF của SophosAI chặn nỗ lực tải xuống sau khi phát hiện thấy nó độc hại.

Các cách khác bao gồm chặn các tệp PDF trong điểm cuối, gửi các tệp độc hại đến hộp cát (nơi nó được tính điểm bằng cách sử dụng nhiều mô hình), gửi tệp độc hại đến cơ sở hạ tầng tính điểm và tạo báo cáo bảo mật, v.v.

Động lực

Để xây dựng một công cụ phát hiện dựa trên cây có thể kết tội các tệp PDF độc hại với độ tin cậy cao, đồng thời cho phép tiêu thụ điện năng điểm cuối thấp và phản hồi suy luận nhanh, nhóm SophosAI nhận thấy thuật toán XGBoost là một ứng cử viên hoàn hảo cho nhiệm vụ này. Những con đường nghiên cứu như vậy rất quan trọng đối với Sophos vì hai lý do. Việc triển khai các mô hình nhỏ nhưng mạnh mẽ ở cấp độ điểm cuối của khách hàng có tác động lớn đến đánh giá sản phẩm của công ty bởi các nhà phân tích. Nó cũng, và quan trọng hơn, cung cấp trải nghiệm người dùng tốt hơn về tổng thể.

Thử thách kỹ thuật

Vì mục tiêu là tạo ra một mô hình có dung lượng bộ nhớ nhỏ hơn so với các trình phát hiện phần mềm độc hại PDF hiện có của họ (cả trên đĩa và trong bộ nhớ), SophosAI đã sử dụng XGBoost, một thuật toán phân loại với kỷ lục đã được chứng minh là tạo ra các mô hình nhỏ hơn đáng kể so với mạng nơ-ron trong khi đạt được ấn tượng hiệu suất trên dữ liệu dạng bảng. Trước khi mạo hiểm mô hình hóa các thử nghiệm XGBoost, một điều quan trọng cần cân nhắc là kích thước tuyệt đối của tập dữ liệu. Thật vậy, tập dữ liệu cốt lõi của Sophos về các tệp PDF có tính bằng terabyte.

Do đó, thách thức chính là đào tạo mô hình với một tập dữ liệu lớn mà không cần phải lấy mẫu xuống. Bởi vì điều quan trọng là máy dò tìm cách phát hiện bất kỳ cuộc tấn công dựa trên PDF nào - thậm chí là những cuộc tấn công hoàn toàn mới và hoàn toàn mới để bảo vệ khách hàng của Sophos tốt hơn - điều quan trọng nhất là sử dụng tất cả các bộ dữ liệu đa dạng có sẵn.

Không giống như mạng nơ-ron, nơi bạn có thể đào tạo theo lô, đối với XGBoost, chúng tôi cần toàn bộ tập dữ liệu đào tạo trong bộ nhớ. Tập dữ liệu đào tạo lớn nhất cho dự án này là hơn 1 TB, và không có cách nào để đào tạo trên quy mô như vậy mà không sử dụng các phương pháp luận của khung đào tạo phân tán.

Tổng quan về giải pháp

SageMaker là một dịch vụ ML được quản lý hoàn toàn, cung cấp các công cụ khác nhau để xây dựng, đào tạo, tối ưu hóa và triển khai các mô hình ML. Các Thư viện thuật toán tích hợp sẵn của SageMaker bao gồm 21 thuật toán ML phổ biến, bao gồm cả XGBoost. (Để biết thêm thông tin, hãy xem Đơn giản hóa việc học máy với XGBoost và Amazon SageMaker.) Với thuật toán tích hợp XGBoost, bạn có thể tận dụng lợi thế của nguồn mở Bộ chứa SageMaker XGBoost bằng cách chỉ định phiên bản khung lớn hơn 1.0-1, phiên bản này đã cải thiện tính linh hoạt, khả năng mở rộng, khả năng mở rộng và Đào tạo tại chỗ được quản lý, đồng thời hỗ trợ các định dạng đầu vào như Parquet, là định dạng được sử dụng cho tập dữ liệu PDF.

Lý do chính mà SophosAI chọn SageMaker là khả năng hưởng lợi từ việc đào tạo phân tán được quản lý hoàn toàn trên các phiên bản CPU nhiều nút bằng cách chỉ định nhiều hơn một phiên bản. SageMaker tự động phân chia dữ liệu giữa các nút, tổng hợp kết quả giữa các nút ngang hàng và tạo ra một mô hình duy nhất. Các phiên bản có thể là Phiên bản Spot, do đó giảm đáng kể chi phí đào tạo. Với thuật toán tích hợp cho XGBoost, bạn có thể làm điều này mà không cần bất kỳ tập lệnh tùy chỉnh bổ sung nào. Các phiên bản phân tán của XGBoost cũng tồn tại dưới dạng mã nguồn mở, chẳng hạn như XGBoost-Ray và XGBoost4J-Spark, nhưng việc sử dụng chúng đòi hỏi phải xây dựng, bảo mật, điều chỉnh và tự quản lý các cụm máy tính phân tán, điều này thể hiện nỗ lực đáng kể bổ sung cho sự phát triển khoa học.

Ngoài ra, SageMaker điều chỉnh mô hình tự động, còn được gọi là điều chỉnh siêu tham số, tìm phiên bản tốt nhất của mô hình bằng cách chạy nhiều công việc đào tạo với phạm vi siêu tham số mà bạn chỉ định. Sau đó, nó chọn các giá trị siêu tham số dẫn đến mô hình hoạt động tốt nhất, như được đo bằng số liệu cho nhiệm vụ ML đã cho.

Sơ đồ sau minh họa kiến ​​trúc giải pháp.

Cần lưu ý rằng, khi SophosAI bắt đầu thử nghiệm XGBoost trước khi chuyển sang SageMaker, các nỗ lực đã được thực hiện để sử dụng bộ nhớ lớn Đám mây điện toán đàn hồi Amazon (Amazon EC2) các phiên bản (ví dụ: r5a.24xlarge và x1.32xlarge) để đào tạo mô hình trên một mẫu dữ liệu càng lớn càng tốt. Tuy nhiên, những lần thử này trung bình mất hơn 10 giờ và thường không thành công do hết bộ nhớ.

Ngược lại, bằng cách sử dụng thuật toán SageMaker XGBoost và cơ chế đào tạo phân tán không phức tạp, SophosAI có thể đào tạo một mô hình tăng cường trên quy mô lớn trên tập dữ liệu đào tạo PDF khổng lồ trong vòng 20 phút. Nhóm chỉ phải lưu trữ dữ liệu trên Dịch vụ lưu trữ đơn giản của Amazon (Amazon S3) dưới dạng tệp Parquet có kích thước tương tự và chọn loại phiên bản EC2 và số lượng phiên bản mong muốn và SageMaker quản lý cơ sở hạ tầng cụm máy tính cơ bản và đào tạo phân tán giữa nhiều nút của cụm. Dưới mui xe, SageMaker chia nhỏ dữ liệu qua các nút bằng cách sử dụng ShardedByS3Key để phân phối các đối tượng tệp như nhau giữa mỗi phiên bản và sử dụng triển khai XGBoost của Giao thức Rabit (AllReduce và giao diện quảng bá đáng tin cậy) để khởi chạy quá trình xử lý phân tán và giao tiếp giữa các nút chính và nút ngang hàng. (Để biết thêm chi tiết về tổng hợp biểu đồ và phát trên các nút, hãy tham khảo XGBoost: Hệ thống tăng cường cây có thể mở rộng.)

Ngoài việc chỉ đào tạo một mô hình, với SageMaker, Điều chỉnh siêu tham số XGBoost cũng được thực hiện nhanh chóng và dễ dàng với khả năng chạy đồng thời các thí nghiệm khác nhau để tinh chỉnh sự kết hợp tốt nhất của các siêu thông số. Các siêu tham số có thể điều chỉnh bao gồm cả siêu tham số cụ thể tăng cường và siêu tham số cụ thể cho chức năng mục tiêu. Hai chiến lược tìm kiếm được cung cấp: ngẫu nhiên hoặc Bayesian. Chiến lược tìm kiếm Bayes đã được chứng minh là có giá trị vì nó giúp tìm các siêu tham số tốt hơn so với tìm kiếm ngẫu nhiên đơn thuần, với ít lần lặp lại thử nghiệm hơn.

Thông tin tập dữ liệu

Mô hình phát hiện phần mềm độc hại PDF của SophosAI dựa trên nhiều tính năng như biểu đồ n-gram và các tính năng entropy byte (Để biết thêm thông tin, hãy tham khảo MEADE: Hướng tới Công cụ phát hiện tệp đính kèm email độc hại). Siêu dữ liệu và các tính năng được trích xuất từ ​​các tệp PDF đã thu thập được lưu trữ trong kho dữ liệu phân tán. Sau đó, một tập dữ liệu gồm hơn 3,500 tính năng sẽ được tính toán, chia nhỏ hơn theo thời gian thành các tập huấn luyện và kiểm tra và được lưu trữ theo lô dưới dạng tệp Parquet trong Amazon S3 để SageMaker có thể dễ dàng truy cập cho các công việc đào tạo.

Bảng sau cung cấp thông tin về dữ liệu đào tạo và kiểm tra.

Bộ dữ liệu	Số lượng mẫu	Số lượng tập tin ván sàn	Tổng kích thước
Hội thảo	70,391,634	5,500	~ 1010 GB
Thử nghiệm	1,242,283	98	~ 18 GB

Kích thước dữ liệu đã được tính theo công thức:

Kích thước dữ liệu = N × (n_F + không_L) × 4

Công thức có các tham số sau:

• N là số lượng mẫu trong tập dữ liệu
• n_F là số lượng tính năng, với n_F = 3585
• n_L là số nhãn chân lý cơ bản, với n_L = 1
• 4 là số byte cần thiết cho kiểu dữ liệu của các tính năng: float32

Ngoài ra, các biểu đồ hình tròn sau đây cung cấp phân phối nhãn của cả tập huấn luyện và thử nghiệm, gây ra sự mất cân bằng lớp gặp phải trong nhiệm vụ phát hiện phần mềm độc hại PDF.

Sự phân bổ thay đổi từ tập huấn luyện sang tập kiểm tra một tháng. Việc phân chia theo thời gian của tập dữ liệu thành đào tạo và kiểm tra được áp dụng để mô phỏng kịch bản triển khai trong đời thực và tránh theo dõi tạm thời. Chiến lược này cũng cho phép SophosAI đánh giá khả năng tổng quát hóa thực sự của mô hình khi đối mặt với các cuộc tấn công PDF hoàn toàn mới chưa từng thấy trước đây.

Thử nghiệm và kết quả

Để bắt đầu các thử nghiệm, nhóm SophosAI đã đào tạo một mô hình XGBoost cơ bản với các thông số mặc định. Sau đó, họ bắt đầu thực hiện tinh chỉnh siêu tham số với SageMaker bằng cách sử dụng chiến lược Bayes, đơn giản như chỉ định siêu tham số được điều chỉnh và phạm vi giá trị mong muốn, chỉ số đánh giá (ROC (Đặc tính hoạt động của máy thu) AUC trong trường hợp này) và các bộ đào tạo và xác nhận. Đối với trình phát hiện phần mềm độc hại PDF, SophosAI ưu tiên các siêu tham số bao gồm số vòng tăng cường (num_round), độ sâu cây tối đa (max_depth), tỷ lệ học tập (eta) và tỷ lệ lấy mẫu cột khi xây dựng cây xanh (colsample_bytree). Cuối cùng, các siêu tham số tốt nhất đã được thu thập và sử dụng để đào tạo một mô hình trên tập dữ liệu đầy đủ, và cuối cùng được đánh giá trên tập kiểm tra giữ nguyên.

Biểu đồ sau đây cho thấy chỉ số mục tiêu (ROC AUC) so với 15 công việc đào tạo chạy trong công việc điều chỉnh. Các siêu tham số tốt nhất là những siêu thông số tương ứng với công việc đào tạo thứ chín.

Khi bắt đầu thử nghiệm của SophosAI trên SageMaker, một câu hỏi đặc biệt quan trọng cần trả lời là: loại trường hợp nào và bao nhiêu trường hợp trong số đó là cần thiết để đào tạo XGBoost về dữ liệu có trong tay? Điều này rất quan trọng vì việc sử dụng sai số lượng hoặc loại phiên bản có thể gây lãng phí thời gian và tiền bạc; quá trình đào tạo chắc chắn không thành công do hết bộ nhớ, hoặc, nếu sử dụng quá nhiều phiên bản quá lớn, điều này có thể trở nên tốn kém một cách không cần thiết.

XGBoost là một thuật toán giới hạn bộ nhớ (trái ngược với giới hạn máy tính). Vì vậy, một phiên bản máy tính có mục đích chung (ví dụ: M5) là lựa chọn tốt hơn một phiên bản được tối ưu hóa cho máy tính (ví dụ: C4). Để đưa ra quyết định sáng suốt, có một hướng dẫn đơn giản của SageMaker để chọn số trường hợp cần thiết để chạy đào tạo trên tập dữ liệu đầy đủ:

Tổng kích thước dữ liệu đào tạo × Hệ số an toàn^(*) <Số lượng phiên bản × Tổng bộ nhớ của Loại phiên bản

Trong trường hợp này: Tổng kích thước dữ liệu đào tạo × Hệ số an toàn (12) = 12120 GB

Bảng sau đây tóm tắt các yêu cầu khi loại phiên bản đã chọn là ml.m5.24xlarge.

Kích thước đào tạo × Hệ số an toàn (12)	Bộ nhớ phiên bản ml.m5.24xlarge	Số lượng phiên bản tối thiểu được yêu cầu để đào tạo
12120 GB	384 GB	32

^*Do bản chất của đào tạo phân tán XGBoost, đòi hỏi toàn bộ tập dữ liệu đào tạo phải được tải vào đối tượng DMatrix trước khi đào tạo và bộ nhớ trống bổ sung, hệ số an toàn là 10–12 được khuyến nghị.

Để xem xét kỹ hơn việc sử dụng bộ nhớ cho khóa đào tạo SageMaker đầy đủ về XGBoost trên tập dữ liệu được cung cấp, chúng tôi cung cấp biểu đồ tương ứng thu được từ khóa đào tạo amazoncloudwatch giám sát. Đối với công việc đào tạo này, các phiên bản phóng to 40 ml.m5.24 đã được sử dụng và mức sử dụng bộ nhớ tối đa đạt khoảng 62%.

Chi phí kỹ thuật tiết kiệm được bằng cách tích hợp một dịch vụ ML được quản lý như SageMaker vào đường ống dữ liệu là khoảng 50%. Tùy chọn sử dụng Phiên bản Spot cho các công việc đào tạo và điều chỉnh siêu thông số giúp cắt giảm chi phí thêm 63%.

Kết luận

Với SageMaker, nhóm SophosAI có thể giải quyết thành công một dự án phức tạp có mức độ ưu tiên cao bằng cách xây dựng mô hình XGBoost phát hiện phần mềm độc hại PDF nhẹ hơn nhiều trên đĩa (nhỏ hơn tới 25 lần) và trong bộ nhớ (nhỏ hơn 5 lần) so với máy dò tiền thân. Đó là một trình phát hiện phần mềm độc hại nhỏ nhưng mạnh mẽ với ~ 0.99 AUC và tỷ lệ dương tính thực sự là 0.99 và tỷ lệ dương tính giả là . Mô hình này có thể được đào tạo lại nhanh chóng và hiệu suất của nó có thể dễ dàng theo dõi theo thời gian, vì chỉ mất chưa đầy 20 phút để đào tạo nó trên hơn 1 TB dữ liệu.

Bạn có thể tận dụng thuật toán tích hợp SageMaker XGBoost để xây dựng mô hình với dữ liệu dạng bảng của bạn trên quy mô lớn. Ngoài ra, bạn cũng có thể thử các thuật toán Amazon SageMaker tích hợp mới LightGBM, CatBoost, AutoGluon-Tabular và Tab Transformer như được mô tả trong phần này  Blog của chúng tôi..

---

Giới thiệu về tác giả

Salma Taoufiq là Nhà khoa học dữ liệu cấp cao tại Sophos, làm việc tại điểm giao thoa giữa máy học và an ninh mạng. Với nền tảng đại học về khoa học máy tính, cô tốt nghiệp Thạc sĩ tại Đại học Trung Âu. trong Toán học và Ứng dụng của nó. Khi không phát triển công cụ phát hiện phần mềm độc hại, Salma là một người đam mê đi bộ đường dài, thích du lịch và thích thể loại phim kinh dị.

Harini Kannan là Nhà khoa học dữ liệu tại SophosAI. Cô ấy đã làm trong lĩnh vực khoa học dữ liệu bảo mật được ~ 4 năm. Cô ấy trước đây là Nhà khoa học dữ liệu chính tại Capsule8, được Sophos mua lại. Cô đã nói chuyện tại CAMLIS, BlackHat (Hoa Kỳ), Hội nghị Khoa học Dữ liệu Mở (Miền Đông), Viện Khoa học Dữ liệu, PyData (Boston) và Trình kết nối Dữ liệu. Các lĩnh vực nghiên cứu của cô bao gồm phát hiện các cuộc tấn công dựa trên phần cứng bằng cách sử dụng bộ đếm hiệu suất, phân tích hành vi của người dùng, ML có thể diễn giải và phát hiện bất thường không được giám sát.

Hasan Poonawala là Kiến trúc sư Giải pháp Chuyên gia về AI / ML Cấp cao tại AWS, có trụ sở tại London, Vương quốc Anh. Hasan giúp khách hàng thiết kế và triển khai các ứng dụng học máy trong sản xuất trên AWS. Anh ấy có hơn 12 năm kinh nghiệm làm việc với tư cách là nhà khoa học dữ liệu, người thực hành máy học và nhà phát triển phần mềm. Khi rảnh rỗi, Hasan thích khám phá thiên nhiên và dành thời gian cho bạn bè và gia đình.

Patel Digant là Trưởng nhóm hỗ trợ doanh nghiệp tại AWS. Anh ấy làm việc với khách hàng để thiết kế, triển khai và vận hành trên nền tảng đám mây ở quy mô lớn. Lĩnh vực quan tâm của anh ấy là thực hành MLOps và DevOps và cách nó có thể giúp khách hàng trong hành trình trên đám mây của họ. Ngoài công việc, anh thích chụp ảnh, chơi bóng chuyền và dành thời gian cho bạn bè và gia đình.